Coverage for python/gsfit/gsfit.py: 94%

1import os 

2import time as time_py 

3import typing 

4 

5import gsfit_rs 

6import numpy as np 

7from diagnostic_and_simulation_base import DiagnosticAndSimulationBase 

8 

9from .database_readers import get_database_reader 

10from .database_writers import get_database_writer 

11 

12np.set_printoptions(linewidth=200) 

13 

14 

15class Gsfit(DiagnosticAndSimulationBase): 

16 """ 

17 GSFit: Grad-Shafranov Fit 

18 

19 Example usage: 

20 ```python 

21 # Note, this example will only run inside Tokamak Energy's network 

22 # This is a reproducton of `examples/example_01_st40_with_experimental_data.py` 

23 

24 from gsfit import Gsfit 

25 

26 pulseNo = 12050 # A real experimental shot 

27 pulseNo_write = pulseNo + 11_000_000 # Write to a "million" modelling pulse number 

28 

29 gsfit_controller = Gsfit( 

30 pulseNo=pulseNo, 

31 run_name="TEST01", 

32 run_description="Test run", 

33 write_to_mds=True, 

34 pulseNo_write=pulseNo_write, 

35 ) 

36 

37 gsfit_controller.run() 

38 ``` 

39 """ 

40 

41 coils: gsfit_rs.Coils 

42 passives: gsfit_rs.Passives 

43 plasma: gsfit_rs.Plasma 

44 bp_probes: gsfit_rs.BpProbes 

45 flux_loops: gsfit_rs.FluxLoops 

46 rogowski_coils: gsfit_rs.RogowskiCoils 

47 isoflux: gsfit_rs.Isoflux 

48 isoflux_boundary: gsfit_rs.IsofluxBoundary 

49 pressure_sensors: gsfit_rs.Pressure 

50 stationary_point: gsfit_rs.StationaryPoint 

51 dialoop: gsfit_rs.Dialoop 

52 

53 # TODO: move to DiagnosticAndSimulationBase 

54 def __getitem__(self, key: str) -> typing.Any: 

55 return self.results[key] 

56 

57 # TODO: move to DiagnosticAndSimulationBase 

58 def print_keys(self) -> None: 

59 """ 

60 Print the keys of the `self.results` nested dictionary, including subkeys. 

61 `self.results` is a 1:1 mapping to the MDSplus database structure. 

62 """ 

63 

64 self.results.print_keys() 

65 

66 # TODO: move to DiagnosticAndSimulationBase 

67 def keys(self, search: str | None = None) -> typing.KeysView[str]: 

68 """ 

69 Return the keys of the `self.results` nested dictionary, only the top level keys. 

70 `self.results` is a 1:1 mapping to the MDSplus database structure. 

71 """ 

72 

73 return self.results.keys() 

74 

75 def run(self, **kwargs: typing.Any) -> None: 

76 """ 

77 Run all components of GSFit 

78 (reading databases & initialisation, solving GS equation, and writing to data store). 

79 

80 :param kwargs: Additional arguments to be passed to the database_reader. This is for FreeGS and FreeGNSKE 

81 

82 This will perform the following steps: 

83 1. Set the environment variables 

84 2. Setup the timeslices to reconstruct 

85 3. Read in all the machine settings and initalise the following Rust implementations: 

86 `coils`, `passives`, `plasma`, `bp_probes`, `flux_loops`, `rogowski_coils`, `isoflux`, `isoflux_boundary`, and `stationary_point` 

87 4. Initialise the Greens functions 

88 5. Solve the GS equation 

89 6. Map the results to the MDSplus database structure and store in `self.results` 

90 7. Write the results to MDSplus 

91 """ 

92 

93 self.logger.info(f"Running Gsfit, for pulseNo={self.pulseNo}") 

94 

95 self.set_environment_variables() 

96 

97 self.setup_timeslices() 

98 

99 # Read in all the machine settings and initalise the following Rust implementations: 

100 # `coils`, `passives`, `plasma`, `bp_probes`, `flux_loops`, `rogowski_coils`, `isoflux`, `isoflux_boundary`, and `stationary_point` 

101 self.setup_objects(**kwargs) 

102 

103 # Calculate the Greens functions for all permutations between current source objects and sensors. 

104 self.calculate_greens() 

