Skip to content

Commit

Permalink
Update mappings and trasforms for MAST-U
Browse files Browse the repository at this point in the history
  • Loading branch information
@samueljackson92
samueljackson92 committed Jan 2, 2025
1 parent 5e5d558 commit c2ed67b
Show file tree
Hide file tree
Showing 9 changed files with 152 additions and 20 deletions.
10 changes: 6 additions & 4 deletions mappings/mast/groups.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,9 @@
"ahx": {"name": "hard_x_rays", "imas": "hard_x_rays"},
"ait": {"name": "camera_ir", "imas": "camera_ir"},
"alp": {"name": "langmuir_probes", "imas": "langmuir_probes"},
"ama": {"name": "magnetics_a", "imas": "magnetics"},
"amb": {"name": "magnetics_b", "imas": "magnetics"},
"amc": {"name": "magnetics", "imas": "magnetics"},
"amc": {"name": "magnetics_c", "imas": "magnetics"},
"anb": {"name": "nbi", "imas": "nbi"},
"ane": {"name": "interferometer", "imas": "interferometer"},
"anu": {"name": "neutron_diagnostic", "imas": "neutron_diagnostic"},
Expand All @@ -28,9 +29,10 @@
"xbt": {"name": "bes"},
"xdc": {"name": "pulse_schedule", "imas": "pulse_schedule"},
"xim": {"name": "spectrometer_visible", "imas": "spectrometer_visible"},
"xma": {"name": "magnetics_a", "imas": "magnetics"},
"xmc": {"name": "magnetics_c", "imas": "magnetics"},
"xmo": {"name": "magnetics_mirnov","imas": "magnetics"},
"xma": {"name": "magnetics_raw_a", "imas": "magnetics"},
"xmb": {"name": "magnetics_raw_b", "imas": "magnetics"},
"xmc": {"name": "magnetics_raw_c", "imas": "magnetics"},
"xmo": {"name": "magnetics_omaha", "imas": "magnetics"},
"xms": {"name": "mse", "imas": "mse"},
"xsx": {"name": "soft_x_rays", "imas": "soft_x_rays"}
}
16 changes: 16 additions & 0 deletions mappings/mastu/dimensions.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1079,5 +1079,21 @@
},
"epq/input_constraints_mse_a1": {
"dim_0": "mse_n"
},

