Coverage for tcvx21/grillix_post/components/equi_m.py: 90%

1""" 

2Interface to the magnetic equilibrium 

3""" 

4from pathlib import Path 

5import numpy as np 

6import xarray as xr 

7from scipy.interpolate import RectBivariateSpline 

8from tcvx21.units_m import Quantity, convert_xarray_to_quantity, Dimensionless 

9from tcvx21.analysis.contour_finding_m import find_contours 

10from .vector_m import poloidal_vector, cylindrical_vector 

11 

12 

13def array_coords(array, grid, r_norm, z_norm): 

14 """Sets the attributes of an array based on whether it returns with grid=True""" 

15 

16 if grid: 

17 return array.rename({'dim_0': 'Z', 'dim_1': 'R'}).assign_coords(R=r_norm, Z=z_norm) 

18 else: 

19 return array.rename({'dim_0': 'points'}) 

20 

21 

22class Equi: 

23 

24 def __init__(self, equi_file: Path, penalisation_file: Path = None, flip_z: bool = False): 

25 """ 

26 Initialises an object which stores information about the magnetic geometry and boundaries 

27 """ 

28 assert equi_file.exists() 

29 

30 magnetic_geometry = xr.open_dataset(equi_file, group='Magnetic_geometry') 

31 psi_limits = xr.open_dataset(equi_file, group='Psi_limits') 

32 

33 if penalisation_file is not None: 

34 assert penalisation_file.exists() 

35 penalisation_file = xr.open_dataset(penalisation_file) 

36 self.penalisation_available = True 

37 

38 self.penalisation_characteristic = penalisation_file['pen_chi'].rename({'dim_nl': 'points'}) 

39 self.penalisation_direction = penalisation_file['pen_xi'].rename({'dim_nl': 'points'}) 

40 

41 else: 

42 self.penalisation_available = False 

43 self.penalisation_characteristic, self.penalisation_direction = None, None 

44 

45 self.magnetic_geometry = magnetic_geometry 

46 self.psi_limits = psi_limits 

47 

48 self.poloidal_field_factor = +1.0 

49 self.flipped_z = False 

50 

51 self.axis_r = xr.DataArray(magnetic_geometry.magnetic_axis_R).assign_attrs( 

52 norm=Quantity(1, magnetic_geometry.magnetic_axis_units)) 

53 self.axis_z = xr.DataArray(magnetic_geometry.magnetic_axis_Z).assign_attrs( 

54 norm=Quantity(1, magnetic_geometry.magnetic_axis_units)) 

55 

56 self.x_point_r = xr.DataArray(np.atleast_1d(magnetic_geometry.x_point_R)).assign_attrs( 

57 norm=Quantity(1, magnetic_geometry.x_point_units)) 

58 self.x_point_z = xr.DataArray(np.atleast_1d(magnetic_geometry.x_point_Z)).assign_attrs( 

59 norm=Quantity(1, magnetic_geometry.x_point_units)) 

60 

61 self.field_on_axis = xr.DataArray(magnetic_geometry.axis_Btor).assign_attrs( 

62 norm=Quantity(1, magnetic_geometry.axis_Btor_units)) 

63 

64 # Read the basis vectors for the spline (in normalised units) 

65 self.spline_basis_r = xr.DataArray(magnetic_geometry.R).assign_attrs(norm='R0') 

66 self.spline_basis_z = xr.DataArray(magnetic_geometry.Z).assign_attrs(norm='R0') 

67 

68 # Read the psi data (in Weber) 

69 try: 

70 self.psi_units = self.magnetic_geometry.psi_units 

71 except AttributeError: 

72 self.psi_units = self.magnetic_geometry.psi.units 

73 

74 self.psi_data = xr.DataArray(self.magnetic_geometry.psi).assign_attrs( 

75 norm=Quantity(1, self.psi_units)) 

76 

77 # Note that the axis ordering is inverted relative to the output of 

78 # meshgrid. 

