diff --git a/jobs/freia_write_datasets.qsub b/jobs/freia_write_datasets.qsub index 076437e..3a11678 100755 --- a/jobs/freia_write_datasets.qsub +++ b/jobs/freia_write_datasets.qsub @@ -27,11 +27,10 @@ num_workers=$3 export PATH="/home/rt2549/dev/:$PATH" random_string=$(head /dev/urandom | tr -dc A-Za-z0-9 | head -c 16) + temp_dir="/common/tmp/sjackson/local_cache/$random_string" +metadata_dir="/common/tmp/sjackson/data/uda/" # Run script -# --force --source_names abm ada adg aga ahx aim air ait alp ama amb amc amh amm ams anb ane ant anu aoe arp asb asm asx ayc aye efm esm esx rba rbb rbc rca rco rgb rgc rir rit xmo xpc xsx -# --force --source_names amc ayc efm xmo xsx time mpirun -np $num_workers \ - python3 -m src.archive.main $temp_dir $summary_file $bucket_path --force \ - --source_names ${@:4} \ No newline at end of file + python3 -m src.main $temp_dir $summary_file --metadata_dir $metadata_dir --bucket_path $bucket_path --file_format zarr --upload --force --source_names ${@:4} diff --git a/src/task.py b/src/task.py index 2cafde5..9f0fb60 100644 --- a/src/task.py +++ b/src/task.py @@ -1,10 +1,11 @@ -from pathlib import Path import os -import xarray as xr -import pandas as pd +import traceback import shutil import subprocess import logging +from pathlib import Path +import xarray as xr +import pandas as pd from src.transforms import MASTPipelineRegistry, MASTUPipelineRegistry from src.mast import MASTClient @@ -83,30 +84,68 @@ def __init__( self.pipelines = MASTUPipelineRegistry() def __call__(self): - signal_infos = self.read_signal_info() - source_infos = self.read_source_info() - - if len(self.signal_names) > 0: - signal_infos = signal_infos.loc[signal_infos.name.isin(self.signal_names)] + try: + self._main() + except Exception as e: + trace = traceback.format_exc() + logging.error(f"Error reading sources for shot {self.shot}: {e}\n{trace}") + + def _main(self): + signal_infos, source_infos = self._read_metadata() - if len(self.source_names) > 0: - signal_infos = signal_infos.loc[signal_infos.source.isin(self.source_names)] + if signal_infos is None or signal_infos is None: + return + signal_infos = self._filter_signals(signal_infos) self.writer.write_metadata() - for key, group_index in signal_infos.groupby("source").groups.items(): - signal_infos_for_source = signal_infos.loc[group_index] - signal_datasets = self.load_source(signal_infos_for_source) - pipeline = self.pipelines.get(key) - dataset = pipeline(signal_datasets) - source_info = source_infos.loc[source_infos["name"] == key].iloc[0] - source_info = source_info.to_dict() - dataset.attrs.update(source_info) - self.writer.write_dataset(dataset) + for source_name, source_group_index in signal_infos.groupby("source").groups.items(): + source_info = self._get_source_metadata(source_name, source_infos) + signal_infos_for_source = self._get_signals_for_source(source_name, source_group_index, signal_infos) + self._process_source(source_name, signal_infos_for_source, source_info) self.writer.consolidate_dataset() + def _process_source(self, source_name: str, signal_infos: pd.DataFrame, source_info: dict): + signal_datasets = self.load_source(signal_infos) + pipeline = self.pipelines.get(source_name) + dataset = pipeline(signal_datasets) + dataset.attrs.update(source_info) + self.writer.write_dataset(dataset) + + def _get_source_metadata(self, source_name, source_infos: pd.DataFrame) -> dict: + source_info = source_infos.loc[source_infos["name"] == source_name].iloc[0] + source_info = source_info.to_dict() + return source_info + + def _get_signals_for_source(self, source_name: str, source_group_index: pd.Series, signal_infos: pd.DataFrame): + signal_infos_for_source = signal_infos.loc[source_group_index] + if source_name == 'xdc': + signal_infos_for_source = signal_infos_for_source.loc[signal_infos_for_source.name == 'xdc/ip_t_ipref'] + return signal_infos_for_source + + def _read_metadata(self) -> tuple[pd.DataFrame, pd.DataFrame]: + try: + signal_infos = self.read_signal_info() + source_infos = self.read_source_info() + except FileNotFoundError: + message = f"Could not find source/signal metadata file for shot {self.shot}" + logging.warning(message) + return None, None + + return signal_infos, source_infos + + + def _filter_signals(self, signal_infos: pd.DataFrame) -> pd.DataFrame: + if len(self.signal_names) > 0: + signal_infos = signal_infos.loc[signal_infos.name.isin(self.signal_names)] + + if len(self.source_names) > 0: + signal_infos = signal_infos.loc[signal_infos.source.isin(self.source_names)] + + return signal_infos + def load_source(self, group: pd.DataFrame) -> dict[str, xr.Dataset]: datasets = {} for _, info in group.iterrows(): @@ -126,7 +165,8 @@ def load_source(self, group: pd.DataFrame) -> dict[str, xr.Dataset]: shot_num=self.shot, name=info["uda_name"] ) except Exception as e: - logging.error(f"Error reading dataset {name} for shot {self.shot}: {e}") + uda_name = info["uda_name"] + logging.warning(f"Could not read dataset {name} ({uda_name}) for shot {self.shot}: {e}") continue dataset.attrs.update(info) diff --git a/src/transforms.py b/src/transforms.py index eeadb6b..c0b59bd 100644 --- a/src/transforms.py +++ b/src/transforms.py @@ -158,10 +158,15 @@ def __init__( self.assign_coords = assign_coords def __call__(self, dataset: xr.Dataset) -> xr.Dataset: - group_keys = self._get_group_keys(dataset) + + # If we couldn't find any matching keys, do nothing. + if len(group_keys) == 0: + return dataset + channels = [dataset[key] for key in group_keys] - dataset[self.stem] = xr.combine_nested(channels, concat_dim=self.dim_name) + combined = xr.combine_nested(channels, concat_dim=self.dim_name) + dataset[self.stem] = combined if self.assign_coords: dataset[self.stem] = dataset[self.stem].assign_coords( @@ -171,7 +176,7 @@ def __call__(self, dataset: xr.Dataset) -> xr.Dataset: dataset[self.stem] = dataset[self.stem].chunk("auto") dataset[self.stem] = self._update_attributes(dataset[self.stem], channels) dataset = dataset.drop_vars(group_keys) - dataset = dataset.compute() + dataset: xr.Dataset = dataset.compute() return dataset def _update_attributes( @@ -300,13 +305,18 @@ def __init__(self, stem: str, path: str): cam_data.drop("name", inplace=True, axis=1) cam_data.drop("comment", inplace=True, axis=1) cam_data.columns = [stem + "_" + c for c in cam_data.columns] - name = stem.split("/")[-1] - cam_data.index.name = name + "_channel" self.stem = stem + index_name = f'{self.stem}_channel' + cam_data[index_name] = [stem + '_' + str(index+1) for index in range(len(cam_data))] + cam_data = cam_data.set_index(index_name) self.cam_data = cam_data.to_xarray() def __call__(self, dataset: xr.Dataset) -> xr.Dataset: - dataset = xr.merge([dataset, self.cam_data], combine_attrs="drop_conflicts") + cam_data = self.cam_data.copy() + # if camera data in not in dataset, then skip and do nothing + if self.stem not in dataset: + return dataset + dataset = xr.merge([dataset, cam_data], combine_attrs="drop_conflicts", join='left') dataset = dataset.compute() return dataset @@ -696,318 +706,318 @@ def __init__(self) -> None: MapDict(XDCRenameDimensions()), MapDict(StandardiseSignalDataset("xdc")), MergeDatasets(), - TensoriseChannels( - "ai_cpu1_ccbv", - dim_name="ai_ccbv_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu1_flcc", - dim_name="ai_flcc_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu1_incon", - dim_name="ai_incon_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu1_lhorw", - dim_name="ai_lhorw_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu1_mid", - dim_name="ai_mid_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu1_obr", - dim_name="ai_obr_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu1_obv", - dim_name="ai_obv_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu1_ring", - dim_name="ai_ring_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu1_rodgr", - dim_name="ai_rodgr_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu1_uhorw", - dim_name="ai_uhorw_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu1_vertw", - dim_name="ai_vertw_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu2_ccbv", - dim_name="ai_ccbv_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu2_flcc", - dim_name="ai_flcc_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu2_incon", - dim_name="ai_incon_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu2_lhorw", - dim_name="ai_lhorw_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu2_mid", - dim_name="ai_mid_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu2_obr", - dim_name="ai_obr_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu2_obv", - dim_name="ai_obv_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu2_ring", - dim_name="ai_ring_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu2_rodgr", - dim_name="ai_rodgr_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu2_uhorw", - dim_name="ai_uhorw_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu2_vertw", - dim_name="ai_vertw_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu3_ccbv", - dim_name="ai_ccbv_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu3_flcc", - dim_name="ai_flcc_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu3_incon", - dim_name="ai_incon_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu3_lhorw", - dim_name="ai_lhorw_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu3_mid", - dim_name="ai_mid_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu3_obr", - dim_name="ai_obr_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu3_obv", - dim_name="ai_obv_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu3_ring", - dim_name="ai_ring_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu3_rodgr", - dim_name="ai_rodgr_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu3_uhorw", - dim_name="ai_uhorw_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu3_vertw", - dim_name="ai_vertw_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu4_ccbv", - dim_name="ai_ccbv_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu4_flcc", - dim_name="ai_flcc_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu4_incon", - dim_name="ai_incon_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu4_lhorw", - dim_name="ai_lhorw_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu4_mid", - dim_name="ai_mid_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu4_obr", - dim_name="ai_obr_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu4_obv", - dim_name="ai_obv_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu4_ring", - dim_name="ai_ring_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu4_rodgr", - dim_name="ai_rodgr_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu4_uhorw", - dim_name="ai_uhorw_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_cpu4_vertw", - dim_name="ai_vertw_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_raw_ccbv", dim_name="ai_ccbv", assign_coords=False - ), - TensoriseChannels( - "ai_raw_flcc", - dim_name="ai_flcc_channel", - assign_coords=False, - ), - TensoriseChannels( - "ai_raw_obv", dim_name="ai_obv_channel", assign_coords=False - ), - TensoriseChannels( - "ai_raw_obr", dim_name="ai_obr_channel", assign_coords=False - ), - TensoriseChannels( - "equil_s_seg", - regex=r"equil_s_seg(\d+)$", - dim_name="equil_seg_channel", - assign_coords=False, - ), - TensoriseChannels( - "equil_s_seg_at", - regex=r"equil_s_seg(\d+)at$", - dim_name="equil_seg_channel", - assign_coords=False, - ), - TensoriseChannels( - "equil_s_seg_rt", - regex=r"equil_s_seg(\d+)rt$", - dim_name="equil_seg_channel", - assign_coords=False, - ), - TensoriseChannels( - "equil_s_seg_zt", - regex=r"equil_s_seg(\d+)zt$", - dim_name="equil_seg_channel", - assign_coords=False, - ), - TensoriseChannels( - "equil_s_segb", - dim_name="equil_seg_channel", - assign_coords=False, - ), - TensoriseChannels( - "equil_t_seg", - regex=r"equil_t_seg(\d+)$", - dim_name="equil_seg_channel", - assign_coords=False, - ), - TensoriseChannels( - "equil_t_seg_u", - regex=r"equil_t_seg(\d+)u$", - dim_name="equil_seg_channel", - assign_coords=False, - ), - TensoriseChannels("isoflux_e_seg"), - TensoriseChannels( - "isoflux_t_rpsh_n", - regex=r"isoflux_t_rpsh(\d+)n", - ), - TensoriseChannels( - "isoflux_t_rpsh_p", - regex=r"isoflux_t_rpsh(\d+)p", - ), - TensoriseChannels("isoflux_t_seg", regex=r"isoflux_t_seg(\d+)$"), - TensoriseChannels( - "isoflux_t_seg_gd", regex=r"isoflux_t_seg(\d+)gd$" - ), - TensoriseChannels( - "isoflux_t_seg_gi", regex=r"isoflux_t_seg(\d+)gi$" - ), - TensoriseChannels( - "isoflux_t_seg_gp", regex=r"isoflux_t_seg(\d+)gp$" - ), - TensoriseChannels( - "isoflux_t_seg_td", regex=r"isoflux_t_seg(\d+)td$" - ), - TensoriseChannels( - "isoflux_t_seg_ti", regex=r"isoflux_t_seg(\d+)ti$" - ), - TensoriseChannels( - "isoflux_t_seg_tp", regex=r"isoflux_t_seg(\d+)tp$" - ), - TensoriseChannels("isoflux_t_seg_u", regex=r"isoflux_t_seg(\d+)u$"), - TensoriseChannels( - "isoflux_t_zpsh_n", - regex=r"isoflux_t_zpsh(\d+)n", - ), - TensoriseChannels( - "isoflux_t_zpsh_p", - regex=r"isoflux_t_zpsh(\d+)p", - ), + # TensoriseChannels( + # "ai_cpu1_ccbv", + # dim_name="ai_ccbv_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu1_flcc", + # dim_name="ai_flcc_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu1_incon", + # dim_name="ai_incon_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu1_lhorw", + # dim_name="ai_lhorw_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu1_mid", + # dim_name="ai_mid_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu1_obr", + # dim_name="ai_obr_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu1_obv", + # dim_name="ai_obv_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu1_ring", + # dim_name="ai_ring_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu1_rodgr", + # dim_name="ai_rodgr_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu1_uhorw", + # dim_name="ai_uhorw_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu1_vertw", + # dim_name="ai_vertw_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu2_ccbv", + # dim_name="ai_ccbv_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu2_flcc", + # dim_name="ai_flcc_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu2_incon", + # dim_name="ai_incon_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu2_lhorw", + # dim_name="ai_lhorw_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu2_mid", + # dim_name="ai_mid_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu2_obr", + # dim_name="ai_obr_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu2_obv", + # dim_name="ai_obv_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu2_ring", + # dim_name="ai_ring_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu2_rodgr", + # dim_name="ai_rodgr_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu2_uhorw", + # dim_name="ai_uhorw_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu2_vertw", + # dim_name="ai_vertw_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu3_ccbv", + # dim_name="ai_ccbv_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu3_flcc", + # dim_name="ai_flcc_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu3_incon", + # dim_name="ai_incon_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu3_lhorw", + # dim_name="ai_lhorw_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu3_mid", + # dim_name="ai_mid_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu3_obr", + # dim_name="ai_obr_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu3_obv", + # dim_name="ai_obv_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu3_ring", + # dim_name="ai_ring_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu3_rodgr", + # dim_name="ai_rodgr_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu3_uhorw", + # dim_name="ai_uhorw_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu3_vertw", + # dim_name="ai_vertw_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu4_ccbv", + # dim_name="ai_ccbv_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu4_flcc", + # dim_name="ai_flcc_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu4_incon", + # dim_name="ai_incon_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu4_lhorw", + # dim_name="ai_lhorw_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu4_mid", + # dim_name="ai_mid_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu4_obr", + # dim_name="ai_obr_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu4_obv", + # dim_name="ai_obv_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu4_ring", + # dim_name="ai_ring_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu4_rodgr", + # dim_name="ai_rodgr_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu4_uhorw", + # dim_name="ai_uhorw_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_cpu4_vertw", + # dim_name="ai_vertw_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_raw_ccbv", dim_name="ai_ccbv", assign_coords=False + # ), + # TensoriseChannels( + # "ai_raw_flcc", + # dim_name="ai_flcc_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "ai_raw_obv", dim_name="ai_obv_channel", assign_coords=False + # ), + # TensoriseChannels( + # "ai_raw_obr", dim_name="ai_obr_channel", assign_coords=False + # ), + # TensoriseChannels( + # "equil_s_seg", + # regex=r"equil_s_seg(\d+)$", + # dim_name="equil_seg_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "equil_s_seg_at", + # regex=r"equil_s_seg(\d+)at$", + # dim_name="equil_seg_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "equil_s_seg_rt", + # regex=r"equil_s_seg(\d+)rt$", + # dim_name="equil_seg_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "equil_s_seg_zt", + # regex=r"equil_s_seg(\d+)zt$", + # dim_name="equil_seg_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "equil_s_segb", + # dim_name="equil_seg_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "equil_t_seg", + # regex=r"equil_t_seg(\d+)$", + # dim_name="equil_seg_channel", + # assign_coords=False, + # ), + # TensoriseChannels( + # "equil_t_seg_u", + # regex=r"equil_t_seg(\d+)u$", + # dim_name="equil_seg_channel", + # assign_coords=False, + # ), + # TensoriseChannels("isoflux_e_seg"), + # TensoriseChannels( + # "isoflux_t_rpsh_n", + # regex=r"isoflux_t_rpsh(\d+)n", + # ), + # TensoriseChannels( + # "isoflux_t_rpsh_p", + # regex=r"isoflux_t_rpsh(\d+)p", + # ), + # TensoriseChannels("isoflux_t_seg", regex=r"isoflux_t_seg(\d+)$"), + # TensoriseChannels( + # "isoflux_t_seg_gd", regex=r"isoflux_t_seg(\d+)gd$" + # ), + # TensoriseChannels( + # "isoflux_t_seg_gi", regex=r"isoflux_t_seg(\d+)gi$" + # ), + # TensoriseChannels( + # "isoflux_t_seg_gp", regex=r"isoflux_t_seg(\d+)gp$" + # ), + # TensoriseChannels( + # "isoflux_t_seg_td", regex=r"isoflux_t_seg(\d+)td$" + # ), + # TensoriseChannels( + # "isoflux_t_seg_ti", regex=r"isoflux_t_seg(\d+)ti$" + # ), + # TensoriseChannels( + # "isoflux_t_seg_tp", regex=r"isoflux_t_seg(\d+)tp$" + # ), + # TensoriseChannels("isoflux_t_seg_u", regex=r"isoflux_t_seg(\d+)u$"), + # TensoriseChannels( + # "isoflux_t_zpsh_n", + # regex=r"isoflux_t_zpsh(\d+)n", + # ), + # TensoriseChannels( + # "isoflux_t_zpsh_p", + # regex=r"isoflux_t_zpsh(\d+)p", + # ), TransformUnits(), ] ), @@ -1036,6 +1046,11 @@ def __init__(self) -> None: TensoriseChannels("hcam_l", regex=r"hcam_l_(\d+)"), TensoriseChannels("hcam_u", regex=r"hcam_u_(\d+)"), TensoriseChannels("tcam", regex=r"tcam_(\d+)"), + TensoriseChannels("hcam_l", regex=r"hcaml#(\d+)"), + TensoriseChannels("hcam_u", regex=r"hcamu#(\d+)"), + TensoriseChannels("hpzr", regex=r"hpzr_(\d+)"), + TensoriseChannels("v_ste29", regex=r"v_ste29_(\d+)"), + TensoriseChannels("v_ste36", regex=r"v_ste36_(\d+)"), TransformUnits(), AddXSXCameraParams("hcam_l", "parameters/xsx_camera_l.csv"), AddXSXCameraParams("hcam_u", "parameters/xsx_camera_u.csv"), diff --git a/src/workflow.py b/src/workflow.py index 96cb7f3..fe108ff 100644 --- a/src/workflow.py +++ b/src/workflow.py @@ -122,9 +122,8 @@ def __call__(self, shot: int): create() except Exception as e: import traceback - traceback.print_exc() - print(traceback.format_exc()) - logging.error(f"Failed to run workflow with error {type(e)}: {e}") + trace = traceback.format_exc() + logging.error(f"Failed to run workflow with error {type(e)}: {e}\n{trace}") diff --git a/tests/test_transforms.py b/tests/test_transforms.py index d95225e..6b1733d 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -1,3 +1,4 @@ +import numpy as np import xarray as xr from src.transforms import ( AddXSXCameraParams, @@ -85,7 +86,14 @@ def test_standardise_dataset(fake_dataset): assert "plasma_current_error" in dataset.data_vars -def test_xsx_camera_params(fake_dataset): +def test_xsx_camera_params(): + fake_dataset = xr.Dataset( + data_vars=dict( + tcam=(("time", 'tcam_channels'), np.random.random((100, 18))), + time=("time", np.random.random(100)), + ), + attrs={"name": "xsx/tcam", "shot_id": 30420}, + ) transform = AddXSXCameraParams("tcam", "parameters/xsx_camera_t.csv") dataset = transform(fake_dataset)