Coverage for python/gsfit/database_writers/tokamak_energy_mdsplus/map_results_to_database.py: 66%

1from typing import TYPE_CHECKING 

2 

3import numpy as np 

4 

5# from st40_database import GetData 

6 

7if TYPE_CHECKING: 

8 from ...gsfit import Gsfit 

9 from . import DatabaseWriterTokamakEnergyMDSplus 

10 

11 

12def map_results_to_database( 

13 self: "DatabaseWriterTokamakEnergyMDSplus", 

14 gsfit_controller: "Gsfit", 

15) -> None: 

16 """Map the results to MDSplus structure. 

17 `gsfit_controller.results` is a `NestedDict` object, which is a 1:1 mapping to the MDSplus structure. 

18 This function will mutate the `gsfit_controller` object. 

19 """ 

20 

21 # Take class object out of the `gsfit_controller` object 

22 pulseNo = gsfit_controller.pulseNo 

23 settings = gsfit_controller.settings 

24 plasma = gsfit_controller.plasma 

25 bp_probes = gsfit_controller.bp_probes 

26 flux_loops = gsfit_controller.flux_loops 

27 rogowski_coils = gsfit_controller.rogowski_coils 

28 passives = gsfit_controller.passives 

29 results = gsfit_controller.results 

30 pressure_sensors = gsfit_controller.pressure_sensors 

31 

32 # Two-d 

33 results["TWO_D"]["BR"] = plasma.get_array3(["two_d", "br"]) 

34 results["TWO_D"]["BT"] = plasma.get_array3(["two_d", "bt"]) 

35 results["TWO_D"]["BZ"] = plasma.get_array3(["two_d", "bz"]) 

36 results["TWO_D"]["MASK"] = plasma.get_array3(["two_d", "mask"]) 

37 results["TWO_D"]["P"] = plasma.get_array3(["two_d", "p"]) 

38 results["TWO_D"]["PSI"] = plasma.get_array3(["two_d", "psi"]) 

39 results["TWO_D"]["RGRID"] = plasma.get_array1(["grid", "r"]) 

40 results["TWO_D"]["ZGRID"] = plasma.get_array1(["grid", "z"]) 

41 

42 # Global 

43 results["GLOBAL"]["BETA_N"] = plasma.get_array1(["global", "beta_n"]) 

44 results["GLOBAL"]["BETA_P_1"] = plasma.get_array1(["global", "beta_p_1"]) 

45 results["GLOBAL"]["BETA_P_2"] = plasma.get_array1(["global", "beta_p_2"]) 

46 results["GLOBAL"]["BETA_P_3"] = plasma.get_array1(["global", "beta_p_3"]) 

47 results["GLOBAL"]["BETA_T"] = plasma.get_array1(["global", "beta_t"]) 

48 results["GLOBAL"]["BT_VAC_RGEO"] = plasma.get_array1(["global", "bt_vac_at_r_geo"]) 

49 results["GLOBAL"]["CHI_MAG"] = plasma.get_array1(["global", "chi_mag"]) 

50 results["GLOBAL"]["LI_1"] = plasma.get_array1(["global", "li_1"]) 

51 results["GLOBAL"]["LI_2"] = plasma.get_array1(["global", "li_2"]) 

52 results["GLOBAL"]["LI_3"] = plasma.get_array1(["global", "li_3"]) 

53 results["GLOBAL"]["DELTA_Z"] = plasma.get_array1(["global", "delta_z"]) 

54 results["GLOBAL"]["ELON"] = plasma.get_array1(["global", "elongation"]) 

55 results["GLOBAL"]["PHI_DIA"] = plasma.get_array1(["global", "phi_dia"]) 

56 results["GLOBAL"]["GS_ERROR"] = plasma.get_array1(["global", "gs_error"]) 

57 results["GLOBAL"]["I_ROD"] = plasma.get_array1(["global", "i_rod"]) 

58 results["GLOBAL"]["IP"] = plasma.get_array1(["global", "ip"]) 

59 results["GLOBAL"]["N_ITER"] = np.array(plasma.get_vec_usize(["global", "n_iter"])).astype(np.int32) 

60 results["GLOBAL"]["P"] = plasma.get_array1(["global", "p"]) 