79 self.psi_interpolator = RectBivariateSpline( 

80 self.spline_basis_r, self.spline_basis_z, self.psi_data.T) 

81 

82 if flip_z: 

83 self.flip_z() 

84 

85 def flip_z(self): 

86 """ 

87 Flips the vertical direction of the equilibrium, to affect a toroidal field reversal 

88 """ 

89 

90 self.flipped_z = not self.flipped_z 

91 self.poloidal_field_factor *= -1.0 

92 

93 self.axis_z *= -1.0 

94 self.x_point_z *= -1.0 

95 

96 self.spline_basis_z = -self.spline_basis_z[::-1] 

97 self.psi_data = self.psi_data[::-1, :] 

98 

99 self.psi_interpolator = RectBivariateSpline( 

100 self.spline_basis_r, self.spline_basis_z, self.psi_data.T) 

101 

102 @property 

103 def poloidal_flux_on_axis(self): 

104 """Poloidal flux in Wb on axis""" 

105 return xr.DataArray(self.psi_limits.psi_axis).assign_attrs(norm=Quantity(1, self.psi_units)) 

106 

107 @property 

108 def poloidal_flux_on_separatrix(self): 

109 """Poloidal flux in Wb on separatrix/at the x-point""" 

110 return xr.DataArray(self.psi_limits.psi_seperatrix).assign_attrs(norm=Quantity(1, self.psi_units)) 

111 

112 @property 

113 def axis_r_norm(self): 

114 """Magnetic axis radial position in normalised units (1 by definition)""" 

115 return xr.DataArray(1).assign_attrs(norm='R0') 

116 

117 @property 

118 def axis_z_norm(self): 

119 """Magnetic axis vertical position in normalised units (usually close to 0)""" 

120 return xr.DataArray(self.axis_z / self.axis_r).assign_attrs(norm='R0') 

121 

122 @property 

123 def x_point_r_norm(self): 

124 """x-point radial position in normalised units""" 

125 return xr.DataArray(self.x_point_r / self.axis_r).assign_attrs(norm='R0') 

126 

127 @property 

128 def x_point_z_norm(self): 

129 """x-point vertical position in normalised units""" 

130 return xr.DataArray(self.x_point_z / self.axis_r).assign_attrs(norm='R0') 

131 

132 @property 

133 def axis_btor(self): 

134 """Returns the on-axis field as a quantity""" 

135 return convert_xarray_to_quantity(self.field_on_axis) 

136 

137 def poloidal_flux(self, r_norm, z_norm, grid: bool = True): 

138 """Returns the poloidal flux in Wb at a point r_norm, z_norm""" 

139 

140 return array_coords(xr.DataArray(self.psi_interpolator(x=r_norm, y=z_norm, grid=grid).T).assign_attrs( 

141 norm=Quantity(1, self.psi_units)), 

142 grid=grid, r_norm=r_norm, z_norm=z_norm) 

143 

144 

145 def magnetic_field_r(self, r_norm, z_norm, grid: bool = True): 

146 """ 

147 Returns the normalised radial magnetic field at a point r_norm, z_norm 

148 

149 Evaluate the poloidal flux, taking the 0th x derivative and the 1st y derivative 

150 Then calculate the radial magnetic field according to eqn 2.25 of H. Zohm, 'MHD stability of Tokamaks', p24 

151 

152 N.b. One factor of self%axis_r comes from evaluating the derivative on the (r_norm', z_norm') normalised coordinate 

153 and the second comes from using normalised r_norm', z_norm' for the axis 

154 """ 

155 psi_dz = xr.DataArray(np.atleast_1d(self.psi_interpolator(x=r_norm, y=z_norm, dx=0, dy=1, grid=grid).T))\ 

156 .assign_attrs(norm=Quantity(1, 'Wb')) 

157 

158 b_r = psi_dz / \ 

159 (2 * np.pi * (r_norm[None, ...] if grid else r_norm) * 

160 self.axis_r.values * self.axis_r.values) 

161 

