Coverage for python/gsfit/database_writers/imas/map_results_to

1"""

2Translate the GSFit results into an IMAS `equilibrium` IDS, using the official `imas-python`

3package.

5**How it works.** Since the IMAS migration, GSFit already stores its equilibrium results in the

6IMAS data model: `plasma` holds an `equilibrium` IDS on the Rust side, read in place with `plasma.get`,

7and `gsfit_rs.imas.equilibrium_paths` is a path builder over the same schema. So this module does not

8hand-map several hundred leaves. It **walks the Rust schema and copies each leaf into the node of

9the same name** in the official IDS.

11Each time-slice is traversed on its own, reading `time_slice[i_time]` rather than the accumulator

12`time_slice[:]`. So every leaf comes back exactly as that time-slice stores it: a leaf whose length

13differs between time-slices, such as the boundary outline, keeps its own length, and there is no

14padding to remove.

16That matters for maintenance: a leaf added on the Rust side appears in the output automatically,

17and a leaf the official data dictionary does not have is skipped rather than crashing the run.

18Which leaves those are depends on the data dictionary in use, so they are logged at `debug` level

19at the end of every run rather than written down here.

21**Data retrieval only.** No physics, no unit conversions, no derived quantities. The only

22re-shaping is:

24* a leaf that is entirely non-finite is left unset, rather than written as `NaN`, because an unset

25 node is how IMAS represents "not calculated";

26* the 2D maps are transposed if needed so that `dim1` is R and `dim2` is Z, checked against the

27 lengths of `grid/dim1` and `grid/dim2` rather than assumed;

28* `bool` and `usize` are cast to the `int` the data dictionary expects.

30**Nothing is written to a database.** The IDS is returned, and `Gsfit.write_results_to_database`

31keeps it on `gsfit_controller.equilibrium_ids`. Selecting this writer also stops GSFit writing to

32MDSplus at all, whatever `write_to_mds` is set to - see `Gsfit.writes_to_mdsplus`.

34Conventions to be aware of

35--------------------------

361. **COCOS.** The IMAS data dictionary fixes its own convention: DD major version 4 is COCOS 17

37 (`<cocos>17</cocos>` in `IDSDef.xml`; DD 3 was COCOS 11). There is no field in which to declare

38 a different one, so the data has to be *in* COCOS 17.

40 GSFit is intended to be COCOS 13. Using the Sauter & Medvedev (2013) table:

43 | -------- | ------ | ---------- | ------------- | ------------------- |

44 | COCOS 13 | 1 | -1 | +1 | -1 |

45 | COCOS 17 | 1 | -1 | +1 | +1 |

47 They agree on everything except `sigma_rhothetaphi`, the direction of the poloidal angle. Both

48 have `e_Bp = 1`, so `psi` is the *total* poloidal flux [weber] in both and there is **no factor

49 of 2 * pi**; both have `sigma_Bp = -1`, so `psi`, `ip`, `f`, `phi` and the current densities

50 need no sign change either.

52 What does change is anything whose sign follows the poloidal angle direction: the DD's `q_like`

53 (`.fact_q`) and `pol_angle_like` (`.fact_dtheta`) groups. The affected leaves written here are

54 `time_slice/profiles_1d/q`, `global_quantities/q_axis`, `global_quantities/q_95` and

55 `time_slice/profiles_2d/theta`.

56 TODO (rust): flip the sign of `q` when going to COCOS 17, or store `q` in the IMAS convention

57 directly. It is not done here, because negating a number is a calculation.

59 `profiles_2d/theta` is the exception that is *already* in the IMAS convention: `epp_theta_2d`

60 in `rust/gsfit_rs/src/plasma.rs` computes `theta = -atan2(z - z_mag, r - r_mag)`, the clockwise

61 (`sigma_rhothetaphi = +1`) sense COCOS 17 requires, so it passes through unchanged. That makes

62 `theta` inconsistent with GSFit's COCOS 13 intent for everything else. If GSFit settles on

63 COCOS 13 throughout, `epp_theta_2d` is where the sense should be flipped, and `q` and `theta`

64 would then be handled the same way here.

66 Note: `imas-python` only implements COCOS handling as part of its DD3 -> DD4 conversion (a sign

67 flip on the `psi_like` / `dodpsi_like` groups); there is no general "convert from COCOS N" entry

68 point to lean on.

692. **Constraint ordering.** The IMAS constraint arrays are positional and are meant to line up

70 with the `magnetics` / `pf_active` / `pf_passive` IDSs. Those IDSs are not written here, so the

71 GSFit sensor name is stored in `source` to keep the mapping recoverable.

723. **Rogowski coils are not all passive.** GSFit keeps every Rogowski coil in one `rogowski_coils`

73 object, but on ST40 they measure two different things: the `*WIRE` coils (`BVLWIRE`, `BVUBWIRE`,

74 `BVUTWIRE`, `DIVBWIRE`, `DIVTWIRE`, `MCWIRE`, `PSHBWIRE`, `PSHTWIRE`, `SOLWIRE`, `TFWIRE`)

75 measure the current in the *active* coil windings, while the rest (coil cases, `INIVC*`,

76 `HFSPSR*`, `DIVPSR*`, `GASBFL*`) measure *passive* structure currents. The whole family is

77 written to `constraints/pf_passive_current`, because splitting it here would mean inferring the

78 sensor type from its name. `source` holds the GSFit name, so the split stays possible.

79 TODO (rust): tag each Rogowski coil with what it measures, then the `*WIRE` coils can be routed

80 to `constraints/pf_current` instead.

814. **Missing sensor values are passed through as-is.** Sensors GSFit did not fit carry

82 `fit_settings/weight = NAN` (14 of the 27 ST40 Rogowski coils), un-measured channels carry `0.0`

83 rather than `NAN` (e.g. flux loop `L003`), and `DIALOOP` has `measured = NAN` on some pulses.

84 None of it is cleaned up here.

85"""