105 

106 # Solve the GS equation 

107 if self.settings["GSFIT_code_settings.json"]["type_of_run"]["inverse"]: 

108 self.inverse_solver_rust() 

109 elif self.settings["GSFIT_code_settings.json"]["type_of_run"]["forward"]: 

110 self.logger.warning("Forward GS solver not implemented yet. Skipping forward run!!") 

111 else: 

112 raise ValueError(f"Unknown type_of_run={self.settings['GSFIT_code_settings.json']['type_of_run']}") 

113 

114 self.write_results_to_mdsplus() 

115 

116 def write_results_to_mdsplus(self) -> None: 

117 """ 

118 Write the results to MDSplus: 

119 1. Results are collected from the Rust objects and stored in `self.results`,which is similar 

120 to a nested dictionary, and has a 1:1 mapping to the MDSplus database structure. 

121 2. The results are then written to MDSplus. 

122 """ 

123 

124 # Map the results to MDSplus. 

125 # `self.results` is a 1:1 mapping to MDSplus 

126 database_writer_method = self.settings["GSFIT_code_settings.json"]["database_writer"]["method"] 

127 database_writer = get_database_writer(database_writer_method) 

128 database_writer.map_results_to_database(self) 

129 

130 # Do the writing to MDSplus 

131 self.logger.info(f"pulseNo = {self.pulseNo} pulseNo_write = {self.pulseNo_write} run_name = {self.run_name}") 

132 if self.write_to_mds: 

133 self.logger.info("Writing to MDSplus") 

134 self._write_to_mds() 

135 

136 def setup_timeslices(self) -> None: 

137 """ 

138 Calculates the timeslices to reconstruct, and stores them in `self.results["TIME"]` 

139 """ 

140 

141 # Extract the timeslice settings from the JSON file 

142 timeslices_settings = self.settings["GSFIT_code_settings.json"]["timeslices"] 

143 

144 # Calculate the timeslices 

145 match timeslices_settings["method"]: 

146 case "arange": 

147 time = np.arange( 

148 timeslices_settings["arange"]["time_start"], 

149 timeslices_settings["arange"]["time_end"], 

150 timeslices_settings["arange"]["dt"], 

151 ) 

152 case "linspace": 

153 time = np.linspace( 

154 timeslices_settings["linspace"]["time_start"], 

155 timeslices_settings["linspace"]["time_end"], 

156 timeslices_settings["linspace"]["n_time"], 

157 ) 

158 case "user_defined": 

159 time = np.array(timeslices_settings["user_defined"]) 

160 case _: 

161 raise ValueError(f"Unknown timeslices method: {timeslices_settings['method']}") 

162 

163 # Store the times to reconstruct 

164 self.results["TIME"] = time 

165 

166 def set_environment_variables(self) -> None: 

167 """ 

168 Set the system environment variables before running. 

169 Presently, this defines the numer of cores for the Rust code 

170 """ 

171 

172 # Set the number of cores (for Rayon, Rust's parallelisation library) 

173 os.environ["RAYON_NUM_THREADS"] = str(self.settings["GSFIT_code_settings.json"]["RAYON_NUM_THREADS"]) 

174 

175 def inverse_solver_rust(self) -> None: 

176 """ 

177 Solve the "inverse" problem, i.e. reconstruction 

178 """ 

179 

180 # Extract objects from class 

181 coils = self.coils 

182 passives = self.passives 

183 plasma = self.plasma 

184 bp_probes = self.bp_probes 

185 flux_loops = self.flux_loops 

186 rogowski_coils = self.rogowski_coils 

187 isoflux = self.isoflux 

188 isoflux_boundary = self.isoflux_boundary 

189 pressure_sensors = self.pressure_sensors 

190 stationary_point = self.stationary_point 

191 dialoop = self.dialoop 

192 

193 times_to_reconstruct = self.results["TIME"] 

194 

195 self.logger.info(msg="About to call: `gsfit_rs.solve_grad_shafranov`") 

196 # Note: the solution to the GS equation is stored inside: `plasma`, `passives`, `bp_probes`, `flux_loops`, and `rogowski_coils` 

197 tic = time_py.time() 