162 return array_coords(xr.DataArray(-self.poloidal_field_factor * b_r / self.field_on_axis).assign_attrs( 

163 norm=self.axis_btor), 

164 grid=grid, r_norm=r_norm, z_norm=z_norm) 

165 

166 def magnetic_field_z(self, r_norm, z_norm, grid: bool = True): 

167 """ 

168 Returns the normalised vertical magnetic field at a point r_norm, z_norm 

169 

170 Evaluate the poloidal flux, taking the 1st x derivative and the 0th y derivative 

171 Then calculate the radial magnetic field according to eqn 2.24 of H. Zohm, 'MHD stability of Tokamaks', p24 

172 

173 N.b. One factor of self%axis_r comes from evaluating the derivative on the (r_norm', z_norm') normalised coordinate 

174 and the second comes from using normalised r_norm', z_norm' for the axis 

175 """ 

176 psi_dr = xr.DataArray(np.atleast_1d(self.psi_interpolator(x=r_norm, y=z_norm, dx=1, dy=0, grid=grid).T))\ 

177 .assign_attrs(norm=Quantity(1, 'Wb')) 

178 

179 b_z = psi_dr / \ 

180 (2 * np.pi * (r_norm[None, ...] if grid else r_norm) * 

181 self.axis_r.values * self.axis_r.values) 

182 

183 return array_coords(xr.DataArray(self.poloidal_field_factor * b_z / self.field_on_axis).assign_attrs( 

184 norm=self.axis_btor), 

185 grid=grid, r_norm=r_norm, z_norm=z_norm) 

186 

187 def magnetic_field_toroidal(self, r_norm, z_norm, grid: bool = False) -> xr.DataArray: 

188 """Normalised magnetic field strength in the phi direction""" 

189 

190 if not grid: 

191 # Scalar case 

192 return array_coords(xr.DataArray(1.0 / np.atleast_1d(r_norm)).assign_attrs(norm='B0'), 

193 grid=grid, r_norm=r_norm, z_norm=z_norm) 

194 else: 

195 # Basis vector case 

196 r_mesh, _ = np.meshgrid(r_norm, z_norm) 

197 return array_coords(xr.DataArray(1.0 / np.atleast_1d(r_mesh)).assign_attrs(norm='B0'), 

198 grid=grid, r_norm=r_norm, z_norm=z_norm) 

199 

200 def magnetic_field_poloidal(self, r_norm, z_norm, grid: bool = False) -> xr.DataArray: 

201 """Normalised magnetic field strength in the poloidal plane. N.b. norm is inherited from components""" 

202 b_x = self.magnetic_field_r(r_norm, z_norm, grid) 

203 b_y = self.magnetic_field_z(r_norm, z_norm, grid) 

204 return xr.DataArray(np.sqrt(b_x**2 + b_y**2)).assign_attrs(norm=self.axis_btor) 

205 

206 def magnetic_field_absolute(self, r_norm, z_norm, grid: bool = False) -> xr.DataArray: 

207 """Total normalised magnetic field strength. N.b. norm is inherited from components""" 

208 b_x = self.magnetic_field_r(r_norm, z_norm, grid) 

209 b_y = self.magnetic_field_z(r_norm, z_norm, grid) 

210 b_t = self.magnetic_field_toroidal(r_norm, z_norm, grid) 

211 return xr.DataArray(np.sqrt(b_x**2 + b_y**2 + b_t**2)).assign_attrs(norm=self.axis_btor) 

212 

213 def normalised_flux_surface_label(self, r_norm, z_norm, grid: bool = True) -> xr.DataArray: 

214 """ 

215 Flux surface which is 0 at the magnetic axis and 1 at the separatrix 

216 """ 

217 

218 return array_coords(self.normalise_flux_surface_label(psi_value=self.poloidal_flux(r_norm, z_norm, grid)), 

219 grid=grid, r_norm=r_norm, z_norm=z_norm) 

220 

221 def normalise_flux_surface_label(self, psi_value): 

