Coverage for tcvx21/grillix_post/components/vector_m.py: 95%

1""" 

2Cylindrical vector operations based on xarray labelled axes 

3""" 

4import xarray as xr 

5import numpy as np 

6 

7 

8def _add_vector_dim(dims, coords): 

9 """Adds a 'vector' dimension with coords (eR, ePhi, eZ)""" 

10 dims = list(dims) 

11 dims.append('vector') 

12 

13 coords = dict(coords) 

14 coords['vector'] = ['eR', 'ePhi', 'eZ'] 

15 

16 return dims, coords 

17 

18 

19def toroidal_vector(input_array: xr.DataArray): 

20 """ 

21 Builds a cylindrical vector-field array from a phi component array, and zeros the (R, Z) components 

22 """ 

23 vector_array = np.zeros(tuple(list(input_array.shape) + [3])) 

24 vector_array[..., 1] = input_array 

25 

26 dims, coords = _add_vector_dim(input_array.dims, input_array.coords) 

27 

28 return xr.DataArray(vector_array, dims=dims, attrs=input_array.attrs, coords=coords) 

29 

30 

31def poloidal_vector(input_r: xr.DataArray, input_z: xr.DataArray, dims: list, coords: dict = {}, attrs: dict = {}): 

32 """ 

33 Builds a cylindrical vector-field array from (R, Z) component arrays, and zeros the phi component 

34 

35 dims should be the names of the dimensions (i.e. from input_array.dims) 

36 Optional: coords and attrs are used to set the coords and attributes of the output array 

37 """ 

38 assert input_r.shape == input_z.shape, f'input_r and input_z shapes are {input_r.shape} and {input_z.shape}.\ 

39 Please broadcast shapes to match before passing to poloidal vector' 

40 

41 vector_array = np.zeros(tuple(list(input_r.shape) + [3])) 

42 vector_array[..., 0] = input_r 

43 vector_array[..., 2] = input_z 

44 

45 dims, coords = _add_vector_dim(dims, coords) 

46 

47 return xr.DataArray(vector_array, dims=dims, attrs=attrs, coords=coords) 

48 

49 

50def cylindrical_vector(input_r: xr.DataArray, input_phi: xr.DataArray, input_z: xr.DataArray, 

51 dims: list, coords: dict = {}, attrs: dict = {}): 

52 """ 

53 Builds a cylindrical vector-field array from (R, phi, Z) component arrays 

54 

55 dims should be the names of the dimensions (i.e. from input_array.dims) 

56 Optional: coords and attrs are used to set the coords and attributes of the output array 

57 """ 

58 assert input_r.shape == input_phi.shape and input_r.shape == input_z.shape,f'input_r, input_phi and input_z shapes are {input_r.shape}, {input_phi.shape} {input_z.shape}.\ 

59 Please broadcast shapes to match before passing to cylindrical vector' 

60 

61 vector_array = np.zeros(tuple(list(input_r.shape) + [3])) 

62 vector_array[..., 0] = input_r 

63 vector_array[..., 1] = input_phi 

64 vector_array[..., 2] = input_z 

65 

66 dims, coords = _add_vector_dim(dims, coords) 

67 

68 return xr.DataArray(vector_array, dims=dims, attrs=attrs, coords=coords) 

69 

70 

71def eR_unit_vector(r_norm, z_norm, grid: bool = False): 

72 """ 

73 Unit vector pointing in the R (radial) direction 

74 """ 

75 

76 if not grid: 

77 return poloidal_vector( 

78 input_r=xr.DataArray(1.0), 

79 input_z=xr.DataArray(0.0), 

80 dims=[]) 

81 else: 

82 return poloidal_vector( 

83 input_r=xr.DataArray(np.ones((z_norm.size, r_norm.size))), 

84 input_z=xr.DataArray(np.zeros((z_norm.size, r_norm.size))), 

85 dims=['Z', 'R'], coords={'R': r_norm, 'Z': z_norm}) 

86 

87 

88def ePhi_unit_vector(r_norm, z_norm, grid: bool = False): 

89 """ 

90 Unit vector pointing in the phi (toroidal) direction 

91 """ 

92 

93 if not grid: 

94 return toroidal_vector(xr.DataArray(1.0)) 

95 else: 

