Coverage for python/gsfit/database_readers/st40_mdsplus_with_digital_filter/setup_pressure_sensors.py: 0%
50 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-14 15:25 +0000
1import typing
2from typing import TYPE_CHECKING
4import mdsthin
5import numpy as np
6import numpy.typing as npt
7from gsfit_rs import Pressure
9if TYPE_CHECKING:
10 from . import DatabaseReader
13def _get_workflow_settings(
14 settings: dict[str, typing.Any],
15) -> dict[str, typing.Any]:
16 """
17 Return the `workflow` section for the active `database_reader` method.
19 The pressure sensor set-up is shared between the `st40_mdsplus` and `st40_mdsplus_with_digital_filter`
20 readers, so the method-specific `workflow` section (which holds the PPTS and TS `run_name`s) is looked
21 up from the active `method` rather than being hard-coded.
23 :param settings: Dictionary containing the JSON settings read from the `settings` directory
24 """
26 database_reader_settings: dict[str, typing.Any] = settings["GSFIT_code_settings.json"]["database_reader"]
27 method: str = database_reader_settings["method"]
28 workflow: dict[str, typing.Any] = database_reader_settings[method]["workflow"]
30 return workflow
33def _read_good_time_slices(
34 pulseNo: int,
35 settings: dict[str, typing.Any],
36) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float64]]:
37 """
38 Find the "good" time-slices for the pressure sensors using the `BAD_MA` flag in the PPTS tree.
40 :param pulseNo: Pulse number, used to read from the database
41 :param settings: Dictionary containing the JSON settings read from the `settings` directory
43 :return: A tuple of `(good_indices, good_times)` where `good_indices` are the indices into the PPTS
44 time vector where the magnetics are considered "good" (i.e. `BAD_MA` is `False`), and `good_times`
45 are the corresponding times [second].
47 The PPTS `run_name` is read from the active method's `workflow` section in `GSFIT_code_settings.json`.
49 **This method is specific to ST40's experimental MDSplus database.**
50 """
52 ppts_run_name: str = _get_workflow_settings(settings)["ppts"]["run_name"]
54 with mdsthin.Connection("smaug") as conn:
55 conn.openTree("PPTS", pulseNo)
56 time_vector = np.asarray(conn.get(f"\\PPTS::TOP.{ppts_run_name}:TIME").data(), dtype=np.float64)
57 bad_ma = np.asarray(conn.get(f"\\PPTS::TOP.{ppts_run_name}.GLOBAL:BAD_MA").data(), dtype=bool)
59 good_indices = np.where(~bad_ma)[0].astype(np.int64)
60 good_times = time_vector[good_indices]
62 return good_indices, good_times
65def get_good_pressure_sensor_times(
66 pulseNo: int,
67 settings: dict[str, typing.Any],
68) -> npt.NDArray[np.float64]:
69 """
70 Return the times [second] at which the pressure sensors (Thomson scattering) are considered "good".
72 This is used by `Gsfit.setup_timeslices` when `timeslices.method == "good_pressure_sensors"`, so that
73 GSFit only reconstructs the time-slices where good pressure data exists.
75 :param pulseNo: Pulse number, used to read from the database
76 :param settings: Dictionary containing the JSON settings read from the `settings` directory
78 :return: The "good" times [second].
80 **This method is specific to ST40's experimental MDSplus database.**
81 """
83 _, good_times = _read_good_time_slices(pulseNo, settings)
85 return good_times
88def setup_pressure_sensors(
89 self: "DatabaseReader",
90 pulseNo: int,
91 settings: dict[str, typing.Any],
92) -> Pressure:
93 """
94 This method initialises the Rust `Pressure` class using ST40's Thomson scattering (TS) data.
96 :param pulseNo: Pulse number, used to read from the database
97 :param settings: Dictionary containing the JSON settings read from the `settings` directory
99 The pressure sensors are only added when `sensor_weights_pressure.json["include"]` is `True`.
100 The measured pressure is the electron pressure from the TS `PROFILES` node, scaled by
101 `pressure_multiplier` (typically 2.0, to convert electron pressure into total pressure).
102 Measurements below `minimum_pressure` are set to NaN so that GSFit ignores them.
104 The PPTS and TS `run_name`s are read from the active method's `workflow` section in
105 `GSFIT_code_settings.json` (the MDSplus node paths are built from these run names).
107 The sensors are added on the "good" time-base (where the `BAD_MA` flag in the PPTS tree is `False`).
108 To reconstruct only these time-slices, set `timeslices.method == "good_pressure_sensors"` in
109 `GSFIT_code_settings.json`.
111 **This method is specific to ST40's experimental MDSplus database.**
113 See `python/gsfit/database_readers/interface.py` for more details on how a new database_reader should be implemented.
114 """
116 # Initialise the Pressure Rust class
117 pressure = Pressure()
119 pressure_settings = settings["sensor_weights_pressure.json"]
121 # Pressure sensors are opt-in
122 if not pressure_settings.get("include", False):
123 return pressure
125 # Find the "good" time-slices using the `BAD_MA` flag in the PPTS tree
126 good_indices, good_times = _read_good_time_slices(pulseNo, settings)
128 # Read the Thomson scattering geometry and electron pressure profiles
129 thomson_scattering_settings = pressure_settings["thomson_scattering"]
130 pressure_multiplier = thomson_scattering_settings["pressure_multiplier"]
131 minimum_pressure = thomson_scattering_settings["minimum_pressure"]
133 # The TS `run_name` is read from the active method's workflow (e.g. "BEST")
134 ts_run_name: str = _get_workflow_settings(settings)["ts"]["run_name"]
136 with mdsthin.Connection("smaug") as conn:
137 conn.openTree("TS", pulseNo)
138 sensors_geometry_r = np.asarray(conn.get(f"\\TS::TOP.{ts_run_name}:R").data(), dtype=np.float64)
139 sensors_geometry_z = np.asarray(conn.get(f"\\TS::TOP.{ts_run_name}:Z").data(), dtype=np.float64)
140 # electron_pressure has shape = [n_time, n_sensors]
141 electron_pressure = np.asarray(conn.get(f"\\TS::TOP.{ts_run_name}.PROFILES:PE").data(), dtype=np.float64)
143 # Fit settings, shared by all Thomson scattering channels
144 sensor_name_prefix = pressure_settings["sensor_name_prefix"]
145 fit_settings_comment = pressure_settings["fit_settings"]["comment"]
146 fit_settings_expected_value = pressure_settings["fit_settings"]["expected_value"]
147 fit_settings_weight = pressure_settings["fit_settings"]["weight"]
149 n_sensors = len(sensors_geometry_r)
150 for i_sensor in range(n_sensors):
151 # Convert electron pressure into total pressure at the "good" time-slices
152 measured = pressure_multiplier * electron_pressure[good_indices, i_sensor]
153 measured = measured.copy()
155 # Ignore (set to NaN) any non-physical / below-threshold measurements
156 measured[measured <= minimum_pressure] = np.nan
158 # Only add the sensor if it has at least one valid measurement
159 if np.any(~np.isnan(measured)):
160 pressure.add_sensor(
161 name=f"{sensor_name_prefix}{i_sensor + 1:02d}",
162 geometry_r=sensors_geometry_r[i_sensor],
163 geometry_z=sensors_geometry_z[i_sensor],
164 fit_settings_comment=fit_settings_comment,
165 fit_settings_expected_value=fit_settings_expected_value,
166 fit_settings_include=True,
167 fit_settings_weight=fit_settings_weight,
168 time=good_times,
169 measured=measured,
170 )
172 return pressure