61 results["GLOBAL"]["PSI_A"] = plasma.get_array1(["global", "psi_a"]) 

62 results["GLOBAL"]["PSI_B"] = plasma.get_array1(["global", "psi_b"]) 

63 results["GLOBAL"]["Q0"] = plasma.get_array1(["global", "q0"]) 

64 results["GLOBAL"]["Q95"] = plasma.get_array1(["global", "q95"]) 

65 results["GLOBAL"]["R_CUR"] = plasma.get_array1(["global", "r_cur"]) 

66 results["GLOBAL"]["Z_CUR"] = plasma.get_array1(["global", "z_cur"]) 

67 results["GLOBAL"]["R_GEO"] = plasma.get_array1(["global", "r_geo"]) 

68 results["GLOBAL"]["Z_GEO"] = plasma.get_array1(["global", "z_geo"]) 

69 results["GLOBAL"]["R_MAG"] = plasma.get_array1(["global", "r_mag"]) 

70 results["GLOBAL"]["Z_MAG"] = plasma.get_array1(["global", "z_mag"]) 

71 results["GLOBAL"]["R_MINOR"] = plasma.get_array1(["global", "r_minor"]) 

72 results["GLOBAL"]["V_LOOP"] = plasma.get_array1(["global", "v_loop"]) 

73 results["GLOBAL"]["VPLASMA"] = plasma.get_array1(["global", "plasma_volume"]) 

74 results["GLOBAL"]["W_MHD"] = plasma.get_array1(["global", "w_mhd"]) 

75 results["GLOBAL"]["XPT_DIVERTED"] = np.array(plasma.get_vec_bool(["global", "xpt_diverted"])).astype(np.int32) 

76 

77 # Bp probes (note, this is all the sensors, both the ones we fit and the ones we don't) 

78 bp_names = bp_probes.keys() # list of strings; len(bp_names) = n_sensors 

79 bp_names = [bp_name.replace("P", "B_BPPROBE_") for bp_name in bp_names] 

80 results["CONSTRAINTS"]["BPPROBE"]["NAME"] = np.array(bp_names) # MDSplus requires numpy objects, not lists of strings; shape = [n_sensors] 

81 results["CONSTRAINTS"]["BPPROBE"]["CVALUE"] = bp_probes.get_array2(["*", "b", "calculated", "value"]) # shape = [n_time, n_sensors] 

82 results["CONSTRAINTS"]["BPPROBE"]["INCLUDE"] = np.array(bp_probes.get_vec_bool(["*", "fit_settings", "include"])).astype(np.int32) 

83 results["CONSTRAINTS"]["BPPROBE"]["MVALUE"] = bp_probes.get_array2(["*", "b", "measured", "value"]) # shape = [n_time, n_sensors] 

84 results["CONSTRAINTS"]["BPPROBE"]["WEIGHT"] = bp_probes.get_array1(["*", "fit_settings", "weight"]) # shape = [n_sensors] 

85 

86 # Flux loops (note, this is all the sensors, both the ones we fit and the ones we don't) 

87 fl_names = flux_loops.keys() # list of strings; len(fl_names) = n_sensors 

88 fl_names = [fl_name.replace("L", "PSI_FLOOP_") for fl_name in fl_names] 

89 results["CONSTRAINTS"]["FLOOP"]["NAME"] = np.array(fl_names) # MDSplus requires numpy objects, not lists of strings; shape = [n_sensors] 

90 results["CONSTRAINTS"]["FLOOP"]["CVALUE"] = flux_loops.get_array2(["*", "psi", "calculated", "value"]) # shape = [n_time, n_sensors] 

91 results["CONSTRAINTS"]["FLOOP"]["INCLUDE"] = np.array(flux_loops.get_vec_bool(["*", "fit_settings", "include"])).astype(np.int32) 

92 results["CONSTRAINTS"]["FLOOP"]["MVALUE"] = flux_loops.get_array2(["*", "psi", "measured", "value"]) # shape = [n_time, n_sensors] 

93 results["CONSTRAINTS"]["FLOOP"]["WEIGHT"] = flux_loops.get_array1(["*", "fit_settings", "weight"]) # shape = [n_sensors] 

94 