87import typing

88from typing import TYPE_CHECKING

90import gsfit_rs.imas as gsfit_imas

91import imas

92import imas.ids_defs

93import numpy as np

94import numpy.typing as npt

95from diagnostic_and_simulation_base import version_storage

97if TYPE_CHECKING:

98 import gsfit_rs

99 from imas.ids_toplevel import IDSToplevel

100

101 from ...gsfit import Gsfit

102 from . import DatabaseWriterIMAS

103

104equilibrium_paths = gsfit_imas.equilibrium_paths

105

106

107def _is_leaf(path: typing.Any) -> bool:

108 """A path node holding data, as opposed to a structure or an array of structures.

109

110 Both a leaf and an array of structures have no children in `dir()`. They are told apart by

111 `data_type`, which an array of structures reports as an empty string.

112 """

113

114 if [name for name in dir(path) if not name.startswith("_")]:

115 return False

116 try:

117 return bool(path.data_type)

118 except Exception:

119 return False

120

121

122def _is_array_of_structures(path: typing.Any) -> bool:

123 """A path node which cannot be read until it is indexed."""

124

125 return not [name for name in dir(path) if not name.startswith("_")] and not _is_leaf(path)

126

127

128def _read_leaf(plasma: "gsfit_rs.Plasma", path: typing.Any) -> typing.Any:

129 """Read one leaf out of the Rust IDS held by `plasma`, or `None` if its path does not exist.

130

131 Every path read here selects a single element of each array of structures on it, e.g.

132 `time_slice[i_time]` and `profiles_2d[0]`, so the leaf comes back exactly as it is stored in that

133 element, at its own length. None of the accumulators (`time_slice[:]`) is used, because gathering

134 elements of different lengths into one array has to pad them.

135

136 A leaf GSFit never set reads back as its IMAS empty value - `NaN`, `EMPTY_INT`, an empty string

137 or an empty array - and `_write_leaf` leaves such a node unset in the output. The read itself

138 only fails when an index is beyond the end of an array of structures, e.g. `profiles_2d[0]` on

139 a time-slice which has no `profiles_2d`.

140 """

141

142 try:

143 return plasma.get(path)

144 except IndexError:

145 return None

146

147

148def _write_leaf(destination: typing.Any, name: str, values: typing.Any, data_type: str) -> None:

149 """Write one leaf into the official IDS, doing the minimum reshaping the DD requires."""

150

151 if data_type.startswith("STR"):

152 if values is not None and str(values) != "":

153 setattr(destination, name, str(values))

154 return

155

156 array = np.asarray(values)

157

158 if data_type.startswith("INT"):

159 # `EMPTY_INT` is how an unset integer reads back, e.g. from a time-slice which did not

160 # converge; an unset integer array is empty

161 if array.ndim == 0:

162 if int(array) != imas.ids_defs.EMPTY_INT:

163 setattr(destination, name, int(array))

164 elif array.size > 0:

165 setattr(destination, name, array.astype(np.int32))

166 return

167

168 array = array.astype(np.float64)

