Coverage for python/gsfit/gsfit.py: 87%
212 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-14 15:25 +0000
1import os
2import time as time_py
3import typing
5import gsfit_rs
6import numpy as np
7from diagnostic_and_simulation_base import DiagnosticAndSimulationBase
9from .database_readers import get_database_reader
10from .database_writers import get_database_writer
12if typing.TYPE_CHECKING:
13 from imas.ids_toplevel import IDSToplevel
15np.set_printoptions(linewidth=200)
18class Gsfit(DiagnosticAndSimulationBase):
19 """
20 GSFit: Grad-Shafranov Fit
22 Example usage:
23 ```python
24 # Note, this example will only run inside Tokamak Energy's network
25 # This is a reproducton of `examples/example_01_st40_with_experimental_data.py`
27 from gsfit import Gsfit
29 pulseNo = 12050 # A real experimental shot
30 pulseNo_write = pulseNo + 11_000_000 # Write to a "million" modelling pulse number
32 gsfit_controller = Gsfit(
33 pulseNo=pulseNo,
34 run_name="TEST01",
35 run_description="Test run",
36 write_to_mds=True,
37 pulseNo_write=pulseNo_write,
38 )
40 gsfit_controller.run()
41 ```
42 """
44 coils: gsfit_rs.Coils
45 passives: gsfit_rs.Passives
46 plasma: gsfit_rs.Plasma
47 bp_probes: gsfit_rs.BpProbes
48 flux_loops: gsfit_rs.FluxLoops
49 rogowski_coils: gsfit_rs.RogowskiCoils
50 isoflux: gsfit_rs.Isoflux
51 isoflux_boundary: gsfit_rs.IsofluxBoundary
52 pressure_sensors: gsfit_rs.Pressure
53 stationary_point: gsfit_rs.StationaryPoint
54 dialoop: gsfit_rs.Dialoop
55 tf: gsfit_rs.Tf
56 wall: gsfit_rs.Wall
58 # Optional: only a `database_reader` which reads gap definitions builds one. `None` means
59 # the equilibrium has no gaps
60 pulse_schedule: gsfit_rs.PulseSchedule | None = None
62 # Set by `write_results_to_database` when the `imas` database_writer is selected: the
63 # populated IMAS `equilibrium` IDS. `None` for every other writer
64 equilibrium_ids: "IDSToplevel | None" = None
66 # TODO: move to DiagnosticAndSimulationBase
67 def __getitem__(self, key: str) -> typing.Any:
68 return self.results[key]
70 # TODO: move to DiagnosticAndSimulationBase
71 def print_keys(self) -> None:
72 """
73 Print the keys of the `self.results` nested dictionary, including subkeys.
74 `self.results` is a 1:1 mapping to the MDSplus database structure.
75 """
77 self.results.print_keys()
79 # TODO: move to DiagnosticAndSimulationBase
80 def keys(self, search: str | None = None) -> typing.KeysView[str]:
81 """
82 Return the keys of the `self.results` nested dictionary, only the top level keys.
83 `self.results` is a 1:1 mapping to the MDSplus database structure.
84 """
86 return self.results.keys()
88 def run(self, **kwargs: typing.Any) -> None:
89 """
90 Run all components of GSFit
91 (reading databases & initialisation, solving GS equation, and writing to data store).
93 :param kwargs: Additional arguments to be passed to the database_reader. This is for FreeGS and FreeGNSKE
95 This will perform the following steps:
96 1. Start creating the MDSplus nodes, in a separate process
97 2. Set the environment variables
98 3. Setup the timeslices to reconstruct
99 4. Read in all the machine settings and initalise the following Rust implementations:
100 `coils`, `passives`, `plasma`, `tf`, `wall`, `bp_probes`, `flux_loops`, `rogowski_coils`, `isoflux`, `isoflux_boundary`, and `stationary_point`,
101 and the optional `pulse_schedule`
102 5. Initialise the Greens functions
103 6. Solve the GS equation
104 7. Map the results to the MDSplus database structure and store in `self.results`
105 8. Wait for the MDSplus nodes, then write the results into them
107 Steps 1 and 8 are skipped for the `imas` database_writer, which writes nothing.
108 """
110 self.logger.info(f"Running Gsfit, for pulseNo={self.pulseNo}")
112 # Creating the MDSplus nodes needs nothing but the settings, which have already been read,
113 # so start it now and let it run while we read the databases and solve the GS equation.
114 # The settings are the user's to change up to this point, so the values in force now are
115 # the ones the node creation is given; from here on they are fixed.
116 # `write_results_to_database` waits for it, and checks that it worked, before writing
117 if self.writes_to_mdsplus():
118 self.logger.info("Creating the MDSplus nodes in a separate process")
119 self.start_mds_node_creation(workflows=self.get_workflow_names())
121 self.set_environment_variables()
123 self.setup_timeslices()
125 # Read in all the machine settings and initalise the following Rust implementations:
126 # `coils`, `passives`, `plasma`, `tf`, `wall`, `bp_probes`, `flux_loops`, `rogowski_coils`, `isoflux`, `isoflux_boundary`, and `stationary_point`,
127 # and the optional `pulse_schedule`
128 self.setup_objects(**kwargs)
130 # Calculate the Greens functions for all permutations between current source objects and sensors.
131 self.calculate_greens()
133 # Solve the GS equation
134 if self.settings["GSFIT_code_settings.json"]["type_of_run"]["inverse"]:
135 self.inverse_solver_rust()
136 elif self.settings["GSFIT_code_settings.json"]["type_of_run"]["forward"]:
137 self.logger.warning("Forward GS solver not implemented yet. Skipping forward run!!")
138 else:
139 raise ValueError(f"Unknown type_of_run={self.settings['GSFIT_code_settings.json']['type_of_run']}")
141 self.write_results_to_database()
143 def writes_to_mdsplus(self) -> bool:
144 """
145 Whether this run will write to MDSplus at all.
147 `write_to_mds` is the user's switch, but the `imas` database_writer does not write to
148 MDSplus whatever it is set to: it builds an IMAS `equilibrium` IDS instead. Asking this
149 rather than `write_to_mds` keeps the MDSplus node creation, which starts before anything
150 else in `run`, in step with the writing that happens at the end.
151 """
153 database_writer_method = self.settings["GSFIT_code_settings.json"]["database_writer"]["method"]
155 return self.write_to_mds and database_writer_method != "imas"
157 def write_results_to_database(self) -> None:
158 """
159 Write the results to the database:
160 1. Results are collected from the Rust objects and stored in `self.results`, which is
161 similar to a nested dictionary, and has a 1:1 mapping to the MDSplus database structure.
162 2. We wait for the MDSplus nodes, which `run` started creating in a separate process, and
163 check that the creation did not fail. If `run` was not used, the nodes are created here.
164 3. The results are then written to MDSplus.
166 The `imas` database_writer does neither 2 nor 3. It returns an IMAS `equilibrium` IDS,
167 which is kept on `self.equilibrium_ids` for the caller to use, and nothing is written.
168 """
170 # Map the results to the database structure.
171 # For the MDSplus writers `self.results` is a 1:1 mapping to MDSplus
172 database_writer_method = self.settings["GSFIT_code_settings.json"]["database_writer"]["method"]
173 database_writer = get_database_writer(database_writer_method)
174 if database_writer_method != "imas":
175 database_writer.map_results_to_database(self)
176 else:
177 self.equilibrium_ids = database_writer.map_results_to_database(self)
178 self.logger.info("IMAS IDS populated; access it via `self.equilibrium_ids`")
180 # Do the writing to MDSplus
181 self.logger.info(f"pulseNo = {self.pulseNo} pulseNo_write = {self.pulseNo_write} run_name = {self.run_name}")
182 if self.writes_to_mdsplus():
183 # The nodes must exist, and must have been created without error, before we write into them
184 self.wait_for_mds_node_creation(workflows=self.get_workflow_names())
186 self.logger.info("Writing to database")
187 self._write_data_to_mds()
188 if database_writer_method == "tokamak_energy_mdsplus_new":
189 from .database_writers.tokamak_energy_mdsplus_new.create_mdsplus_links import create_mdsplus_links
191 create_mdsplus_links(pulseNo_write=self.pulseNo_write, run_name=self.run_name, tree_name=self.analysis_name)
193 def get_workflow_names(self) -> list[str]:
194 """
195 The names of the input codes which will be stored under `INPUT.WORKFLOW`,
196 for example `["elmag", "mag", "psu2coil"]`.
198 Which codes these are depends on the `database_reader`, and whether they are stored at all
199 depends on the `database_writer`, so the selected `database_writer` is asked. The answer
200 comes from the settings alone, which means the `INPUT.WORKFLOW` MDSplus nodes can be
201 created before GSFit has read any data.
202 """
204 database_writer_method = self.settings["GSFIT_code_settings.json"]["database_writer"]["method"]
205 database_writer = get_database_writer(database_writer_method)
207 return database_writer.get_workflow_names(self)
209 def setup_timeslices(self) -> None:
210 """
211 Calculates the timeslices to reconstruct, and stores them in `self.results["TIME"]`
212 """
214 # Extract the timeslice settings from the JSON file
215 timeslices_settings = self.settings["GSFIT_code_settings.json"]["timeslices"]
217 # Calculate the timeslices
218 match timeslices_settings["method"]:
219 case "arange":
220 time = np.arange(
221 timeslices_settings["arange"]["time_start"],
222 timeslices_settings["arange"]["time_end"],
223 timeslices_settings["arange"]["dt"],
224 )
225 case "linspace":
226 time = np.linspace(
227 timeslices_settings["linspace"]["time_start"],
228 timeslices_settings["linspace"]["time_end"],
229 timeslices_settings["linspace"]["n_time"],
230 )
231 case "user_defined":
232 time = np.array(timeslices_settings["user_defined"])
233 case "good_pressure_sensors":
234 # Reconstruct only the time-slices where the pressure sensors are "good".
235 # The `BAD_MA` reading lives in the `st40_mdsplus_with_digital_filter` reader.
236 from .database_readers.st40_mdsplus_with_digital_filter.setup_pressure_sensors import get_good_pressure_sensor_times
238 time = get_good_pressure_sensor_times(pulseNo=self.pulseNo, settings=self.settings)
239 case _:
240 raise ValueError(f"Unknown timeslices method: {timeslices_settings['method']}")
242 # Store the times to reconstruct
243 self.results["TIME"] = time
245 def set_environment_variables(self) -> None:
246 """
247 Set the system environment variables before running.
248 Presently, this defines the numer of cores for the Rust code
249 """
251 # Set the number of cores (for Rayon, Rust's parallelisation library)
252 n_rayon_threads: int = self.settings["GSFIT_code_settings.json"]["RAYON_NUM_THREADS"]
253 os.environ["RAYON_NUM_THREADS"] = str(n_rayon_threads)
254 self.logger.info(msg=f"RAYON_NUM_THREADS = {n_rayon_threads}")
256 def inverse_solver_rust(self) -> None:
257 """
258 Solve the "inverse" problem, i.e. reconstruction
259 """
261 # Extract objects from class
262 coils = self.coils
263 passives = self.passives
264 plasma = self.plasma
265 tf = self.tf
266 wall = self.wall
267 bp_probes = self.bp_probes
268 flux_loops = self.flux_loops
269 rogowski_coils = self.rogowski_coils
270 isoflux = self.isoflux
271 isoflux_boundary = self.isoflux_boundary
272 pressure_sensors = self.pressure_sensors
273 stationary_point = self.stationary_point
274 dialoop = self.dialoop
275 pulse_schedule = self.pulse_schedule
277 self.logger.info(msg="About to call: `gsfit_rs.solve_grad_shafranov`")
278 # Note: the solution to the GS equation is stored inside: `plasma`, `passives`, `bp_probes`, `flux_loops`, and `rogowski_coils`
279 tic = time_py.time()
280 gsfit_rs.solve_grad_shafranov(
281 plasma,
282 wall,
283 tf,
284 coils,
285 passives,
286 bp_probes,
287 flux_loops,
288 rogowski_coils,
289 isoflux,
290 isoflux_boundary,
291 pressure_sensors,
292 stationary_point,
293 dialoop,
294 pulse_schedule=pulse_schedule,
295 )
296 toc = time_py.time()
297 self.logger.info(msg=f"Finished: `gsfit_rs.solve_inverse_problem` time = {(toc - tic) * 1e3:,.2f}ms")
299 def calculate_greens(self) -> None:
300 """
301 Calculates the Greens table for all permutations between current source objects and sensors.
302 """
304 # Get Rust classes out of self
305 coils = self.coils
306 passives = self.passives
307 plasma = self.plasma
308 bp_probes = self.bp_probes
309 flux_loops = self.flux_loops
310 rogowski_coils = self.rogowski_coils
311 isoflux = self.isoflux
312 isoflux_boundary = self.isoflux_boundary
313 pressure_sensors = self.pressure_sensors
314 stationary_point = self.stationary_point
316 # Greens with coils
317 tic = time_py.time()
318 plasma.greens_with_coils(coils)
319 bp_probes.greens_with_coils(coils)
320 flux_loops.greens_with_coils(coils)
321 rogowski_coils.greens_with_coils(coils)
322 isoflux.greens_with_coils(coils)
323 isoflux_boundary.greens_with_coils(coils)
324 pressure_sensors.greens_with_coils(coils)
325 stationary_point.greens_with_coils(coils)
326 toc = time_py.time()
327 self.logger.info(f"Finished Greens with coils; {(toc - tic) * 1e3:,.2f}ms")
329 # Greens with passives
330 tic = time_py.time()
331 plasma.greens_with_passives(passives)
332 bp_probes.greens_with_passives(passives)
333 flux_loops.greens_with_passives(passives)
334 rogowski_coils.greens_with_passives(passives)
335 isoflux.greens_with_passives(passives)
336 isoflux_boundary.greens_with_passives(passives)
337 pressure_sensors.greens_with_passives(passives)
338 stationary_point.greens_with_passives(passives)
339 toc = time_py.time()
340 self.logger.info(f"Finished Greens with passives; {(toc - tic) * 1e3:,.2f}ms")
342 # Greens with plasma
343 tic = time_py.time()
344 bp_probes.greens_with_plasma(plasma)
345 flux_loops.greens_with_plasma(plasma)
346 rogowski_coils.greens_with_plasma(plasma)
347 isoflux.greens_with_plasma(plasma)
348 isoflux_boundary.greens_with_plasma(plasma)
349 pressure_sensors.greens_with_plasma(plasma)
350 stationary_point.greens_with_plasma(plasma)
351 toc = time_py.time()
352 self.logger.info(f"Finished Greens with plasma; {(toc - tic) * 1e3:,.2f}ms")
354 def setup_objects(self, **kwargs: dict[str, typing.Any]) -> None:
355 """
356 Initialises the Rust objects needed to run the GSFit inverse solver:
357 `coils`, `passives`, `plasma`, `tf`, `wall`, `bp_probes`, `flux_loops`, `rogowski_coils`, `isoflux`, `isoflux_boundary`, and `stationary_point`,
358 and the optional `pulse_schedule`, which is `None` unless the `database_reader` implements `setup_pulse_schedule`
360 Different machines will use different data stores (e.g. MDSplus, or FreeGNSKE object).
361 New readers for different devices / forward GS solvers can be added to:
362 `python/gsfit/database_readers/__init__.py` and `python/gsfit/database_readers/<new_reader_name>`
364 See: `gsfit/database_readers/interface.py` for a description of the interfaces.
365 """
367 # Get the database_reader
368 database_reader_method = self.settings["GSFIT_code_settings.json"]["database_reader"]["method"]
369 database_reader = get_database_reader(database_reader_method)
371 # Initialise and store the Rust implementations
372 tic = time_py.time()
373 self.coils = database_reader.setup_coils(pulseNo=self.pulseNo, settings=self.settings, **kwargs)
374 toc = time_py.time()
375 self.logger.info(msg=f"`coils` initialised; {(toc - tic) * 1e3:,.2f}ms")
377 tic = time_py.time()
378 self.bp_probes = database_reader.setup_bp_probes(pulseNo=self.pulseNo, settings=self.settings, **kwargs)
379 toc = time_py.time()
380 self.logger.info(msg=f"`bp_probes` initialised; {(toc - tic) * 1e3:,.2f}ms")
382 tic = time_py.time()
383 self.flux_loops = database_reader.setup_flux_loops(pulseNo=self.pulseNo, settings=self.settings, **kwargs)
384 toc = time_py.time()
385 self.logger.info(msg=f"`flux_loops` initialised; {(toc - tic) * 1e3:,.2f}ms")
387 tic = time_py.time()
388 self.rogowski_coils = database_reader.setup_rogowski_coils(pulseNo=self.pulseNo, settings=self.settings, **kwargs)
389 toc = time_py.time()
390 self.logger.info(msg=f"`rogowski_coils` initialised; {(toc - tic) * 1e3:,.2f}ms")
392 tic = time_py.time()
393 self.passives = database_reader.setup_passives(pulseNo=self.pulseNo, settings=self.settings, **kwargs)
394 toc = time_py.time()
395 self.logger.info(msg=f"`passives` initialised; {(toc - tic) * 1e3:,.2f}ms")
397 times_to_reconstruct = self.results["TIME"]
399 tic = time_py.time()
400 self.plasma = database_reader.setup_plasma(
401 pulseNo=self.pulseNo, settings=self.settings, times_to_reconstruct=times_to_reconstruct, **kwargs
402 )
403 toc = time_py.time()
404 self.logger.info(msg=f"`plasma` initialised; {(toc - tic) * 1e3:,.2f}ms")
406 tic = time_py.time()
407 self.wall = database_reader.setup_wall(pulseNo=self.pulseNo, settings=self.settings, **kwargs)
408 toc = time_py.time()
409 self.logger.info(msg=f"`wall` initialised; {(toc - tic) * 1e3:,.2f}ms")
411 tic = time_py.time()
412 self.tf = database_reader.setup_tf(pulseNo=self.pulseNo, settings=self.settings, **kwargs)
413 toc = time_py.time()
414 self.logger.info(msg=f"`tf` initialised; {(toc - tic) * 1e3:,.2f}ms")
416 tic = time_py.time()
417 self.isoflux = database_reader.setup_isoflux_sensors(pulseNo=self.pulseNo, settings=self.settings, times_to_reconstruct=times_to_reconstruct, **kwargs)
418 toc = time_py.time()
419 self.logger.info(msg=f"`isoflux` initialised; {(toc - tic) * 1e3:,.2f}ms")
421 tic = time_py.time()
422 self.isoflux_boundary = database_reader.setup_isoflux_boundary_sensors(pulseNo=self.pulseNo, settings=self.settings, **kwargs)
423 toc = time_py.time()
424 self.logger.info(msg=f"`isoflux_boundary` initialised; {(toc - tic) * 1e3:,.2f}ms")
426 tic = time_py.time()
427 self.pressure_sensors = database_reader.setup_pressure_sensors(pulseNo=self.pulseNo, settings=self.settings, **kwargs)
428 toc = time_py.time()
429 self.logger.info(msg=f"`pressure_sensors` initialised; {(toc - tic) * 1e3:,.2f}ms")
431 tic = time_py.time()
432 self.stationary_point = database_reader.setup_stationary_point_sensors(
433 pulseNo=self.pulseNo,
434 settings=self.settings,
435 times_to_reconstruct=times_to_reconstruct,
436 **kwargs,
437 )
438 toc = time_py.time()
439 self.logger.info(msg=f"`stationary_point` initialised; {(toc - tic) * 1e3:,.2f}ms")
441 tic = time_py.time()
442 self.dialoop = database_reader.setup_dialoop(pulseNo=self.pulseNo, settings=self.settings, **kwargs)
443 toc = time_py.time()
444 self.logger.info(msg=f"`dialoop` initialised; {(toc - tic) * 1e3:,.2f}ms")
446 tic = time_py.time()
447 self.pulse_schedule = database_reader.setup_pulse_schedule(pulseNo=self.pulseNo, settings=self.settings, **kwargs)
448 toc = time_py.time()
449 if self.pulse_schedule is None:
450 self.logger.info(msg=f"`pulse_schedule` not used by this database_reader, so there are no gaps; {(toc - tic) * 1e3:,.2f}ms")
451 else:
452 self.logger.info(msg=f"`pulse_schedule` initialised; {(toc - tic) * 1e3:,.2f}ms")