95 # Rogowski coils (note, this is all the sensors, both the ones we fit and the ones we don't) 

96 rog_names = rogowski_coils.keys() # list of strings; len(rog_names) = n_sensors 

97 rog_names = [f"I_ROG_{rog_name}" for rog_name in rog_names] 

98 results["CONSTRAINTS"]["ROG"]["NAME"] = np.array(rog_names) # MDSplus requires numpy objects, not lists of strings; shape = [n_sensors] 

99 results["CONSTRAINTS"]["ROG"]["CVALUE"] = rogowski_coils.get_array2(["*", "i", "calculated", "value"]) # shape = [n_time, n_sensors] 

100 results["CONSTRAINTS"]["ROG"]["INCLUDE"] = np.array(rogowski_coils.get_vec_bool(["*", "fit_settings", "include"])).astype(np.int32) 

101 results["CONSTRAINTS"]["ROG"]["MVALUE"] = rogowski_coils.get_array2(["*", "i", "measured", "value"]) # shape = [n_time, n_sensors] 

102 results["CONSTRAINTS"]["ROG"]["WEIGHT"] = rogowski_coils.get_array1(["*", "fit_settings", "weight"]) # shape = [n_sensors] 

103 

104 # Plasma boundary 

105 results["P_BOUNDARY"]["NBND"] = np.array(plasma.get_vec_usize(["p_boundary", "nbnd"])) 

106 results["P_BOUNDARY"]["RBND"] = plasma.get_array2(["p_boundary", "rbnd"]) 

107 results["P_BOUNDARY"]["ZBND"] = plasma.get_array2(["p_boundary", "zbnd"]) 

108 results["P_BOUNDARY"]["BOUNDING_R"] = plasma.get_array1(["p_boundary", "bounding_r"]) 

109 results["P_BOUNDARY"]["BOUNDING_Z"] = plasma.get_array1(["p_boundary", "bounding_z"]) 

110 

111 # X-points 

112 results["XPOINTS"]["UPPER"]["R"] = plasma.get_array1(["xpoints", "upper", "r"]) 

113 results["XPOINTS"]["UPPER"]["Z"] = plasma.get_array1(["xpoints", "upper", "z"]) 

114 results["XPOINTS"]["LOWER"]["R"] = plasma.get_array1(["xpoints", "lower", "r"]) 

115 results["XPOINTS"]["LOWER"]["Z"] = plasma.get_array1(["xpoints", "lower", "z"]) 

116 

117 # Profiles 

118 results["PROFILES"]["RHO"]["AREA"] = plasma.get_array2(["profiles", "area"]) 

119 results["PROFILES"]["RHO"]["AREA_PRIME"] = plasma.get_array2(["profiles", "area_prime"]) 

120 results["PROFILES"]["RHO"]["F"] = plasma.get_array2(["profiles", "f"]) 

121 results["PROFILES"]["RHO"]["FF_PRIME"] = plasma.get_array2(["profiles", "ff_prime"]) 

122 results["PROFILES"]["RHO"]["FLUX_TOR"] = plasma.get_array2(["profiles", "flux_tor"]) 

123 results["PROFILES"]["RHO"]["P"] = plasma.get_array2(["profiles", "p"]) 

124 results["PROFILES"]["RHO"]["P_PRIME"] = plasma.get_array2(["profiles", "p_prime"]) 

125 results["PROFILES"]["RHO"]["Q"] = plasma.get_array2(["profiles", "q"]) 

126 results["PROFILES"]["RHO"]["RHO_POL"] = plasma.get_array2(["profiles", "rho_pol"]) 

127 results["PROFILES"]["RHO"]["RHO_TOR"] = plasma.get_array2(["profiles", "rho_tor"]) 

128 results["PROFILES"]["RHO"]["PSI_N"] = plasma.get_array1(["profiles", "psi_n"]) 

129 results["PROFILES"]["RHO"]["VOL"] = plasma.get_array2(["profiles", "vol"]) 

130 results["PROFILES"]["RHO"]["VOL_PRIME"] = plasma.get_array2(["profiles", "vol_prime"]) 

131 

132 # Mid-plane profiles 