198 gsfit_rs.solve_grad_shafranov( 

199 plasma, 

200 coils, 

201 passives, 

202 bp_probes, 

203 flux_loops, 

204 rogowski_coils, 

205 isoflux, 

206 isoflux_boundary, 

207 pressure_sensors, 

208 stationary_point, 

209 dialoop, 

210 times_to_reconstruct, 

211 self.settings["GSFIT_code_settings.json"]["numerics"]["n_iter_max"], 

212 self.settings["GSFIT_code_settings.json"]["numerics"]["n_iter_min"], 

213 self.settings["GSFIT_code_settings.json"]["numerics"]["n_iter_no_vertical_feedback"], 

214 self.settings["GSFIT_code_settings.json"]["numerics"]["gs_error"], 

215 self.settings["GSFIT_code_settings.json"]["numerics"]["anderson_mixing"]["use"], 

216 self.settings["GSFIT_code_settings.json"]["numerics"]["anderson_mixing"]["mixing_from_previous_iter"], 

217 ) 

218 toc = time_py.time() 

219 self.logger.info(msg=f"Finished: `gsfit_rs.solve_inverse_problem` time = {(toc - tic) * 1e3:,.2f}ms") 

220 

221 def calculate_greens(self) -> None: 

222 """ 

223 Calculates the Greens table for all permutations between current source objects and sensors. 

224 """ 

225 

226 # Get Rust classes out of self 

227 coils = self.coils 

228 passives = self.passives 

229 plasma = self.plasma 

230 bp_probes = self.bp_probes 

231 flux_loops = self.flux_loops 

232 rogowski_coils = self.rogowski_coils 

233 isoflux = self.isoflux 

234 isoflux_boundary = self.isoflux_boundary 

235 pressure_sensors = self.pressure_sensors 

236 stationary_point = self.stationary_point 

237 

238 # Greens with coils 

239 tic = time_py.time() 

240 plasma.greens_with_coils(coils) 

241 bp_probes.greens_with_coils(coils) 

242 flux_loops.greens_with_coils(coils) 

243 rogowski_coils.greens_with_coils(coils) 

244 isoflux.greens_with_coils(coils) 

245 isoflux_boundary.greens_with_coils(coils) 

246 pressure_sensors.greens_with_coils(coils) 

247 stationary_point.greens_with_coils(coils) 

248 toc = time_py.time() 

249 self.logger.info(f"Finished Greens with coils; {(toc - tic) * 1e3:,.2f}ms") 

250 

251 # Greens with passives 

252 tic = time_py.time() 

253 plasma.greens_with_passives(passives) 

254 bp_probes.greens_with_passives(passives) 

255 flux_loops.greens_with_passives(passives) 

256 rogowski_coils.greens_with_passives(passives) 

257 isoflux.greens_with_passives(passives) 

258 isoflux_boundary.greens_with_passives(passives) 

259 pressure_sensors.greens_with_passives(passives) 

260 stationary_point.greens_with_passives(passives) 

261 toc = time_py.time() 

262 self.logger.info(f"Finished Greens with passives; {(toc - tic) * 1e3:,.2f}ms") 

263 

264 # Greens with plasma 

265 tic = time_py.time() 

266 bp_probes.greens_with_plasma(plasma) 

267 flux_loops.greens_with_plasma(plasma) 

268 rogowski_coils.greens_with_plasma(plasma) 

269 isoflux.greens_with_plasma(plasma) 

270 isoflux_boundary.greens_with_plasma(plasma) 

271 pressure_sensors.greens_with_plasma(plasma) 

272 stationary_point.greens_with_plasma(plasma) 

273 toc = time_py.time() 

274 self.logger.info(f"Finished Greens with plasma; {(toc - tic) * 1e3:,.2f}ms") 

275 

276 def setup_objects(self, **kwargs: dict[str, typing.Any]) -> None: 

277 """ 

278 Initialises the Rust objects needed to run the GSFit inverse solver: 

279 `coils`, `passives`, `plasma`, `bp_probes`, `flux_loops`, `rogowski_coils`, `isoflux`, `isoflux_boundary`, and `stationary_point` 

280 

281 Different machines will use different data stores (e.g. MDSplus, or FreeGNSKE object). 

282 New readers for different devices / forward GS solvers can be added to: 

283 `python/gsfit/database_readers/__init__.py` and `python/gsfit/database_readers/<new_reader_name>` 

284 

285 See: `gsfit/database_readers/interface.py` for a description of the interfaces. 

286 """ 