"xsx/devices_limit": {
"dim_0": "dim_1"
},
"xim/devices_limit": {
"dim_0": "dim_1"
},
"xma/devices_limit": {
"dim_0": "dim_1"
},
"xmb/devices_limit": {
"dim_0": "dim_1"
},
"xmc/devices_limit": {
"dim_0": "dim_1"
}
}
13 changes: 7 additions & 6 deletions mappings/mastu/groups.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,11 @@
"aiv": {"name": "camera_ir_v", "imas": "camera_ir"},
"alp": {"name": "langmuir_probes", "imas": "langmuir_probes"},
"amb": {"name": "magnetics_b", "imas": "magnetics"},
"amc": {"name": "magnetics", "imas": "magnetics"},
"amc": {"name": "magnetics_c", "imas": "magnetics"},
"ams": {"name": "mse", "imas": "mse"},
"anb": {"name": "nbi", "imas": "nbi"},
"ane": {"name": "interferometer", "imas": "interferometer"},
"anu": {"name": "neutron_diagnostic", "imas": "neutron_diagnostic"},
"asm": {"name": "magnetics_saddle", "imas": "magnetics"},
"ayc": {"name": "thomson_scattering", "imas": "thomson_scattering"},
"ayd": {"name": "thomson_scattering_divertor", "imas": "thomson_scattering"},
"epm": {"name": "equilibrium", "imas": "equilibrium"},
Expand All @@ -24,9 +23,11 @@
"rgb": {"name": "camera_visible_bremsstrahlung_a", "imas": "camera_visible"},
"rgc": {"name": "camera_visible_bremsstrahlung_b", "imas": "camera_visible"},
"xbt": {"name": "bes"},
"xdc": {"name": "controllers", "imas": "controllers"},
"xdc": {"name": "pulse_schedule", "imas": "pulse_schedule"},
"xim": {"name": "spectrometer_visible", "imas": "spectrometer_visible"},
"xma": {"name": "magnetics_a", "imas": "magnetics"},
"xmc": {"name": "magnetics_c", "imas": "magnetics"},
"xsx": {"name": "soft_x_rays", "imas": "soft_x_rays"}
"xma": {"name": "magnetics_raw_a", "imas": "magnetics"},
"xmb": {"name": "magnetics_raw_b", "imas": "magnetics"},
"xmc": {"name": "magnetics_omaha", "imas": "magnetics"},
"xsx": {"name": "soft_x_rays", "imas": "soft_x_rays"},
"asx": {"name": "soft_x_rays", "imas": "soft_x_rays"}
}
18 changes: 18 additions & 0 deletions mappings/mastu/variables.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,24 @@
"r_error": "major_radius_error",
"arb": "radial_index"
},
"rba": {
"rba": "data"
},
"rbb": {
"rbb": "data"
},
"rbc": {
"rbc": "data"
},
"rco": {
"rco": "data"
},
"rgb": {
"rgb": "data"
},
"rgc": {
"rgc": "data"
},
"xdc": {
"XDC/PLASMA/T/IP_REF": "reference_plasma_current",
"XDC/FUELLING/T/DENSITY_REF_TARGET": "reference_plasma_density"
Expand Down
1 change: 1 addition & 0 deletions mappings/units.json
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
{
"STATE": "",
"W / (sr.m2)": "W / (sr*m^2)",
"S": "s",
"m^-^3": "m^-3",
Expand Down
2 changes: 2 additions & 0 deletions src/builder.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,8 @@ def create(self, shot: int):
pipeline = self.pipelines.get(group_name)

dataset: xr.Dataset = pipeline(datasets)
dataset.attrs["description"] = dataset_info.description
dataset.attrs["quality"] = dataset_info.quality
dataset, group_name = self._rename_group(dataset, group_name)

logger.info(f"Writing {group_name} for shot #{shot}")
Expand Down
12 changes: 12 additions & 0 deletions src/load.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
from src.registry import Registry
from src.utils import harmonise_name

LAST_MAST_SHOT = 30471


class MissingProfileError(Exception):
pass
Expand Down Expand Up @@ -113,6 +115,15 @@ def get_source_infos(self, shot_num: int) -> t.List[DatasetInfo]:
)
for item in signals
]

# Special case: in MAST-U, soft x rays were moved from XSX -> ASX then back to XSX, but XSX contained raw data signals.
# Here we drop XSX if ASX is avilable, otherwise we return ASX
if shot_num > LAST_MAST_SHOT:
sources = {info.name: info for info in infos}
if "asx" in sources:
sources.pop("xsx")
infos = sources.values()

return infos

def lookup_status_code(self, status):
Expand Down Expand Up @@ -247,6 +258,7 @@ def _convert_signal_to_dataset(self, signal_name, signal):
data = xr.DataArray(data, dims=dim_names, coords=coords, attrs=attrs)
if signal_name == "time":
signal_name = "time_"