133 results["PROFILES"]["MID_PLANE"]["P"] = plasma.get_array2(["profiles", "mid_plane", "p"]) 

134 results["PROFILES"]["MID_PLANE"]["R"] = plasma.get_array1(["profiles", "mid_plane", "r"]) 

135 

136 # Passives 

137 for passive_name in passives.keys(): 

138 if passive_name == "IVC": 

139 results["PASSIVES"]["IVC"]["DOF"]["EIG_01"]["CVALUE"] = passives.get_array1(["IVC", "dof", "eig_01", "calculated"]) 

140 results["PASSIVES"]["IVC"]["DOF"]["EIG_01"]["I_DIST"] = passives.get_array1(["IVC", "dof", "eig_01", "current_distribution"]) 

141 results["PASSIVES"]["IVC"]["DOF"]["EIG_02"]["CVALUE"] = passives.get_array1(["IVC", "dof", "eig_02", "calculated"]) 

142 results["PASSIVES"]["IVC"]["DOF"]["EIG_02"]["I_DIST"] = passives.get_array1(["IVC", "dof", "eig_02", "current_distribution"]) 

143 results["PASSIVES"]["IVC"]["DOF"]["EIG_03"]["CVALUE"] = passives.get_array1(["IVC", "dof", "eig_03", "calculated"]) 

144 results["PASSIVES"]["IVC"]["DOF"]["EIG_03"]["I_DIST"] = passives.get_array1(["IVC", "dof", "eig_03", "current_distribution"]) 

145 results["PASSIVES"]["IVC"]["DOF"]["EIG_04"]["CVALUE"] = passives.get_array1(["IVC", "dof", "eig_04", "calculated"]) 

146 results["PASSIVES"]["IVC"]["DOF"]["EIG_04"]["I_DIST"] = passives.get_array1(["IVC", "dof", "eig_04", "current_distribution"]) 

147 results["PASSIVES"]["IVC"]["DOF"]["EIG_05"]["CVALUE"] = passives.get_array1(["IVC", "dof", "eig_05", "calculated"]) 

148 results["PASSIVES"]["IVC"]["DOF"]["EIG_05"]["I_DIST"] = passives.get_array1(["IVC", "dof", "eig_05", "current_distribution"]) 

149 results["PASSIVES"]["IVC"]["DOF"]["EIG_06"]["CVALUE"] = passives.get_array1(["IVC", "dof", "eig_06", "calculated"]) 

150 results["PASSIVES"]["IVC"]["DOF"]["EIG_06"]["I_DIST"] = passives.get_array1(["IVC", "dof", "eig_06", "current_distribution"]) 

151 results["PASSIVES"]["IVC"]["DOF"]["EIG_07"]["CVALUE"] = passives.get_array1(["IVC", "dof", "eig_07", "calculated"]) 

152 results["PASSIVES"]["IVC"]["DOF"]["EIG_07"]["I_DIST"] = passives.get_array1(["IVC", "dof", "eig_07", "current_distribution"]) 

153 results["PASSIVES"]["IVC"]["DOF"]["EIG_08"]["CVALUE"] = passives.get_array1(["IVC", "dof", "eig_08", "calculated"]) 

154 results["PASSIVES"]["IVC"]["DOF"]["EIG_08"]["I_DIST"] = passives.get_array1(["IVC", "dof", "eig_08", "current_distribution"]) 

155 results["PASSIVES"]["IVC"]["DOF"]["EIG_09"]["CVALUE"] = passives.get_array1(["IVC", "dof", "eig_09", "calculated"]) 

156 results["PASSIVES"]["IVC"]["DOF"]["EIG_09"]["I_DIST"] = passives.get_array1(["IVC", "dof", "eig_09", "current_distribution"]) 

157 results["PASSIVES"]["IVC"]["DOF"]["EIG_10"]["CVALUE"] = passives.get_array1(["IVC", "dof", "eig_10", "calculated"]) 

158 results["PASSIVES"]["IVC"]["DOF"]["EIG_10"]["I_DIST"] = passives.get_array1(["IVC", "dof", "eig_10", "current_distribution"]) 

159 results["PASSIVES"]["IVC"]["DOF"]["EIG_11"]["CVALUE"] = passives.get_array1(["IVC", "dof", "eig_11", "calculated"]) 