222 """ 

223 Provides a method to convert from poloidal flux (psi) to the normalised flux coordinate (rho) 

224 """ 

225 rho_squared = (psi_value - self.poloidal_flux_on_axis) \ 

226 / (self.poloidal_flux_on_separatrix - self.poloidal_flux_on_axis) 

227 

228 rho_squared = np.atleast_1d(rho_squared) 

229 rho_squared[np.asarray(rho_squared < 0.0)] = 0.0 

230 

231 return xr.DataArray(np.sqrt(rho_squared)).assign_attrs(norm=Dimensionless) 

232 

233 def get_separatrix(self, r_norm, z_norm, level: float = 1.0) -> list: 

234 """ 

235 Returns the separatrix contour 

236 

237 Can plot a specific contour via plt.plot(*separatrix[index].T) where index is a contour index 

238 """ 

239 assert r_norm.ndim == 1 and z_norm.ndim == 1 

240 

241 return find_contours( 

242 r_norm, z_norm, self.normalised_flux_surface_label(r_norm, z_norm), level=level) 

243 

244 def poloidal_unit_vector(self, r_norm, z_norm, grid: bool = False) -> xr.DataArray: 

245 """Unit vector pointing in the direction of the poloidal magnetic field. i.e. 'along a flux surface'""" 

246 b_x = self.magnetic_field_r(r_norm, z_norm, grid) 

247 b_y = self.magnetic_field_z(r_norm, z_norm, grid) 

248 b_pol = np.sqrt(b_x ** 2 + b_y ** 2) 

249 

250 if grid: 

251 return poloidal_vector(b_x / b_pol, b_y / b_pol, dims=['Z', 'R'], coords={'R': r_norm, 'Z': z_norm}) 

252 elif r_norm.size == 1: 

253 return poloidal_vector(b_x / b_pol, b_y / b_pol, dims=[]) 

254 else: 

255 return poloidal_vector(b_x / b_pol, b_y / b_pol, dims=['points']) 

256 

257 def radial_unit_vector(self, r_norm, z_norm, grid: bool = False) -> xr.DataArray: 

258 """Unit vector pointing normal to the direction of the poloidal magnetic field. i.e. 'across a flux surface'""" 

259 b_x = self.magnetic_field_r(r_norm, z_norm, grid) 

260 b_y = self.magnetic_field_z(r_norm, z_norm, grid) 

261 b_pol = np.sqrt(b_x ** 2 + b_y ** 2) 

262 

263 if grid: 

264 return poloidal_vector(-b_y / b_pol, b_x / b_pol, dims=['Z', 'R'], coords={'R': r_norm, 'Z': z_norm}) 

265 elif r_norm.size == 1: 

266 return poloidal_vector(-b_y / b_pol, b_x / b_pol, dims=[]) 

267 else: 

268 return poloidal_vector(-b_y / b_pol, b_x / b_pol, dims=['points']) 

269 

270 def parallel_unit_vector(self, r_norm, z_norm, grid: bool = False) -> xr.DataArray: 

271 """Unit vector pointing in the direction of the total magnetic field""" 

272 b_x = self.magnetic_field_r(r_norm, z_norm, grid) 

273 b_y = self.magnetic_field_z(r_norm, z_norm, grid) 

274 b_t = self.magnetic_field_toroidal(r_norm, z_norm, grid) 

275 b_abs = np.sqrt(b_x ** 2 + b_y ** 2 + b_t ** 2) 

276 

277 if grid: 

278 return cylindrical_vector(input_r=b_x / b_abs, input_phi=b_t / b_abs, input_z=b_y / b_abs, 

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

280 elif r_norm.size == 1: 

281 return cylindrical_vector(input_r=b_x / b_abs, input_phi=b_t / b_abs, input_z=b_y / b_abs, dims=[]) 

282 else: 

283 return cylindrical_vector(input_r=b_x / b_abs, input_phi=b_t / b_abs, input_z=b_y / b_abs, dims=['points']) 

284