169 if not np.isfinite(array).any():

170 # Nothing was calculated; an unset node says that, a node full of NaN does not

171 return

172

173 setattr(destination, name, array)

174

175

176def _copy_branch(

177 source: typing.Any,

178 destination: typing.Any,

179 plasma: "gsfit_rs.Plasma",

180 path: str,

181 skipped: set[str],

182) -> None:

183 """Traverse one branch of the Rust IDS, and fill the node of the same name in the official IDS.

184

185 :param source: a path into the Rust IDS which selects single elements only, such as

186 `equilibrium_paths.time_slice[i_time]` or `equilibrium_paths.vacuum_toroidal_field`

187 :param destination: the matching node of the official IDS

188 :param path: the IMAS path of `source`, used to report the leaves which are skipped

189 :param skipped: the leaves the data dictionary has no node for, added to in place

190 """

191

192 for name in sorted(n for n in dir(source) if not n.startswith("_")):

193 source_child = getattr(source, name)

194 if not isinstance(source_child, gsfit_imas.Path):

195 continue

196 child_path = f"{path}/{name}"

197

198 if not hasattr(destination, name):

199 skipped.add(child_path)

200 continue

201 destination_child = getattr(destination, name)

202

203 if _is_array_of_structures(source_child):

204 # `profiles_2d` and `boundary/gap` are the only arrays of structures this copier fills.

205 # The constraint arrays are filled from the sensor objects further down, and the rest

206 # (`contour_tree/node`, `ggd`) GSFit does not store. None of them is a schema mismatch,

207 # so none is reported as one

208 if name == "profiles_2d":

209 _copy_profiles_2d(source_child, destination_child, plasma, child_path, skipped)

210 elif name == "gap":

211 _copy_gaps(source_child, destination_child, plasma, child_path, skipped)

212 continue

213

214 if _is_leaf(source_child):

215 values = _read_leaf(plasma, source_child)

216 if values is not None:

217 _write_leaf(destination, name, values, source_child.data_type)

218 continue

219

220 _copy_branch(source_child, destination_child, plasma, child_path, skipped)

221

222

223def _copy_profiles_2d(

224 source: typing.Any,

225 destination: typing.Any,

226 plasma: "gsfit_rs.Plasma",

227 path: str,

228 skipped: set[str],

229) -> None:

230 """Copy `profiles_2d[0]`, orienting the 2D maps as the data dictionary wants them."""

231

232 source_entry = source[0]

233

234 dim1 = _read_leaf(plasma, source_entry.grid.dim1)

235 dim2 = _read_leaf(plasma, source_entry.grid.dim2)

236 if dim1 is None or dim2 is None:

237 # This time-slice has no `profiles_2d`

238 return

239 n_dim1: int = len(dim1)

240 n_dim2: int = len(dim2)

241

242 destination.resize(1)

243 _copy_profiles_2d_branch(source_entry, destination[0], plasma, path, skipped, n_dim1, n_dim2)

244

245

246def _copy_gaps(

247 source: typing.Any,

248 destination: typing.Any,

249 plasma: "gsfit_rs.Plasma",

250 path: str,

251 skipped: set[str],

252) -> None:

253 """Copy `boundary/gap`, one element at a time.

254

255 There are only gaps when `solve_grad_shafranov` was given a pulse schedule defining some, so an

256 empty array is the usual case and leaves the node empty. Every time-slice carries the same gaps,

257 whether or not it converged; one which did not has each `value` unset.

258 """

259

260 gap_names = _read_leaf(plasma, source[:].name)

261 if gap_names is None:

262 return

263 n_gaps: int = len(gap_names)

264 if n_gaps == 0:

265 return

266

267 destination.resize(n_gaps)

268 for i_gap in range(n_gaps):

269 _copy_branch(source[i_gap], destination[i_gap], plasma, path, skipped)

270

271

272def _copy_profiles_2d_branch(

273 source: typing.Any,

274 destination: typing.Any,

275 plasma: "gsfit_rs.Plasma",

276 path: str,

277 skipped: set[str],

278 n_dim1: int,

279 n_dim2: int,

280) -> None:

281 """Copy one branch of `profiles_2d[0]`, orienting every 2D map it holds against the grid.

282

283 Recursive, so that the 2D maps in a sub-structure, such as `grid/volume_element`, are oriented

284 too, not only the direct children of `profiles_2d`.

285

286 :param n_dim1: the length of `grid/dim1`, the major radius [count]

287 :param n_dim2: the length of `grid/dim2`, the height [count]

288 """