160 results["PASSIVES"]["IVC"]["DOF"]["EIG_11"]["I_DIST"] = passives.get_array1(["IVC", "dof", "eig_11", "current_distribution"]) 

161 results["PASSIVES"]["IVC"]["DOF"]["EIG_12"]["CVALUE"] = passives.get_array1(["IVC", "dof", "eig_12", "calculated"]) 

162 results["PASSIVES"]["IVC"]["DOF"]["EIG_12"]["I_DIST"] = passives.get_array1(["IVC", "dof", "eig_12", "current_distribution"]) 

163 results["PASSIVES"]["IVC"]["DOF"]["EIG_13"]["CVALUE"] = passives.get_array1(["IVC", "dof", "eig_13", "calculated"]) 

164 results["PASSIVES"]["IVC"]["DOF"]["EIG_13"]["I_DIST"] = passives.get_array1(["IVC", "dof", "eig_13", "current_distribution"]) 

165 results["PASSIVES"]["IVC"]["DOF"]["EIG_14"]["CVALUE"] = passives.get_array1(["IVC", "dof", "eig_14", "calculated"]) 

166 results["PASSIVES"]["IVC"]["DOF"]["EIG_14"]["I_DIST"] = passives.get_array1(["IVC", "dof", "eig_14", "current_distribution"]) 

167 results["PASSIVES"]["IVC"]["DOF"]["EIG_15"]["CVALUE"] = passives.get_array1(["IVC", "dof", "eig_15", "calculated"]) 

168 results["PASSIVES"]["IVC"]["DOF"]["EIG_15"]["I_DIST"] = passives.get_array1(["IVC", "dof", "eig_15", "current_distribution"]) 

169 results["PASSIVES"]["IVC"]["GEOMETRY"]["ANGLE_1"] = passives.get_array1(["IVC", "geometry", "angle_1"]) 

170 results["PASSIVES"]["IVC"]["GEOMETRY"]["ANGLE_2"] = passives.get_array1(["IVC", "geometry", "angle_2"]) 

171 results["PASSIVES"]["IVC"]["GEOMETRY"]["D_R"] = passives.get_array1(["IVC", "geometry", "d_r"]) 

172 results["PASSIVES"]["IVC"]["GEOMETRY"]["D_Z"] = passives.get_array1(["IVC", "geometry", "d_z"]) 

173 results["PASSIVES"]["IVC"]["GEOMETRY"]["R"] = passives.get_array1(["IVC", "geometry", "r"]) 

174 results["PASSIVES"]["IVC"]["GEOMETRY"]["Z"] = passives.get_array1(["IVC", "geometry", "z"]) 

175 else: 

176 results["PASSIVES"][passive_name]["DOF"]["CONSTANT_J"]["CVALUE"] = passives.get_array1( 

177 [passive_name, "dof", "constant_current_density", "calculated"] 

178 ) 

179 results["PASSIVES"][passive_name]["DOF"]["CONSTANT_J"]["I_DIST"] = passives.get_array1( 

180 [passive_name, "dof", "constant_current_density", "current_distribution"] 

181 ) 

182 results["PASSIVES"][passive_name]["GEOMETRY"]["ANGLE_1"] = passives.get_array1([passive_name, "geometry", "angle_1"]) 

183 results["PASSIVES"][passive_name]["GEOMETRY"]["ANGLE_2"] = passives.get_array1([passive_name, "geometry", "angle_2"]) 

184 results["PASSIVES"][passive_name]["GEOMETRY"]["D_R"] = passives.get_array1([passive_name, "geometry", "d_r"]) 

185 results["PASSIVES"][passive_name]["GEOMETRY"]["D_Z"] = passives.get_array1([passive_name, "geometry", "d_z"]) 

186 results["PASSIVES"][passive_name]["GEOMETRY"]["R"] = passives.get_array1([passive_name, "geometry", "r"]) 

187 results["PASSIVES"][passive_name]["GEOMETRY"]["Z"] = passives.get_array1([passive_name, "geometry", "z"]) 

188 

189 # Scrape off layer (SOL) 

190 results["SOL"]["HFS"]["CONTOUR"]["R"] = plasma.get_array2(["sol", "hfs", "contour", "r"]) # shape = [n_time, n_points] 