287 

288 # Get the database_reader 

289 database_reader_method = self.settings["GSFIT_code_settings.json"]["database_reader"]["method"] 

290 database_reader = get_database_reader(database_reader_method) 

291 

292 # Initialise and store the Rust implementations 

293 tic = time_py.time() 

294 self.coils = database_reader.setup_coils(pulseNo=self.pulseNo, settings=self.settings, **kwargs) 

295 toc = time_py.time() 

296 self.logger.info(msg=f"`coils` initialised; {(toc - tic) * 1e3:,.2f}ms") 

297 

298 tic = time_py.time() 

299 self.bp_probes = database_reader.setup_bp_probes(pulseNo=self.pulseNo, settings=self.settings, **kwargs) 

300 toc = time_py.time() 

301 self.logger.info(msg=f"`bp_probes` initialised; {(toc - tic) * 1e3:,.2f}ms") 

302 

303 tic = time_py.time() 

304 self.flux_loops = database_reader.setup_flux_loops(pulseNo=self.pulseNo, settings=self.settings, **kwargs) 

305 toc = time_py.time() 

306 self.logger.info(msg=f"`flux_loops` initialised; {(toc - tic) * 1e3:,.2f}ms") 

307 

308 tic = time_py.time() 

309 self.rogowski_coils = database_reader.setup_rogowski_coils(pulseNo=self.pulseNo, settings=self.settings, **kwargs) 

310 toc = time_py.time() 

311 self.logger.info(msg=f"`rogowski_coils` initialised; {(toc - tic) * 1e3:,.2f}ms") 

312 

313 tic = time_py.time() 

314 self.passives = database_reader.setup_passives(pulseNo=self.pulseNo, settings=self.settings, **kwargs) 

315 toc = time_py.time() 

316 self.logger.info(msg=f"`passives` initialised; {(toc - tic) * 1e3:,.2f}ms") 

317 

318 tic = time_py.time() 

319 self.plasma = database_reader.setup_plasma(pulseNo=self.pulseNo, settings=self.settings, **kwargs) 

320 toc = time_py.time() 

321 self.logger.info(msg=f"`plasma` initialised; {(toc - tic) * 1e3:,.2f}ms") 

322 times_to_reconstruct = self.results["TIME"] 

323 

324 tic = time_py.time() 

325 self.isoflux = database_reader.setup_isoflux_sensors(pulseNo=self.pulseNo, settings=self.settings, times_to_reconstruct=times_to_reconstruct, **kwargs) 

326 toc = time_py.time() 

327 self.logger.info(msg=f"`isoflux` initialised; {(toc - tic) * 1e3:,.2f}ms") 

328 

329 tic = time_py.time() 

330 self.isoflux_boundary = database_reader.setup_isoflux_boundary_sensors(pulseNo=self.pulseNo, settings=self.settings, **kwargs) 

331 toc = time_py.time() 

332 self.logger.info(msg=f"`isoflux_boundary` initialised; {(toc - tic) * 1e3:,.2f}ms") 

333 

334 tic = time_py.time() 

335 self.pressure_sensors = database_reader.setup_pressure_sensors(pulseNo=self.pulseNo, settings=self.settings, **kwargs) 

336 toc = time_py.time() 

337 self.logger.info(msg=f"`pressure_sensors` initialised; {(toc - tic) * 1e3:,.2f}ms") 

338 

339 tic = time_py.time() 

340 self.stationary_point = database_reader.setup_stationary_point_sensors( 

341 pulseNo=self.pulseNo, 

342 settings=self.settings, 

343 times_to_reconstruct=times_to_reconstruct, 

344 **kwargs, 

345 ) 

346 toc = time_py.time() 

347 self.logger.info(msg=f"`stationary_point` initialised; {(toc - tic) * 1e3:,.2f}ms") 

348 

349 tic = time_py.time() 

350 self.dialoop = database_reader.setup_dialoop(pulseNo=self.pulseNo, settings=self.settings, **kwargs) 

351 toc = time_py.time() 

352 self.logger.info(msg=f"`dialoop` initialised; {(toc - tic) * 1e3:,.2f}ms")