Coverage for tcvx21/grillix_post/lineouts/lineouts_m.py: 100%

1""" 

2Implementations of lineouts specific to the tcvx21 analysis 

3""" 

4import numpy as np 

5import warnings 

6from pathlib import Path 

7from numpy.polynomial import Polynomial 

8import matplotlib.pyplot as plt 

9import xarray as xr 

10import pint 

11 

12from .lineout_m import Lineout, on_grid_interpolator 

13from tcvx21.analysis.contour_finding_m import find_contours 

14from tcvx21.file_io.json_io_m import read_from_json 

15from tcvx21.file_io.matlab_file_reader_m import read_struct_from_file 

16from tcvx21.units_m import Quantity, convert_xarray_to_quantity 

17 

18 

19def outboard_midplane_chord(grid, equi, n_points: int = 1000): 

20 """ 

21 Returns n_points of equally spaced points along the outboard zaxis, which is a horizontal line from the magnetic axis 

22 outwards (to large radius). N.b. the outboard midplane is usually to largest radius, but this is roughly equivalent 

23 """ 

24 

25 r_start = equi.axis_r_norm 

26 z_start = equi.axis_z_norm 

27 

28 r_end = grid.r_s.max() 

29 z_end = equi.axis_z_norm 

30 

31 lineout = Lineout([r_start, r_end], [z_start, z_end]) 

32 lineout.find_points_on_grid(grid, n_points=n_points, trim_points=False) 

33 

34 return lineout 

35 

36 

37def penalisation_contour(grid, equi, level: float, contour_index: int = None, epsilon: float = 1E-10, 

38 smoothing: float = 0.1, n_points: int = 1000): 

39 """ 

40 Returns a lineout along a contour of the penalisation characteristic 

41 

42 Should use torx.specialisations.*.penalisation_m.filled_penalisation_characteristic to get the 

43 penalisation_characteristic filled to include pghost points 

44 

45 smoothing is recommended, since this removes the grid staircase pattern of the contour 

46 """ 

47 penalisation_characteristic = grid.shape(equi.penalisation_characteristic) 

48 

49 assert 'R' in penalisation_characteristic.dims and 'Z' in penalisation_characteristic.dims, "Call vector_to_matrix on\ 

50 penalisation_chacteristic before calling penalisation_contour" 

51 

52 # Shift the level slightly in from [0, 1] to ensure a clear contour 

53 level = max(level, epsilon) 

54 level = min(level, 1.0 - epsilon) 

55 

56 contours = find_contours(grid.r_s, grid.z_s, penalisation_characteristic, level=level) 

57 

58 if equi.flipped_z: 

59 # Reverse the contour list is using flipped_z, since contours are returned from bottom to top 

60 contours.reverse() 

61 

62 if contour_index is None: # pragma: no cover 

63 plt.pcolormesh(grid.r_s, grid.z_s, penalisation_characteristic, shading='nearest') 

64 for i, contour in enumerate(contours): 

65 r_points, z_points = contour.T 

66 

67 plt.plot(r_points, z_points, label=f'contour_index = {i}') 

68 plt.legend() 

69 plt.show() 

70 

71 raise ValueError("Select a contour index for penalisation_contour") 

72 

73 else: 

74 r_points, z_points = contours[contour_index].T 

75 

76 lineout = Lineout(r_points, z_points, smoothing=smoothing) 

77 lineout.find_points_on_grid(grid, trim_points=True, n_points=n_points) 

78 

79 return lineout 

80 

81 

82def thomson_scattering(grid, equi, thomson_position: Path, n_points: int = 1000): 

83 """ 

84 Returns a lineout for the Thomson scattering points below the magnetic axis 

85 """ 

86 thomson_position_ = read_from_json(thomson_position) 

87 

88 # For the simulation, we only want divertor entrance points (below the magnetic axis) 

89 # Use the ts_below_axis mask to limit to this region 

90 ts_below_axis = thomson_position_['Z'] < 0 

91 thomson_position_ = {'R': thomson_position_['R'][ts_below_axis], 

92 'Z': thomson_position_['Z'][ts_below_axis], } 

93 

94 R0 = convert_xarray_to_quantity(equi.axis_r) 

95 

96 ts_r = (Quantity(thomson_position_['R'], 'm') / R0).to('').magnitude 

97 ts_z = (Quantity(thomson_position_['Z'], 'm') / R0).to('').magnitude * (-1.0 if equi.flipped_z else 1.0) 

98 

99 ts = Lineout(ts_r, ts_z) 

100 ts.find_points_on_grid(grid, n_points=n_points, trim_points=False) 

101 

102 return ts 

103 

104 

105def rdpa(grid, equi, omp_map, norm, rdpa_coordinates: Path): 

106 """ 

107 Returns a lineout for the 2D RDPA (reciprocating divertor probe array). Despite being a 2D diagnostic, 

108 since observables are not on a regular rectangular grid we calculate values as tricolumn (R^u-R^_sep, Z, values) 

109 data and then reshape. 

110 """ 

