From 6bbf4cd207e9057e39c56ae25514b7d0853b7c71 Mon Sep 17 00:00:00 2001 From: svoboda Date: Thu, 5 Feb 2026 12:28:01 +0100 Subject: [PATCH] Add support for specifying revision to fiesta loading --- cherab/compass/equilibrium/fiesta.py | 44 ++++++++++++++-------------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/cherab/compass/equilibrium/fiesta.py b/cherab/compass/equilibrium/fiesta.py index 91d7cfa..5d8a46d 100644 --- a/cherab/compass/equilibrium/fiesta.py +++ b/cherab/compass/equilibrium/fiesta.py @@ -2,7 +2,7 @@ import xarray as xr -def fiesta_from_cudb(shot_number): +def fiesta_from_cudb(shot_number, revision=-1): """ Obtains Fiesta equilibrium for COMPASS Upgrade from the cudb database and returns a COMPASSEquilibrium object. :param shot_number: Shot number to load from cudb @@ -18,72 +18,72 @@ def fiesta_from_cudb(shot_number): shot = cdbxr.Shot(shot_number) - data = shot["psi/fiesta_out"].copy() + data = shot[f"psi/fiesta_out:{shot_number}:{revision}"].copy() data.name = "psi_grid" data = data.to_dataset() data = data.rename({"fiesta_r": "R", "fiesta_z": "Z"}) - data = data.merge(shot["psi_0/fiesta_out"].to_dataset()) + data = data.merge(shot[f"psi_0/fiesta_out:{shot_number}:{revision}"].to_dataset()) data = data.rename({"psi_0": "psi_axis"}) - data = data.merge(shot["psi_boundary/fiesta_out"].to_dataset()) + data = data.merge(shot[f"psi_boundary/fiesta_out:{shot_number}:{revision}"].to_dataset()) data = data.rename({"psi_boundary": "psi_lcfs"}) - data = data.merge(shot["r_mag/fiesta_out"].to_dataset()) + data = data.merge(shot[f"r_mag/fiesta_out:{shot_number}:{revision}"].to_dataset()) data = data.rename({"r_mag": "R_magnetic_axis"}) - data = data.merge(shot["z_mag/fiesta_out"].to_dataset()) + data = data.merge(shot[f"z_mag/fiesta_out:{shot_number}:{revision}"].to_dataset()) data = data.rename({"z_mag": "Z_magnetic_axis"}) # add xpoint data - tmp = shot["xp_lower_r/fiesta_out"].expand_dims({"xpoint": 1}) - tmp = xr.concat((tmp, shot["xp_upper_r/fiesta_out"].expand_dims({"xpoint": 1})), dim="xpoint") + tmp = shot[f"xp_lower_r/fiesta_out:{shot_number}:{revision}"].expand_dims({"xpoint": 1}) + tmp = xr.concat((tmp, shot[f"xp_upper_r/fiesta_out:{shot_number}:{revision}"].expand_dims({"xpoint": 1})), dim="xpoint") tmp.name = "R_xpoint" tmp = tmp.to_dataset() data = xr.merge((data, tmp)) - tmp = shot["xp_lower_z/fiesta_out"].expand_dims({"xpoint": 1}) - tmp = xr.concat((tmp, shot["xp_upper_z/fiesta_out"].expand_dims({"xpoint": 1})), dim="xpoint") + tmp = shot[f"xp_lower_z/fiesta_out:{shot_number}:{revision}"].expand_dims({"xpoint": 1}) + tmp = xr.concat((tmp, shot[f"xp_upper_z/fiesta_out:{shot_number}:{revision}"].expand_dims({"xpoint": 1})), dim="xpoint") tmp.name = "Z_xpoint" tmp = tmp.to_dataset() data = xr.merge((data, tmp)) - tmp = shot["sp_hfs_r/fiesta_out"].expand_dims({"strike_point": 1}) - tmp = xr.concat((tmp, shot["sp_hfs_r/fiesta_out"].expand_dims({"strike_point": 1})), dim="strike_point") + tmp = shot[f"sp_hfs_r/fiesta_out:{shot_number}:{revision}"].expand_dims({"strike_point": 1}) + tmp = xr.concat((tmp, shot[f"sp_hfs_r/fiesta_out:{shot_number}:{revision}"].expand_dims({"strike_point": 1})), dim="strike_point") tmp.name = "R_strike_point" tmp = tmp.to_dataset() data = xr.merge((data, tmp)) - tmp = shot["sp_hfs_z/fiesta_out"].expand_dims({"strike_point": 1}) - tmp = xr.concat((tmp, shot["sp_hfs_z/fiesta_out"].expand_dims({"strike_point": 1})), dim="strike_point") + tmp = shot[f"sp_hfs_z/fiesta_out:{shot_number}:{revision}"].expand_dims({"strike_point": 1}) + tmp = xr.concat((tmp, shot[f"sp_hfs_z/fiesta_out:{shot_number}:{revision}"].expand_dims({"strike_point": 1})), dim="strike_point") tmp.name = "Z_strike_point" tmp = tmp.to_dataset() data = xr.merge((data, tmp)) - data = data.merge(shot["f/fiesta_out"].to_dataset()) + data = data.merge(shot[f"f/fiesta_out:{shot_number}:{revision}"].to_dataset()) data = data.rename({"f": "f_profile"}) - data = data.merge(shot["q/fiesta_out"].to_dataset()) + data = data.merge(shot[f"q/fiesta_out:{shot_number}:{revision}"].to_dataset()) data = data.rename({"q": "q_profile"}) - data = data.merge(shot["Bt_vac_mag_axis/fiesta_out"].to_dataset()) + data = data.merge(shot[f"Bt_vac_mag_axis/fiesta_out:{shot_number}:{revision}"].to_dataset()) data = data.rename({"Bt_vac_mag_axis": "Btor_vacuum_magnitude"}) - data = data.merge(shot["Bt_vac_mag_axis/fiesta_out"].to_dataset()) + data = data.merge(shot[f"Bt_vac_mag_axis/fiesta_out:{shot_number}:{revision}"].to_dataset()) data = data.rename({"Bt_vac_mag_axis": "Btor_vacuum_radius"}) - tmp = shot["boundary_closed_r/fiesta_out"].to_dataset().rename({"boundary_closed_r": "R_lcfs", + tmp = shot[f"boundary_closed_r/fiesta_out:{shot_number}:{revision}"].to_dataset().rename({"boundary_closed_r": "R_lcfs", "boundary_closed_r_axis1": "lcfs_vertex"}) - tmp = tmp.merge(shot["boundary_closed_z/fiesta_out"].to_dataset().rename({"boundary_closed_z": "Z_lcfs", + tmp = tmp.merge(shot[f"boundary_closed_z/fiesta_out:{shot_number}:{revision}"].to_dataset().rename({"boundary_closed_z": "Z_lcfs", "boundary_closed_z_axis1": "lcfs_vertex"})) data = data.merge(tmp) - tmp = shot["Z_limiter/fiesta_out"] + tmp = shot[f"Z_limiter/fiesta_out:{shot_number}:{revision}"] tmp = xr.Dataset({"Z_limiter": ("limiter_vertex", tmp.Z_limiter.values)}) data = data.merge(tmp) - tmp = shot["R_limiter/fiesta_out"] + tmp = shot[f"R_limiter/fiesta_out:{shot_number}:{revision}"] tmp = xr.Dataset({"R_limiter": ("limiter_vertex", tmp.R_limiter.values)}) data = data.merge(tmp)