diff --git a/sed/calibrator/delay.py b/sed/calibrator/delay.py index b0506d6f..8aaa78c7 100644 --- a/sed/calibrator/delay.py +++ b/sed/calibrator/delay.py @@ -25,22 +25,92 @@ class DelayCalibrator: def __init__( self, config: dict = None, - ): + ) -> None: """Initialization of the DelayCalibrator class passes the config. Args: config (dict, optional): Config dictionary. Defaults to None. """ - if config is None: - config = {} + self._config: dict = config or {} - self._config = config + self.loader: str = self._config["core"]["loader"] + if self.loader == "hextof": + self._append_delay_axis_docstring: str = self.append_delay_axis_hextof.__doc__ + elif self.loader == "mpes": + self._append_delay_axis_docstring: str = self.append_delay_axis_mpes.__doc__ + else: + raise NotImplementedError(f"Loader '{self.loader}' not implemented.") - self.adc_column = self._config["dataframe"]["adc_column"] - self.delay_column = self._config["dataframe"]["delay_column"] - self.calibration: Dict[Any, Any] = {} + self.adc_column: str = self._config["dataframe"].get("adc_column", None) + self.delay_stage_column: str = self._config["dataframe"].get("delay_stage_column", None) + if self.delay_stage_column is None and self.adc_column is None: + raise ValueError("No delay stage column specified.") + self.delay_column: str = self._config["dataframe"]["delay_column"] + self.calibration: Dict[str, Any] = {} def append_delay_axis( + self, + df: Union[pd.DataFrame, dask.dataframe.DataFrame], + *args, + **kwargs, + ) -> Tuple[Union[pd.DataFrame, dask.dataframe.DataFrame], dict]: + """TODO: docstring""" + method = getattr(self, f"append_delay_axis_{self.loader}") + return method(df, *args, **kwargs) + + def append_delay_axis_hextof( + self, + df: Union[pd.DataFrame, dask.dataframe.DataFrame], + time0: float = None, + flip_time_axis: bool = None, + delay_stage_column: str = None, + delay_column: str = None, + ) -> Tuple[dask.dataframe.DataFrame, dict]: + """Calculate and append the delay axis to the events dataframe. + + Args: + df (Union[pd.DataFrame, dask.dataframe.DataFrame]): The dataframe where + to apply the delay calibration to. + time0 (float, optional): Pump-Probe overlap value of the delay coordinate. + If omitted, it is searched for in the data files. + flip_time_axis (bool, optional): Invert the time axis. + delay_stage_column (str, optional): Source column for delay calibration. + Defaults to config["dataframe"]["delay_stage_column"]. + delay_column (str, optional): Destination column for delay calibration. + Defaults to config["dataframe"]["delay_column"]. + + Returns: + Union[pd.DataFrame, dask.dataframe.DataFrame]: dataframe with added column + and delay calibration metdata dictionary. + """ + assert self.loader == "hextof", "Invalid loader for this method." + # pylint: disable=duplicate-code + delay_stage_column = delay_stage_column or self.delay_stage_column + delay_column = delay_column or self.delay_column + + time0 = time0 or self._config["delay"].get("time0", 0) + flip_time_axis = flip_time_axis or self._config["delay"].get("flip_time_axis", False) + + def calibrate_time(x, time0, flip_time_axis) -> Any: + return (x[delay_stage_column] - time0) * (-1 if flip_time_axis else 1) + + df[delay_column] = df.map_partitions( + calibrate_time, + time0, + flip_time_axis, + meta=(delay_column, np.float64), + ) + + metadata: Dict[str, Any] = { + "applied": True, + "calibration": { + "time0": time0, + "flip_time_axis": flip_time_axis, + }, + } + return df, metadata + + def append_delay_axis_mpes( self, df: Union[pd.DataFrame, dask.dataframe.DataFrame], adc_column: str = None, @@ -93,6 +163,8 @@ def append_delay_axis( Union[pd.DataFrame, dask.dataframe.DataFrame]: dataframe with added column and delay calibration metdata dictionary. """ + assert self.loader == "mpes", "Invalid loader for this method." + # pylint: disable=duplicate-code if calibration is None: if self.calibration: diff --git a/tests/calibrator/test_delay.py b/tests/calibrator/test_delay.py index e702a7cc..570a53d6 100644 --- a/tests/calibrator/test_delay.py +++ b/tests/calibrator/test_delay.py @@ -3,6 +3,9 @@ import os from importlib.util import find_spec +import dask.dataframe +import numpy as np +import pandas as pd import pytest from sed.calibrator.delay import DelayCalibrator @@ -114,3 +117,71 @@ def test_delay_parameters_from_delay_range_mm(): assert "adc_range" in metadata["calibration"] assert "time0" in metadata["calibration"] assert "delay_range_mm" in metadata["calibration"] + + +def test_loader_selection(): + """test that the correct calibration method is used based on the loader in the config""" + config = parse_config( + config={ + "core": {"loader": "mpes"}, + "delay": { + "p1_key": "@trARPES:DelayStage:p1", + "p2_key": "@trARPES:DelayStage:p2", + "t0_key": "@trARPES:DelayStage:t0", + }, + }, + folder_config={}, + user_config={}, + system_config={}, + ) + df, _ = get_loader(loader_name="mpes", config=config).read_dataframe( + files=[file], + collect_metadata=False, + ) + dc = DelayCalibrator(config=config) + assert dc.loader == "mpes" + # assert dc.append_delay_axis.__doc__ == dc.append_delay_axis_mpes.__doc__ + assert getattr(dc, f"append_delay_axis_{dc.loader}") == dc.append_delay_axis_mpes + + config = parse_config( + config={ + "core": {"loader": "hextof"}, + "delay": {"time0": 1}, + }, + folder_config={}, + user_config={}, + system_config={}, + ) + _ = dask.dataframe.from_pandas( + pd.DataFrame({"delayStage": np.linspace(0, 1, 100)}), + npartitions=2, + ) + dc = DelayCalibrator(config=config) + assert dc.loader == "hextof" + # assert dc.append_delay_axis.__doc__ == dc.append_delay_axis_hextof.__doc__ + assert getattr(dc, f"append_delay_axis_{dc.loader}") == dc.append_delay_axis_hextof + + +def test_hextof_append_delay(): + """test functionality of the hextof delay calibration method""" + config = parse_config( + config={ + "core": {"loader": "hextof"}, + "dataframe": {"delay_column": "delay", "delay_stage_column": "delayStage"}, + "delay": {"time0": 1}, + }, + folder_config={}, + user_config={}, + system_config={}, + ) + df = dask.dataframe.from_pandas( + pd.DataFrame({"dldPosX": np.linspace(0, 1, 100), "delayStage": np.linspace(0, 1, 100)}), + npartitions=2, + ) + dc = DelayCalibrator(config=config) + df, metadata = dc.append_delay_axis(df) + assert "delay" in df.columns + assert "time0" in metadata["calibration"] + assert metadata["calibration"]["time0"] == 1 + assert metadata["calibration"]["flip_time_axis"] is False + np.testing.assert_allclose(df["delay"], np.linspace(0, 1, 100) - 1)