289

290 for name in sorted(n for n in dir(source) if not n.startswith("_")):

291 source_child = getattr(source, name)

292 if not isinstance(source_child, gsfit_imas.Path):

293 continue

294 child_path = f"{path}/{name}"

295

296 if not hasattr(destination, name):

297 skipped.add(child_path)

298 continue

299

300 if not _is_leaf(source_child):

301 if _is_array_of_structures(source_child):

302 continue

303 _copy_profiles_2d_branch(source_child, getattr(destination, name), plasma, child_path, skipped, n_dim1, n_dim2)

304 continue

305

306 raw = _read_leaf(plasma, source_child)

307 if raw is None:

308 continue

309 values = np.asarray(raw)

310

311 # An unset map, e.g. on a time-slice which did not converge, is empty; there is nothing to

312 # orient, and `_write_leaf` leaves it unset

313 if values.size == 0:

314 continue

315

316 # The DD requires (dim1, dim2) = (R, Z). Checked against the grid rather than assumed,

317 # so that a change of storage order on the Rust side is caught instead of silently

318 # transposing the wrong way

319 if values.ndim == 2 and values.shape != (n_dim1, n_dim2):

320 if values.shape == (n_dim2, n_dim1):

321 values = values.T

322 else:

323 raise ValueError(

324 f"{child_path}: shape {values.shape} matches neither the grid "

325 f"({n_dim1}, {n_dim2}) nor its transpose"

326 )

327

328 _write_leaf(destination, name, values, source_child.data_type)

329

330

331def _sensor_series(sensor_object: typing.Any, sensor_name: str, quantity: str) -> dict[str, typing.Any]:

332 """The measured / reconstructed / time / weight series for one sensor.

333

334 The sensor objects still carry the flat getters (`get_array1`, `get_f64`); only `plasma` moved

335 to the IMAS path API.

336 """

337

338 return {

339 "measured": sensor_object.get_array1([sensor_name, quantity, "measured", "value"]),

340 "reconstructed": sensor_object.get_array1([sensor_name, quantity, "calculated", "value"]),

341 "time_measurement": sensor_object.get_array1([sensor_name, quantity, "measured", "time"]),

342 "weight": sensor_object.get_f64([sensor_name, "fit_settings", "weight"]),

343 }

344

345

346def _write_sensor_constraint(entry: typing.Any, series: dict[str, typing.Any], sensor_name: str, i_time: int) -> None:

347 """Fill one `constraints/...` entry for one time-slice."""

348

349 entry.source = sensor_name

350 entry.measured = float(series["measured"][i_time])

351 entry.reconstructed = float(series["reconstructed"][i_time])

352 entry.time_measurement = float(series["time_measurement"][i_time])

353 entry.weight = float(series["weight"])

354

355

356def map_results_to_database(

357 self: "DatabaseWriterIMAS",

358 gsfit_controller: "Gsfit",

359) -> "IDSToplevel":

360 """

361 Map the GSFit results into an IMAS `equilibrium` IDS.

362

363 :param gsfit_controller: the `Gsfit` controller, holding the solved Rust objects

364 :return: a populated `equilibrium` IDS. Nothing is written to any backend

365 """

366

367 logger = gsfit_controller.logger

368

369 plasma = gsfit_controller.plasma

370 bp_probes = gsfit_controller.bp_probes

371 flux_loops = gsfit_controller.flux_loops

372 rogowski_coils = gsfit_controller.rogowski_coils

373 dialoop = gsfit_controller.dialoop

374 pressure_sensors = gsfit_controller.pressure_sensors

375

376 time: npt.NDArray[np.float64] = np.asarray(gsfit_controller.results["TIME"], dtype=np.float64) # [second]

377 n_time: int = len(time)

378

379

380 equilibrium = imas.IDSFactory().equilibrium()

381

382 # ------------------------------------------------------------------

383 # ids_properties and code provenance

384 # ------------------------------------------------------------------

385 equilibrium.ids_properties.homogeneous_time = imas.ids_defs.IDS_TIME_MODE_HOMOGENEOUS

386 equilibrium.ids_properties.comment = (

387 f"GSFit reconstruction of pulse {gsfit_controller.pulseNo}, run {gsfit_controller.run_name}: "

388 f"{gsfit_controller.run_description}"

389 )

390 equilibrium.ids_properties.provider = version_storage.__user__

391 equilibrium.ids_properties.creation_date = version_storage.__datetime__

392

