Coverage for python/gsfit/database_readers/interface.py: 100%
27 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 13:12 +0000
1import typing
2from typing import Protocol
4import numpy as np
5import numpy.typing as npt
6from gsfit_rs import BpProbes
7from gsfit_rs import Coils
8from gsfit_rs import Dialoop
9from gsfit_rs import FluxLoops
10from gsfit_rs import Isoflux
11from gsfit_rs import IsofluxBoundary
12from gsfit_rs import Passives
13from gsfit_rs import Plasma
14from gsfit_rs import Pressure
15from gsfit_rs import RogowskiCoils
16from gsfit_rs import StationaryPoint
19class DatabaseReaderProtocol(Protocol):
20 """
21 Protocol for reading experimental data.
22 Each method is responsible for initialising one of the Rust objects:
23 `bp_probes`, `coils`, `dialoop`, `flux_loops`, `isoflux`, `isoflux_boundary`, `stationary_point`, `passives`, `plasma`, and `rogowski_coils`.
25 The Protocol defines the inputs and outputs of each method.
26 New database readers should be implemented **all** methods.
27 """
29 def setup_bp_probes(self, pulseNo: int, settings: dict[str, typing.Any], **kwargs: dict[str, typing.Any]) -> BpProbes:
30 """
31 This method initialises the Rust `BpProbes` class (Mirnov coils with a direction in the poloidal plane).
33 :param pulseNo: Pulse number, used to read from the database
34 :param settings: Dictionary containing the JSON settings read from the `settings` directory
35 :param kwargs: Additional objects, such as FreeGNSKE object
37 Initialising requires reading data from two locations:
38 1. `sensor_weights_bp_probe.json`: Which contains "fitting parameters", e.g. if a probe should be used in the fitting
39 2. Database reading (e.g. MDSplus, or FreeGNSKE object): Which contains the measurements and probe geometry
41 Different machines will use different data stores for the probe geometry and measured signals.
42 This Protocol allows different database readers to be selected.
43 The output of this method must always be a `BpProbes` object.
45 At a minimum this method should look like this:
46 ```python
47 # Initialise the BpProbes Rust class
48 bp_probes = BpProbes()
50 # Add all of the BP probes
51 for i_bp_probe in range(n_bp_probes):
52 bp_probes.add_sensor(
53 name=..., # need to be same in database and `sensor_weights_bp_probe.json` file
54 geometry_angle_pol=..., # read from a database
55 geometry_r=..., # read from a database
56 geometry_z=..., # read from a database
57 fit_settings_comment=..., # read from `sensor_weights_bp_probe.json` file
58 fit_settings_expected_value=..., # read from `sensor_weights_bp_probe.json` file
59 fit_settings_include=..., # read from `sensor_weights_bp_probe.json` file
60 fit_settings_weight=..., # read from `sensor_weights_bp_probe.json` file
61 time=..., # read from a database
62 measured=..., # read from a database
63 )
65 return bp_probes
66 ```
67 """
68 ...
70 def setup_coils(self, pulseNo: int, settings: dict[str, typing.Any], **kwargs: dict[str, typing.Any]) -> Coils:
71 """
72 This method initialises the Rust `Coils` class, which contains both the poloidal fied (PF) and toroidal field (TF) coils.
74 :param pulseNo: Pulse number, used to read from the database
75 :param settings: Dictionary containing the JSON settings read from the `settings` directory
76 :param kwargs: Additional objects, such as FreeGSNKE object
78 Initialising requires reading data from:
79 1. Database reading (e.g. MDSplus, or FreeGSNKE object): Which contains the coil geometry and current measurements
81 Different machines will use different data stores for the coil locations and currents.
82 This Protocol allows different database readers to be selected.
83 The output of this method must always be a `Coils` object.
85 At a minimum this method should look like this:
86 ```python
87 # Initialise the Coils Rust class
88 coils = Coils()
90 # Add all of the PF coils
91 for i_pf_coil in range(n_pf_coils):
92 coils.add_pf_coil(
93 coil_name=..., # read from a database
94 coil_r=..., # read from a database
95 coil_z=..., # read from a database
96 coil_d_r=..., # read from a database
97 coil_d_z=..., # read from a database
98 time=..., # read from a database
99 measured=..., # read from a database
100 )
102 # Add TF coil
103 coils.add_tf_coil(
104 time=..., # read from a database
105 i_rod=..., # read from a database
106 )
108 return coils
109 ```
110 """
111 ...
113 def setup_dialoop(self, pulseNo: int, settings: dict[str, typing.Any], **kwargs: dict[str, typing.Any]) -> Dialoop:
114 """
115 This method initialises the Rust `Dialoop` class (the plasma's diamagnetic flux, with the vaccuum toroidal field subtracted).
117 :param pulseNo: Pulse number, used to read from the database
118 :param settings: Dictionary containing the JSON settings read from the `settings` directory
119 :param kwargs: Additional objects, such as FreeGNSKE object
121 Initialising requires reading data from two locations:
122 1. `sensor_weights_dialoop.json`: Which contains "fitting parameters", e.g. if a probe should be used in the fitting
123 2. Database reading (e.g. MDSplus, or FreeGNSKE object): Which contains the measurements and the probe geometry
125 Different machines will use different data stores for the measured signals.
126 This Protocol allows different database readers to be selected.
127 The output of this method must always be a `Dialoop` object.
129 At a minimum this method should look like this:
130 ```python
131 # Initialise the Dialoop Rust class
132 dialoop = Dialoop()
134 # Add all of the diamagnetic loops
135 for i_dialoop in range(n_dialoops):
136 dialoop.add_sensor(
137 name=..., # read from a database
138 fit_settings_comment=..., # read from `sensor_weights_dialoop.json` file
139 fit_settings_expected_value=..., # read from `sensor_weights_dialoop.json` file
140 fit_settings_include=..., # read from `sensor_weights_dialoop.json` file
141 fit_settings_weight=..., # read from `sensor_weights_dialoop.json` file
142 time=..., # read from a database
143 measured=..., # read from a database
144 )
146 return dialoop
147 ```
148 """
149 ...
151 def setup_flux_loops(self, pulseNo: int, settings: dict[str, typing.Any], **kwargs: dict[str, typing.Any]) -> FluxLoops:
152 """
153 This method initialises the Rust `FluxLoops` class (loops going in the toroidal direction measuring poloidal flux).
155 :param pulseNo: Pulse number, used to read from the database
156 :param settings: Dictionary containing the JSON settings read from the `settings` directory
157 :param kwargs: Additional objects, such as FreeGNSKE object
159 Initialising requires reading data from two locations:
160 1. `sensor_weights_flux_loops.json`: Which contains "fitting parameters", e.g. if a probe should be used in the fitting
161 2. Database reading (e.g. MDSplus, or FreeGNSKE object): Which contains the measurements and the probe geometry
163 Different machines will use different data stores for the probes geometry and measured signals.
164 This Protocol allows different database readers to be selected.
165 The output of this method must always be a `FluxLoops` object.
167 At a minimum this method should look like this:
168 ```python
169 # Initialise the FluxLoops Rust class
170 flux_loops = FluxLoops()
172 # Add all of the BP probes
173 for i_flux_loop in range(n_flux_loops):
174 flux_loops.add_sensor(
175 name=..., # need to be same in database and `sensor_weights_flux_loops.json` file
176 geometry_r=..., # read from a database
177 geometry_z=..., # read from a database
178 fit_settings_comment=..., # read from `sensor_weights_flux_loops.json` file
179 fit_settings_expected_value=..., # read from `sensor_weights_flux_loops.json` file
180 fit_settings_include=..., # read from `sensor_weights_flux_loops.json` file
181 fit_settings_weight=..., # read from `sensor_weights_flux_loops.json` file
182 time=..., # read from a database
183 measured=..., # read from a database
184 )
186 return flux_loops
187 ```
188 """
189 ...
191 def setup_isoflux_sensors(
192 self, pulseNo: int, settings: dict[str, typing.Any], times_to_reconstruct: npt.NDArray[np.float64], **kwargs: dict[str, typing.Any]
193 ) -> Isoflux:
194 """
195 This method initialises the Rust `Isoflux` class (two locations which have equal poloidal flux).
197 :param pulseNo: Pulse number, used to read from the database
198 :param settings: Dictionary containing the JSON settings read from the `settings` directory
199 :param times_to_reconstruct: Times to reconstruct the equilibrium
200 :param kwargs: Additional objects, such as FreeGNSKE object
202 :param pulseNo: Pulse number, used to read from the database
203 :param settings: Dictionary containing the JSON settings read from the `settings` directory
205 Initialising requires reading data from two locations:
206 1. `sensor_weights_isoflux.json`: Which contains "fitting parameters", e.g. if the constraint should be used in the fitting
207 2. Database reading (e.g. MDSplus, or FreeGNSKE object): Which contains the coordinates where the isoflux constraint is applied
209 Different machines will use different data stores for the isoflux coordiantes.
210 This Protocol allows different database readers to be selected.
211 The output of this method must always be a `Isoflux` object.
213 At a minimum this method should look like this:
214 ```python
215 # Initialise the Isoflux Rust class
216 isoflux = Isoflux()
218 # Add all of the isoflux constraints
219 for i_isoflux_constraint in range(n_isoflux_constraints):
220 isoflux.add_sensor(
221 name=..., # need to be same in database and `sensor_weights_isoflux.json` file
222 fit_settings_comment=..., # read from `sensor_weights_isoflux.json` file
223 fit_settings_include=..., # read from `sensor_weights_isoflux.json` file
224 fit_settings_weight=..., # read from `sensor_weights_isoflux.json` file
225 time=..., # read from a database
226 location_1_r=..., # read from a database
227 location_1_z=..., # read from a database
228 location_2_r=..., # read from a database
229 location_2_z=..., # read from a database
230 times_to_reconstruct=..., # read from `GSFIT_code_settings.json` file
231 )
233 return isoflux
234 ```
235 """
236 ...
238 def setup_isoflux_boundary_sensors(self, pulseNo: int, settings: dict[str, typing.Any], **kwargs: dict[str, typing.Any]) -> IsofluxBoundary:
239 """
240 This method initialises the Rust `IsofluxBoundary` class (a location which has the same poloidal flux as the plasma boundary).
242 :param pulseNo: Pulse number, used to read from the database
243 :param settings: Dictionary containing the JSON settings read from the `settings` directory
244 :param kwargs: Additional objects, such as FreeGNSKE object
246 Initialising requires reading data from two locations:
247 1. `sensor_weights_isoflux_boundary.json`: Which contains "fitting parameters", e.g. if the constraint should be used in the fitting
248 2. Database reading (e.g. MDSplus, or FreeGNSKE object): Which contains the coordinates where the isoflux boundary constraint is applied
250 Different machines will use different data stores for the isoflux coordiantes.
251 This Protocol allows different database readers to be selected.
252 The output of this method must always be a `IsofluxBoundary` object.
254 At a minimum this method should look like this:
255 ```python
256 # Initialise the IsofluxBoundary Rust class
257 isoflux_boundary = IsofluxBoundary()
259 # Add all of the isoflux_boundary constraints
260 for i_isoflux_boundary_constraint in range(n_isoflux_boundary_constraints):
261 isoflux_boundary.add_sensor(
262 name=..., # need to be same in database and `sensor_weights_isoflux_boundary.json` file
263 fit_settings_comment=..., # read from `sensor_weights_isoflux_boundary.json` file
264 fit_settings_include=..., # read from `sensor_weights_isoflux_boundary.json` file
265 fit_settings_weight=..., # read from `sensor_weights_isoflux_boundary.json` file
266 time=..., # read from a database
267 location_1_r=..., # read from a database
268 location_1_z=..., # read from a database
269 location_2_r=..., # read from a database
270 location_2_z=..., # read from a database
271 times_to_reconstruct=..., # read from `GSFIT_code_settings.json` file
272 )
274 return isoflux_boundary
275 ```
276 """
277 ...
279 def setup_stationary_point_sensors(self, pulseNo: int, settings: dict[str, typing.Any], **kwargs: dict[str, typing.Any]) -> StationaryPoint:
280 """
281 This method initialises the Rust `StationaryPoint` class (the magnetic axis, and x-point positions).
283 :param pulseNo: Pulse number, used to read from the database
284 :param settings: Dictionary containing the JSON settings read from the `settings` directory
285 :param kwargs: Additional objects, such as FreeGNSKE object
287 Initialising requires reading data from two locations:
288 1. `sensor_weights_magnetic_axis.json`: Which contains "fitting parameters", e.g. if a probe should be used in the fitting
289 2. Database reading (e.g. MDSplus, or FreeGNSKE object): Which contains the measurements and the probe geometry
291 Different machines will use different data stores for the measured signals.
292 This Protocol allows different database readers to be selected.
293 The output of this method must always be a `StationaryPoint` object.
295 At a minimum this method should look like this:
296 ```python
297 # Initialise the StationaryPoint Rust class
298 stationary_point = StationaryPoint()
300 # Add all of the magnetic axis sensors
301 for i_magnetic_axis in range(n_magnetic_axis_sensors):
302 stationary_point.add_sensor(
303 name=..., # read from a database
304 fit_settings_comment=..., # read from `sensor_weights_magnetic_axis.json` file
305 fit_settings_expected_value=..., # read from `sensor_weights_magnetic_axis.json` file
306 fit_settings_include=..., # read from `sensor_weights_magnetic_axis.json` file
307 fit_settings_weight=..., # read from `sensor_weights_magnetic_axis.json` file
308 time=..., # read from a database
309 mag_axis_r=..., # read from a database
310 mag_axis_z=..., # read from a database
311 times_to_reconstruct=..., # read from `GSFIT_code_settings.json` file
312 )
314 return stationary_point
315 ```
316 """
317 ...
319 def setup_passives(self, pulseNo: int, settings: dict[str, typing.Any], **kwargs: dict[str, typing.Any]) -> Passives:
320 """
321 This method initialises the Rust `Passives` class (toroidal conductors).
323 :param pulseNo: Pulse number, used to read from the database
324 :param settings: Dictionary containing the JSON settings read from the `settings` directory
325 :param kwargs: Additional objects, such as FreeGNSKE object
327 Initialising requires reading data from two locations:
328 1. `passive_dof_regularisation.json`: Which specifies how the conductor should be represented (e.g. constant current density, or eigenmode decomposition)
329 2. Database reading (e.g. MDSplus, or FreeGNSKE object): Which contains the coordinates where the passive conductors are located
331 Different machines will use different data stores for the passive conductor coordiantes.
332 This Protocol allows different database readers to be selected.
333 The output of this method must always be a `Passives` object.
335 At a minimum this method should look like this:
336 ```python
337 # Initialise the Passives Rust class
338 passives = Passives()
340 # Add all of the passive toroidal conductors
341 for i_passives in range(n_passives):
342 passives.add_passive(
343 name=..., # need to be same in database and `passive_dof_regularisation.json` file
344 r=..., # read from a database
345 z=..., # read from a database
346 d_r=..., # read from a database
347 d_z=..., # read from a database
348 angle_1=..., # read from a database
349 angle_2=..., # read from a database
350 resistivity=..., # read from a database
351 current_distribution_type=..., # read from `passive_dof_regularisation.json` file
352 n_dof=..., # read from `passive_dof_regularisation.json` file
353 regularisations=..., # read from `passive_dof_regularisation.json` file
354 regularisations_weight=..., # read from `passive_dof_regularisation.json` file
355 )
357 return passives
358 ```
359 """
360 ...
362 def setup_plasma(self, pulseNo: int, settings: dict[str, typing.Any], **kwargs: dict[str, typing.Any]) -> Plasma:
363 """
364 This method initialises the Rust `Plasma` class.
366 :param pulseNo: Pulse number, used to read from the database
367 :param settings: Dictionary containing the JSON settings read from the `settings` directory
368 :param kwargs: Additional objects, such as FreeGNSKE object
370 Initialising requires reading data from three locations:
371 1. `GSFIT_code_settings.json`: Which contains the plasma grid size and the maximum number of iterations
372 2. `source_function_p_prime.json`: Which contains the number of degrees of freedom for p_prime, and regularisation
373 3. `source_function_ff_prime.json`: Which contains the number of degrees of freedom for ff_prime, and regularisation
375 Normally, this method will be the same for all machines.
377 The output of this method must always be a `Plasma` object.
379 At a minimum this method should look like this:
380 ```python
381 # Initialise the Plasma Rust class
382 p_prime_source_function = gsfit_rs.EfitPolynomial(
383 n_dof=..., # read from `source_function_p_prime.json` file
384 regularisations=..., # read from `source_function_p_prime.json` file
385 )
386 ff_prime_source_function = gsfit_rs.EfitPolynomial(
387 n_dof=..., # read from `source_function_ff_prime.json` file
388 regularisations=..., # read from `source_function_ff_prime.json` file
389 )
391 # Initialise the Plasma Rust class
392 plasma = Plasma(
393 n_r=..., # read from `GSFIT_code_settings.json` file
394 n_z=..., # read from `GSFIT_code_settings.json` file
395 r_min=..., # read from `GSFIT_code_settings.json` file
396 r_max=..., # read from `GSFIT_code_settings.json` file
397 z_min=..., # read from `GSFIT_code_settings.json` file
398 z_max=..., # read from `GSFIT_code_settings.json` file
399 psi_n=..., # read from `GSFIT_code_settings.json` file
400 limit_pts_r=..., # read from `GSFIT_code_settings.json` file
401 limit_pts_z=..., # read from `GSFIT_code_settings.json` file
402 vessel_r=..., # read from `GSFIT_code_settings.json` file
403 vessel_z=..., # read from `GSFIT_code_settings.json` file
404 p_prime_source_function=p_prime_source_function, # built above
405 ff_prime_source_function=ff_prime_source_function, # built above
406 )
408 return plasma
409 ```
410 """
411 ...
413 def setup_rogowski_coils(self, pulseNo: int, settings: dict[str, typing.Any], **kwargs: dict[str, typing.Any]) -> RogowskiCoils:
414 """
415 This method initialises the Rust `RogowskiCoils` class.
417 :param pulseNo: Pulse number, used to read from the database
418 :param settings: Dictionary containing the JSON settings read from the `settings` directory
419 :param kwargs: Additional objects, such as FreeGNSKE object
421 Initialising requires reading data from two locations:
422 1. `sensor_weights_rogowski_coils.json`: Which contains "fitting parameters", e.g. if a probe should be used in the fitting
423 2. Database reading (e.g. MDSplus, or FreeGNSKE object): Which contains the measurements and the Rogowski coil path
425 Different machines will use different data stores for the probes geometry and measured signals.
426 This Protocol allows different database readers to be selected.
427 The output of this method must always be a `RogowskiCoils` object.
429 At a minimum this method should look like this:
430 ```python
431 # Initialise the RogowskiCoils Rust class
432 rogowski_coils = RogowskiCoils()
434 # Add all of the BP probes
435 for i_rogowski_coil in range(n_rogowski_coils):
436 rogowski_coils.add_sensor(
437 sensor_name=..., # need to be same in database and `sensor_weights_flux_loops.json` file
438 path_r, # read from a database
439 path_z, # read from a database
440 fit_settings_comment, # read from `sensor_weights_flux_loops.json` file
441 fit_settings_expected_value, # read from `sensor_weights_flux_loops.json` file
442 fit_settings_include, # read from `sensor_weights_flux_loops.json` file
443 fit_settings_weight, # read from `sensor_weights_flux_loops.json` file
444 time, # read from a database
445 measured, # read from a database
446 gaps_r=gaps_r, # read from a database
447 gaps_z=gaps_z, # read from a database
448 gaps_d_r=gaps_d_r, # read from a database
449 gaps_d_z=gaps_d_z, # read from a database
450 gaps_name=gaps_name, # read from a database
451 )
453 return rogowski_coils
454 ```
455 """
456 ...
458 def setup_pressure_sensors(self, pulseNo: int, settings: dict[str, typing.Any], **kwargs: dict[str, typing.Any]) -> Pressure:
459 """
460 This method initialises the Rust `Pressure` class (pressure measurements at specific locations in the poloidal plane).
462 :param pulseNo: Pulse number, used to read from the database
463 :param settings: Dictionary containing the JSON settings read from the `settings` directory
464 :param kwargs: Additional objects, such as FreeGNSKE object
466 Initialising requires reading data from two locations:
467 1. `sensor_weights_pressure.json`: Which contains "fitting parameters", e.g. if a sensor should be used in the fitting
468 2. Database reading (e.g. MDSplus, or FreeGNSKE object): Which contains the measurements and sensor geometry
470 Different machines will use different data stores for the sensor geometry and measured signals.
471 This Protocol allows different database readers to be selected.
472 The output of this method must always be a `Pressure` object.
474 At a minimum this method should look like this:
475 ```python
476 # Initialise the Pressure Rust class
477 pressure = Pressure()
479 # Add all of the pressure sensors
480 for i_pressure_sensor in range(n_pressure_sensors):
481 pressure.add_sensor(
482 name=..., # need to be same in database and `sensor_weights_pressure.json` file
483 geometry_r=..., # read from a database
484 geometry_z=..., # read from a database
485 fit_settings_comment=..., # read from `sensor_weights_pressure.json` file
486 fit_settings_expected_value=..., # read from `sensor_weights_pressure.json` file
487 fit_settings_include=..., # read from `sensor_weights_pressure.json` file
488 fit_settings_weight=..., # read from `sensor_weights_pressure.json` file
489 time=..., # read from a database
490 measured=..., # read from a database
491 )
493 return pressure
494 ```
495 """
496 ...