191 results["SOL"]["HFS"]["CONTOUR"]["Z"] = plasma.get_array2(["sol", "hfs", "contour", "z"]) # shape = [n_time, n_points] 

192 results["SOL"]["HFS"]["CONTOUR"]["N"] = np.array(plasma.get_vec_usize(["sol", "hfs", "contour", "n"])).astype(np.int32) # shape = [n_time] 

193 results["SOL"]["HFS"]["STRIKE_POINT"]["R"] = plasma.get_array1(["sol", "hfs", "strike_point", "r"]) # shape = [n_time] 

194 results["SOL"]["HFS"]["STRIKE_POINT"]["Z"] = plasma.get_array1(["sol", "hfs", "strike_point", "z"]) # shape = [n_time] 

195 results["SOL"]["LFS"]["CONTOUR"]["R"] = plasma.get_array2(["sol", "lfs", "contour", "r"]) # shape = [n_time, n_points] 

196 results["SOL"]["LFS"]["CONTOUR"]["Z"] = plasma.get_array2(["sol", "lfs", "contour", "z"]) # shape = [n_time, n_points] 

197 results["SOL"]["LFS"]["CONTOUR"]["N"] = np.array(plasma.get_vec_usize(["sol", "lfs", "contour", "n"])).astype(np.int32) # shape = [n_time] 

198 results["SOL"]["LFS"]["STRIKE_POINT"]["R"] = plasma.get_array1(["sol", "lfs", "strike_point", "r"]) # shape = [n_time] 

199 results["SOL"]["LFS"]["STRIKE_POINT"]["Z"] = plasma.get_array1(["sol", "lfs", "strike_point", "z"]) # shape = [n_time] 

200 

201 if len(pressure_sensors.keys()) > 0: 

202 results["CONSTRAINTS"]["PRESSURE"]["RECONSTRUCTED"] = pressure_sensors.get_array2(["*", "pressure", "calculated", "value"]) # shape = [n_time, n_points] 

203 results["CONSTRAINTS"]["PRESSURE"]["MEASURED"] = pressure_sensors.get_array2(["*", "pressure", "measured", "value"]) # shape = [n_time, n_points] 

204 results["CONSTRAINTS"]["PRESSURE"]["WEIGHT"] = pressure_sensors.get_array1(["*", "fit_settings", "weight"]) # shape = [n_points] 

205 results["CONSTRAINTS"]["PRESSURE"]["POSITION"]["R"] = pressure_sensors.get_array1(["*", "geometry", "r"]) # shape = [n_points] 

206 results["CONSTRAINTS"]["PRESSURE"]["POSITION"]["Z"] = pressure_sensors.get_array1(["*", "geometry", "z"]) # shape = [n_points] 

207 results["CONSTRAINTS"]["PRESSURE"]["POSITION"]["PSI"] = pressure_sensors.get_array2(["*", "pressure", "calculated", "psi"]) # shape = [n_time, n_points] 

208 

209 # Store "WORKFLOW" 

210 database_reader_method = settings["GSFIT_code_settings.json"]["database_reader"]["method"] 

211 

212 code_names = settings["GSFIT_code_settings.json"]["database_reader"][database_reader_method]["workflow"].keys() 

213 

214 for code_name in code_names: 

215 pulseNo_json = settings["GSFIT_code_settings.json"]["database_reader"][database_reader_method]["workflow"][code_name]["pulseNo"] 

216 if pulseNo_json is not None: 

217 results["INPUT"]["WORKFLOW"][code_name]["PULSE"] = pulseNo_json 

218 else: 

219 results["INPUT"]["WORKFLOW"][code_name]["PULSE"] = pulseNo 

220 

221 run_name = settings["GSFIT_code_settings.json"]["database_reader"][database_reader_method]["workflow"][code_name]["run_name"] 

222 results["INPUT"]["WORKFLOW"][code_name]["RUN"] = run_name 

223 

224 usage = settings["GSFIT_code_settings.json"]["database_reader"][database_reader_method]["workflow"][code_name]["usage"] 

225 results["INPUT"]["WORKFLOW"][code_name]["USAGE"] = usage