DISPATCH
curvilinear_transforms.py
1 # -*- coding: utf-8 -*-
2 
3 """
4  Coordinate transformation functions for DISPATCH2 data.
5 """
6 
7 import numpy as np
8 
9 
10 def rtp_to_xyz(r, theta, phi, theta_rotate=0.0, phi_rotate=0.0):
11  """
12  Convert 1-D spherical (r, theta, phi) coords. to 3-D Cartesian coords.
13  (x, y, z) for plotting purposes.
14 
15  theta_rotate: rotate the coordinate system in the theta-direction by
16  `theta_rotate`; the default gives a r-coord that overlies
17  the -z-coord when theta=0.
18  phi_rotate: rotate the coordinate system in the phi-direction by
19  `phi_rotate`; the default results in a r-coord that overlies
20  the x-coord when theta = pi/2.
21 
22  """
23 
24  # First establish 3-D versions of r, theta, and phi.
25  rr, tt, pp = np.meshgrid(r, theta, phi, indexing='ij')
26 
27  # Calculate 3-D Cartesian coordinate vectors.
28  xx = rr * np.sin(tt + theta_rotate) * np.cos(pp + phi_rotate)
29  yy = -rr * np.sin(tt + theta_rotate) * np.sin(pp + phi_rotate)
30  zz = -rr * np.cos(tt + theta_rotate)
31 
32  return xx, yy, zz
33 
34 def zrp_to_xyz(z, r, phi, phi_rotate=0.0):
35  """
36  Convert 1-D cylindrical (Z, R, Phi) coords. to 3-D Cartesian coords.
37  (x, y, z) for plotting purposes.
38 
39  phi_rotate: rotate the coordinate system in the phi-direction by
40  'phi_rotate'; the the default gives a R-coord. that overlies
41  the x-coord when phi = 0.
42 
43  """
44 
45  # First establish 3-D versions of Z, R, and Phi.
46  ZZ, RR, PP = np.meshgrid(z, r, phi, indexing='ij')
47 
48  # Calculate 3-D Cartesian coordinate vectors.
49  xx = RR * np.cos(PP + phi_rotate)
50  yy = RR * np.sin(PP + phi_rotate)
51  zz = ZZ
52 
53  return xx, yy, zz
54 
55 def rtp_to_zrp(r, theta, phi, theta_rotate=0.0, phi_rotate=0.0):
56  """
57  Convert 1-D spherical (r, theta, phi) coords. to 3-D cylindrical coords.
58  (Z, R, Phi).
59 
60  theta_rotate: rotate the coordinate system in the theta-direction by
61  `theta_rotate`; the default gives a r-coord that overlies
62  the -z-coord when theta=0.
63  phi_rotate: rotate the coordinate system in the phi-direction by
64  `phi_rotate`; the default results in a r-coord that overlies
65  the R-coord when theta = pi/2.
66 
67  """
68 
69  # First establish 3-D versions of r, theta, and phi.
70  rr, tt, pp = np.meshgrid(r, theta, phi, indexing='ij')
71 
72  # Calculate 3-D cylindrical coordinate vectors.
73  ZZ = -rr * np.cos(tt + theta_rotate)
74  RR = rr * abs(np.sin(tt + theta_rotate))
75  PP = -pp - phi_rotate
76 
77  return ZZ, RR, PP
78 
79 def zrp_to_rtp(z, r, phi, phi_rotate=0.0):
80  """
81  Convert 1-D cylindrical (Z, R, Phi) coords. to 3-D spherical coords.
82  (r, theta, phi).
83 
84  phi_rotate: rotate the coordinate system in the phi-direction by
85  'phi_rotate'; the default gives a R-coord which overlies
86  the r-coord when theta=pi/2.
87 
88  """
89 
90  # First establish 3-D versions of Z, R, and Phi.
91  ZZ, RR, PP = np.meshgrid(z, r, phi, indexing='ij')
92 
93  # Calculate 3-D spherical coordinate vectors.
94  rr = np.sqrt(RR**2 + ZZ**2)
95  tt = np.arctan2(RR, ZZ)
96  pp = -PP - phi_rotate
97 
98  return rr, tt, pp
99 
100 def xyz_to_zrp(x, y, z):
101  """
102  Convert 1-D Cartesian (x, y, z) coords. to 3-D cylindrical coords.
103  (Z, R, Phi).
104 
105  The z- and Z-coords. are assumed equivalent.
106 
107  """
108 
109  # First establish 3-D versions of x, y, z.
110  xx, yy, zz = np.meshgrid(x, y, z, indexing='ij')
111 
112  # Calculate 3-D cylindrical coordinate vectors.
113  ZZ = zz
114  RR = np.sqrt(xx**2 + yy**2)
115  PP = np.arctan2(yy, xx)
116 
117  return ZZ, RR, PP
118 
119 def xyz_to_rtp(x, y, z):
120  """
121  Convert 1-D Cartesian (x, y, z) coords. to 3-D spherical coords.
122  (r, theta, phi).
123 
124  The z-coord. is assumed to be anti-parallel to the r-coord. when
125  theta = 0.
126 
127  """
128 
129  # First establish 3-D versions of x, y, z
130  xx, yy, zz = np.meshgrid(x, y, z, indexing='ij')
131 
132  # Calculate 3-D spherical coordinate vectors.
133  rr = np.sqrt(xx**2 + yy**2 + zz**2)
134  tt = np.arccos(zz / rr)
135  pp = np.arccos(xx / np.sqrt(xx**2 + yy**2))
136 
137  return rr, tt, pp