data.name = signal_name
data.attrs["name"] = data.name
data.attrs["uda_name"] = uda_name
Expand Down
93 changes: 88 additions & 5 deletions src/pipelines.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -204,6 +204,20 @@ def __init__(self) -> None:
TransformUnits(),
]
),
"asx": Pipeline(
[
MapDict(RenameDimensions(self.dimension_mapping_file)),
MapDict(RenameVariables(self.variable_mapping_file)),
MapDict(DropZeroDimensions()),
MapDict(DropZeroDataset()),
MergeDatasets(),
TransformUnits(),
TensoriseChannels("hcam_l", regex=r"hcam_l_ch(\d+)"),
TensoriseChannels("hcam_u", regex=r"hcam_u_ch(\d+)"),
TensoriseChannels("tcam", regex=r"tcam_ch(\d+)"),
TensoriseChannels("vcam", regex=r"vcam_ch(\d+)"),
]
),
"ayc": Pipeline(
[
MapDict(RenameDimensions(self.dimension_mapping_file)),
Expand Down Expand Up @@ -245,6 +259,28 @@ def __init__(self) -> None:
"input_constraints_fluxloops_fluxloopdim",
]
),
DropCoordinates("output_radialprofiles_totalpressure", ["r"]),
DropCoordinates("output_radialprofiles_toroidalflux", ["r"]),
DropCoordinates("output_radialprofiles_staticpressure", ["r"]),
DropCoordinates("output_radialprofiles_staticpprime", ["r"]),
DropCoordinates("output_radialprofiles_rotationalpressure", ["r"]),
DropCoordinates("output_radialprofiles_radialcoord", ["r"]),
DropCoordinates("output_radialprofiles_r", ["r"]),
DropCoordinates("output_radialprofiles_q", ["r"]),
DropCoordinates("output_radialprofiles_poloidalarea", ["r"]),
DropCoordinates("output_radialprofiles_plasmavolume", ["r"]),
DropCoordinates("output_radialprofiles_plasmadensity", ["r"]),
DropCoordinates(
"output_radialprofiles_normalizedtoroidalflux", ["r"]
),
DropCoordinates(
"output_radialprofiles_normalizedpoloidalflux", ["r"]
),
DropCoordinates("output_radialprofiles_jphi", ["r"]),
DropCoordinates("output_radialprofiles_ffprime", ["r"]),
DropCoordinates("output_radialprofiles_bz", ["r"]),
DropCoordinates("output_radialprofiles_bt", ["r"]),
DropCoordinates("output_radialprofiles_br", ["r"]),
MergeDatasets(),
TransformUnits(),
]
Expand All @@ -270,6 +306,38 @@ def __init__(self) -> None:
"input_constraints_fluxloops_fluxloopdim",
]
),
DropCoordinates("output_radialprofiles_totalpressure", ["r"]),
DropCoordinates("output_radialprofiles_toroidalflux", ["r"]),
DropCoordinates("output_radialprofiles_staticpressure", ["r"]),
DropCoordinates("output_radialprofiles_staticpprime", ["r"]),
DropCoordinates("output_radialprofiles_rotationalpressure", ["r"]),
DropCoordinates("output_radialprofiles_radialcoord", ["r"]),
DropCoordinates("output_radialprofiles_r", ["r"]),
DropCoordinates("output_radialprofiles_q", ["r"]),
DropCoordinates("output_radialprofiles_poloidalarea", ["r"]),
DropCoordinates("output_radialprofiles_plasmavolume", ["r"]),
DropCoordinates("output_radialprofiles_plasmadensity", ["r"]),
DropCoordinates(
"output_radialprofiles_normalizedtoroidalflux", ["r"]
),
DropCoordinates(
"output_radialprofiles_normalizedpoloidalflux", ["r"]
),
DropCoordinates("output_radialprofiles_jphi", ["r"]),
DropCoordinates("output_radialprofiles_ffprime", ["r"]),
DropCoordinates("output_radialprofiles_bz", ["r"]),
DropCoordinates("output_radialprofiles_bt", ["r"]),
DropCoordinates("output_radialprofiles_br", ["r"]),
MergeDatasets(),
TransformUnits(),
]
),
"esm": Pipeline(
[
MapDict(RenameDimensions(self.dimension_mapping_file)),
MapDict(RenameVariables(self.variable_mapping_file)),
MapDict(DropZeroDimensions()),
MapDict(DropZeroDataset()),
MergeDatasets(),
TransformUnits(),
]
Expand All @@ -289,17 +357,18 @@ def __init__(self) -> None:
"rgc": Pipeline(
[MapDict(RenameVariables(self.variable_mapping_file)), ProcessImage()]
),
"esm": Pipeline(
"xbt": Pipeline(
[
MapDict(RenameDimensions(self.dimension_mapping_file)),
MapDict(RenameVariables(self.variable_mapping_file)),
MapDict(DropZeroDimensions()),
MapDict(DropZeroDataset()),
MergeDatasets(),
TransformUnits(),
TensoriseChannels("bes", regex=r"channel(\d+)"),
]
),
"xbt": Pipeline(
"xdc": Pipeline(
[
MapDict(RenameDimensions(self.dimension_mapping_file)),
MapDict(RenameVariables(self.variable_mapping_file)),
Expand All @@ -309,7 +378,7 @@ def __init__(self) -> None:
TransformUnits(),
]
),
"xdc": Pipeline(
"xim": Pipeline(
[
MapDict(RenameDimensions(self.dimension_mapping_file)),
MapDict(RenameVariables(self.variable_mapping_file)),
Expand All @@ -319,17 +388,29 @@ def __init__(self) -> None:
TransformUnits(),
]
),
"xim": Pipeline(
"xma": Pipeline(
[
MapDict(RenameDimensions(self.dimension_mapping_file)),
MapDict(RenameVariables(self.variable_mapping_file)),
MapDict(DropZeroDimensions()),
MapDict(DropZeroDataset()),
MergeDatasets(),
TransformUnits(),
TensoriseChannels("rtdi_01", regex=r"rtdi_01_ch(\d+)$"),
TensoriseChannels("rtdi_02", regex=r"rtdi_02_ch(\d+)$"),
TensoriseChannels("rtdi_03", regex=r"rtdi_03_ch(\d+)$"),
TensoriseChannels("rtdi_04", regex=r"rtdi_04_ch(\d+)$"),
TensoriseChannels("rtdi_05", regex=r"rtdi_05_ch(\d+)$"),
TensoriseChannels("rtdi_06", regex=r"rtdi_06_ch(\d+)$"),
TensoriseChannels("rtdi_07", regex=r"rtdi_07_ch(\d+)$"),
TensoriseChannels("rtdi_08", regex=r"rtdi_08_ch(\d+)$"),
TensoriseChannels("sanx20_01", regex=r"sanx20_01_ch(\d+)$"),
TensoriseChannels("sanx20_02", regex=r"sanx20_02_ch(\d+)$"),
TensoriseChannels("sanx21_01", regex=r"sanx21_01_ch(\d+)$"),
TensoriseChannels("sanx21_02", regex=r"sanx21_02_ch(\d+)$"),
]
),
"xma": Pipeline(
"xmb": Pipeline(
[
MapDict(RenameDimensions(self.dimension_mapping_file)),
MapDict(RenameVariables(self.variable_mapping_file)),
Expand All @@ -347,6 +428,7 @@ def __init__(self) -> None:
MapDict(DropZeroDataset()),
MergeDatasets(),
TransformUnits(),
TensoriseChannels("acq216_202", regex=r"acq216_202_ch(\d+)"),
]
),
"xsx": Pipeline(
Expand All @@ -360,6 +442,7 @@ def __init__(self) -> None:
TensoriseChannels("hcam_l", regex=r"hcam_l_ch(\d+)"),
TensoriseChannels("hcam_u", regex=r"hcam_u_ch(\d+)"),
TensoriseChannels("tcam", regex=r"tcam_ch(\d+)"),
TensoriseChannels("vcam", regex=r"vcam_ch(\d+)"),
]
),
}
Expand Down
7 changes: 2 additions & 5 deletions src/transforms.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -156,9 +156,9 @@ def __call__(self, dataset: xr.Dataset) -> xr.Dataset:
class MergeDatasets(BaseTransform):
def __call__(self, dataset_dict: dict[str, xr.Dataset]) -> xr.Dataset:
for name, item in dataset_dict.items():
if "dim_0" in item.dims:
if "r" in item.dims:
print(f"{item.dims}, {name}")
# print(item)
print(item["r"].values)

dataset = xr.merge(dataset_dict.values())
dataset = dataset.compute()
Expand Down Expand Up @@ -208,9 +208,6 @@ def __init__(
def __call__(self, dataset: xr.Dataset) -> xr.Dataset:
group_keys = self._get_group_keys(dataset)

for key in group_keys:
self._parse_digits(key)

# If we couldn't find any matching keys, do nothing.
if len(group_keys) == 0:
return dataset
Expand Down

0 comments on commit c2ed67b

Please sign in to comment.