393 equilibrium.time = time

394 equilibrium.time_slice.resize(n_time)

395

396 # ------------------------------------------------------------------

397 # The equilibrium itself, copied leaf by leaf from the Rust IDS

398 # ------------------------------------------------------------------

399 skipped: set[str] = set()

400

401 for i_time in range(n_time):

402 _copy_branch(equilibrium_paths.time_slice[i_time], equilibrium.time_slice[i_time], plasma, "time_slice", skipped)

403

404 _copy_branch(equilibrium_paths.vacuum_toroidal_field, equilibrium.vacuum_toroidal_field, plasma, "vacuum_toroidal_field", skipped)

405 _copy_branch(equilibrium_paths.code, equilibrium.code, plasma, "code", skipped)

406

407 # GSFit stores some things the data dictionary has no home for. Which ones depends on the

408 # dictionary in use, so they are reported from the run rather than listed in the source

409 if skipped:

410 logger.debug(f"IMAS writer: {len(skipped)} GSFit leaves have no node in the data dictionary: {sorted(skipped)}")

411

412 # ------------------------------------------------------------------

413 # Constraints, which live in the sensor objects rather than in the equilibrium IDS

414 # ------------------------------------------------------------------

415 bp_probe_names: list[str] = list(bp_probes.keys())

416 flux_loop_names: list[str] = list(flux_loops.keys())

417 rogowski_coil_names: list[str] = list(rogowski_coils.keys())

418 dialoop_names: list[str] = list(dialoop.keys())

419 pressure_sensor_names: list[str] = list(pressure_sensors.keys())

420

421 bp_probe_series = {name: _sensor_series(bp_probes, name, "b") for name in bp_probe_names}

422 flux_loop_series = {name: _sensor_series(flux_loops, name, "psi") for name in flux_loop_names}

423 rogowski_coil_series = {name: _sensor_series(rogowski_coils, name, "i") for name in rogowski_coil_names}

424 dialoop_series = {name: _sensor_series(dialoop, name, "b") for name in dialoop_names}

425

426 pressure_series = {name: _sensor_series(pressure_sensors, name, "pressure") for name in pressure_sensor_names}

427 pressure_psi = {

428 name: pressure_sensors.get_array1([name, "pressure", "calculated", "psi"]) for name in pressure_sensor_names

429 }

430 pressure_r = {name: pressure_sensors.get_f64([name, "geometry", "r"]) for name in pressure_sensor_names}

431 pressure_z = {name: pressure_sensors.get_f64([name, "geometry", "z"]) for name in pressure_sensor_names}

432

433 for i_time in range(n_time):

434 constraints = equilibrium.time_slice[i_time].constraints

435

436 constraints.b_field_pol_probe.resize(len(bp_probe_names))

437 for i_sensor, sensor_name in enumerate(bp_probe_names):

438 _write_sensor_constraint(constraints.b_field_pol_probe[i_sensor], bp_probe_series[sensor_name], sensor_name, i_time)

439

440 constraints.flux_loop.resize(len(flux_loop_names))

441 for i_sensor, sensor_name in enumerate(flux_loop_names):

442 _write_sensor_constraint(constraints.flux_loop[i_sensor], flux_loop_series[sensor_name], sensor_name, i_time)

443

444 # See conventions note 3: the whole Rogowski family goes to `pf_passive_current`

445 constraints.pf_passive_current.resize(len(rogowski_coil_names))

446 for i_sensor, sensor_name in enumerate(rogowski_coil_names):

447 _write_sensor_constraint(

448 constraints.pf_passive_current[i_sensor], rogowski_coil_series[sensor_name], sensor_name, i_time

449 )

450

451 for sensor_name in dialoop_names:

452 # The DD has a single diamagnetic flux constraint, not an array

453 _write_sensor_constraint(constraints.diamagnetic_flux, dialoop_series[sensor_name], sensor_name, i_time)

454

455 constraints.pressure.resize(len(pressure_sensor_names))

456 for i_sensor, sensor_name in enumerate(pressure_sensor_names):

457 entry = constraints.pressure[i_sensor]

458 _write_sensor_constraint(entry, pressure_series[sensor_name], sensor_name, i_time)

459 entry.position.r = float(pressure_r[sensor_name])

460 entry.position.z = float(pressure_z[sensor_name])

461 entry.position.psi = float(pressure_psi[sensor_name][i_time])

462

463 return equilibrium

Coverage for python/gsfit/database_writers/imas/map_results_to_database.py: 0%

173 statements