111 warnings.simplefilter('ignore', category=pint.errors.UnitStrippedWarning) 

112 

113 rdpa_coords = read_from_json(rdpa_coordinates) 

114 

115 R0 = norm.R0 

116 

117 rdpa_r = xr.DataArray(Quantity(rdpa_coords['R'] , rdpa_coords['R_units'] ) / R0).assign_attrs(norm=R0) 

118 rdpa_zx = xr.DataArray(Quantity(rdpa_coords['Zx'] , rdpa_coords['Zx_units'] ) / R0).assign_attrs(norm=R0) 

119 rdpa_rsep = xr.DataArray(Quantity(rdpa_coords['Rsep'], rdpa_coords['Rsep_units']) / R0).assign_attrs(norm=R0) 

120 rdpa_z = xr.DataArray(Quantity(rdpa_coords['Z'] , rdpa_coords['Z_units'] ) / R0).assign_attrs(norm=R0) 

121 

122 if equi.flipped_z: 

123 rdpa_z *= -1 

124 

125 # Some points cannot be matched in the reference equilibrium. These are marked by NaN 

126 nan_value = ~np.isnan(rdpa_r).values 

127 rdpa_on_grid = on_grid_interpolator(grid)(rdpa_r, rdpa_z, grid=False) > 0.99 

128 mask = np.logical_and.reduce((nan_value, rdpa_on_grid)) 

129 print(f"Masking {np.count_nonzero(~mask)} of {rdpa_r.size} RDPA points") 

130 rdpa_r, rdpa_z, rdpa_zx, rdpa_rsep = rdpa_r[mask], rdpa_z[mask], rdpa_zx[mask], rdpa_rsep[mask] 

131 

132 rdpa_rho = equi.normalised_flux_surface_label(rdpa_r, rdpa_z, grid=False) 

133 assert np.allclose(rdpa_rsep, omp_map.convert_rho_to_distance(rdpa_rho), atol=1E-6), \ 

134 f"Max difference for rdpa Rsep was {np.max(np.abs(rdpa_rsep - omp_map.convert_rho_to_distance(rdpa_rho)))}" 

135 

136 rdpa_lineout = Lineout(rdpa_r, rdpa_z) 

137 rdpa_lineout.setup_interpolation_matrix(grid, use_source_points=True) 

138 

139 # Store the observable positions to allow mapping to 2D data 

140 # We store according to machine coordinates, regardless of whether we have flipped the equilibrium 

141 rdpa_lineout.coords = { 

142 'R': convert_xarray_to_quantity(rdpa_r), 

143 'Rsep': convert_xarray_to_quantity(rdpa_rsep), 

144 'Z': convert_xarray_to_quantity(rdpa_z) * (-1.0 if equi.flipped_z else 1.0), 

145 'Zx': convert_xarray_to_quantity(rdpa_zx), 

146 } 

147 

148 return rdpa_lineout 

149 

150 

151def xpoint(grid, equi, norm, r_range: float = 0.1, z_range: float = 0.1, 

152 rhomin=0.97, rhomax=1.03, r_samples: int = 50, z_samples: int = 50): 

153 """ 

154 Returns a lineout of the region around the X-point 

155 """ 

156 warnings.simplefilter('ignore', category=pint.errors.UnitStrippedWarning) 

157 

158 r_min = float(equi.x_point_r_norm.values - r_range) 

159 r_max = float(equi.x_point_r_norm.values + r_range) 

160 z_min = float(equi.x_point_z_norm.values - z_range) 

161 z_max = float(equi.x_point_z_norm.values + z_range) 

162 

163 r_sample, z_sample = np.linspace(r_min, r_max, num=r_samples), np.linspace(z_min, z_max, num=z_samples) 

164 

165 xpoint_on_grid = (on_grid_interpolator(grid)(r_sample, z_sample, grid=True) > 0.99).T 

166 xpoint_rho = equi.normalised_flux_surface_label(r_sample, z_sample, grid=True).values 

167 

168 r_mesh, z_mesh = np.meshgrid(r_sample, z_sample) 

169 mask = np.logical_and.reduce((xpoint_on_grid, xpoint_rho > rhomin, xpoint_rho < rhomax)) 

170 

171 xpoint_lineout = Lineout(r_mesh[mask].ravel(), z_mesh[mask].ravel()) 

172 xpoint_lineout.setup_interpolation_matrix(grid, use_source_points=True) 

173 

174 xpoint_lineout.coords = { 

175 'R': xpoint_lineout.r_points * norm.R0, 

176 'Z': xpoint_lineout.z_points * norm.R0 * (-1.0 if equi.flipped_z else 1.0), 

177 'Zx': xpoint_lineout.z_points * norm.R0 * (-1.0 if equi.flipped_z else 1.0) \ 

178 - convert_xarray_to_quantity(equi.axis_r) 

179 } 

180 

181 return xpoint_lineout