96 return toroidal_vector(xr.DataArray(np.ones((z_norm.size, r_norm.size)))) 

97 

98 

99def eZ_unit_vector(r_norm, z_norm, grid: bool = False): 

100 """ 

101 Unit vector pointing in the Z (vertical) direction 

102 """ 

103 

104 if not grid: 

105 return poloidal_vector( 

106 input_r=xr.DataArray(0.0), 

107 input_z=xr.DataArray(1.0), 

108 dims=[]) 

109 else: 

110 return poloidal_vector( 

111 input_r=xr.DataArray(np.zeros((z_norm.size, r_norm.size))), 

112 input_z=xr.DataArray(np.ones((z_norm.size, r_norm.size))), 

113 dims=['Z', 'R'], coords={'R': r_norm, 'Z': z_norm}) 

114 

115 

116def vector_magnitude(input_array: xr.DataArray): 

117 """ 

118 Returns the vector magnitude of input_array 

119 

120 3.39 ms ± 276 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) 

121 """ 

122 return np.sqrt(vector_dot(input_array, input_array)) 

123 

124 

125def poloidal_vector_magnitude(input_array: xr.DataArray): 

126 """ 

127 Returns the vector magnitude of the poloidal component of the input array 

128 """ 

129 return vector_magnitude(input_array.isel(vector=[0, 2])) 

130 

131 

132def vector_dot(input_a: xr.DataArray, input_b: xr.DataArray): 

133 """ 

134 Returns the dot product of the input arrays 

135 

136 vec(a) dot vec(b) = mag(a) mag(b) cos(theta) 

137 where theta is the angle between vec(a) and vec(b) 

138 

139 2.71 ms ± 149 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) 

140 """ 

141 return input_a.dot(input_b, dims=['vector']) 

142 

143 

144def vector_cross(input_a, input_b): 

145 """ 

146 Returns the cross product of the input arrays 

147 

148 vec(a) cross vec(b) = mag(a) mag(b) sin(theta) vec(n) 

149 where theta is the angle between vec(a) and vec(b) and 

150 where vec(n) is a unit vector perpendicular to both vec(a) and vec(b) 

151 The sign of vec(b) is given by right-hand-rule (i.e. if 'a' is index finger, 'b' is middle finger, 'n' is thumb, or 

152 alternatively if 'a' is thumb, 'b' is index finger, 'n' is middle finger) 

153 

154 Compatible with dask, parallelization uses input_a.dtype as output_dtype 

155 Taken from https://github.com/pydata/xarray/issues/3279 

156 """ 

157 return xr.apply_ufunc(np.cross, input_a, input_b, 

158 input_core_dims=[['vector'], ['vector']], 

159 output_core_dims=[['vector']], 

160 dask='parallelized', output_dtypes=[input_a.dtype] 

161 ) 

162 

163 

164def unit_vector(input_array: xr.DataArray): 

165 """ 

166 Returns the unit vector in the direction of input_array 

167 """ 

168 return input_array / vector_magnitude(input_array) 

169 

170 

171def scalar_projection(input_a: xr.DataArray, input_b: xr.DataArray): 

172 """ 

173 Returns the magnitude of input_a which is parallel to input_b 

174 

175 See https://en.wikipedia.org/wiki/Vector_projection for notation. scalar_projection is 'a_1' 

176 """ 

177 return vector_dot(input_a, unit_vector(input_b)) 

178 

179 

180def vector_projection(input_a: xr.DataArray, input_b: xr.DataArray): 

181 """ 

182 Returns a vector which is the orthogonal projection of input_a onto a straight line parallel to input_b. 

183 It is parallel to input_b. 

184 

185 See https://en.wikipedia.org/wiki/Vector_projection for notation. vector_rejection is 'vec(a)_1' 

186 """ 

187 return scalar_projection(input_a, input_b) * unit_vector(input_b) 

188 

189 

190def vector_rejection(input_a: xr.DataArray, input_b: xr.DataArray): 

191 """ 

192 Returns a vector which is the projection of input_a onto a straight line orthogonal to input_b. 

193 

194 See https://en.wikipedia.org/wiki/Vector_projection for notation. vector_rejection is 'vec(a)_2' 

195 """ 

196 return input_a - vector_projection(input_a, input_b)