diff --git a/sed/calibrator/energy.py b/sed/calibrator/energy.py index aebfa3a6..a27ba73e 100644 --- a/sed/calibrator/energy.py +++ b/sed/calibrator/energy.py @@ -96,6 +96,7 @@ def __init__( self.calibration: Dict[Any, Any] = {} self.tof_column = self._config["dataframe"]["tof_column"] + self.tof_ns_column = self._config["dataframe"].get("tof_ns_column", None) self.corrected_tof_column = self._config["dataframe"]["corrected_tof_column"] self.energy_column = self._config["dataframe"]["energy_column"] self.x_column = self._config["dataframe"]["x_column"] @@ -880,6 +881,53 @@ def append_energy_axis( return df, metadata + def append_tof_ns_axis( + self, + df: Union[pd.DataFrame, dask.dataframe.DataFrame], + tof_column: str = None, + tof_ns_column: str = None, + **kwds, + ) -> Tuple[Union[pd.DataFrame, dask.dataframe.DataFrame], dict]: + """Converts the time-of-flight time from steps to time in ns. + + # TODO: needs tests + + Args: + df (Union[pd.DataFrame, dask.dataframe.DataFrame]): Dataframe to convert. + tof_column (str, optional): Name of the column containing the + time-of-flight steps. Defaults to config["dataframe"]["tof_column"]. + tof_ns_column (str, optional): Name of the column to store the + time-of-flight in nanoseconds. Defaults to config["dataframe"]["tof_ns_column"]. + + Returns: + dask.dataframe.DataFrame: Dataframe with the new columns. + dict: Metadata dictionary. + """ + binwidth = kwds.pop("binwidth", self.binwidth) + binning = kwds.pop("binning", self.binning) + if tof_column is None: + if self.corrected_tof_column in df.columns: + tof_column = self.corrected_tof_column + else: + tof_column = self.tof_column + + if tof_ns_column is None: + tof_ns_column = self.tof_ns_column + if tof_ns_column is None: + raise AttributeError("tof_ns_column not set!") + + df[tof_ns_column] = tof2ns( + binwidth, + binning, + df[tof_column].astype("float64"), + ) + metadata: Dict[str, Any] = { + "applied": True, + "binwidth": binwidth, + "binning": binning, + } + return df, metadata + def gather_calibration_metadata(self, calibration: dict = None) -> dict: """Collects metadata from the energy calibration @@ -2088,80 +2136,23 @@ def tof2evpoly( return energy -def tof_step_to_ns( - df: Union[pd.DataFrame, dask.dataframe.DataFrame], - tof_ns_column: str = None, - tof_binwidth: float = None, - tof_column: str = None, - tof_binning: int = None, - config: dict = None, -) -> Tuple[Union[pd.DataFrame, dask.dataframe.DataFrame], dict]: - """Converts the time-of-flight time from steps to time in ns. - - Args: - df (Union[pd.DataFrame, dask.dataframe.DataFrame]): Dataframe to convert. - tof_ns_column (str, optional): Name of the column to store the - time-of-flight in nanoseconds. Defaults to config["dataframe"]["tof_ns_column"]. - tof_binwidth (float, optional): Time step size in nanoseconds. - Defaults to config["dataframe"]["tof_binwidth"]. - tof_column (str, optional): Name of the column containing the - time-of-flight steps. Defaults to config["dataframe"]["tof_column"]. - tof_binning (int, optional): Binning of the time-of-flight steps. - - Returns: - dask.dataframe.DataFrame: Dataframe with the new columns. - dict: Metadata dictionary. - """ - if tof_binwidth is None: - if config is None: - raise ValueError("Either tof_binwidth or config must be given.") - tof_binwidth = config["dataframe"]["tof_binwidth"] - if tof_column is None: - if config is None: - raise ValueError("Either tof_column or config must be given.") - tof_column = config["dataframe"]["tof_column"] - if tof_binning is None: - if config is None: - raise ValueError("Either tof_binning or config must be given.") - tof_binning = config["dataframe"]["tof_binning"] - if tof_ns_column is None: - if config is None: - raise ValueError("Either tof_ns_column or config must be given.") - tof_ns_column = config["dataframe"].get("tof_ns_column", None) - if tof_ns_column is None: - raise ValueError("No value for tof_ns_column present in config.") - - def func(x): - return tof2ns(x[tof_column], tof_binwidth, tof_binning) - - df[tof_ns_column] = df.map_partitions(func, meta=(tof_column, np.float64)) - metadata: Dict[str, Any] = { - "applied": True, - "tof_binwidth": tof_binwidth, - "tof_binning": tof_binning, - } - return df, metadata - - def tof2ns( - x: Union[List[float], np.ndarray], - tof_binwidth: float, - tof_binning: int, + binwidth: float, + binning: int, + t: float, ) -> Union[List[float], np.ndarray]: """Converts the time-of-flight steps to time-of-flight in nanoseconds. designed for use with dask.dataframe.DataFrame.map_partitions. Args: - x (Union[List[float], np.ndarray]): Time-of-flight steps. - tof_column (str): Name of the column containing the time-of-flight steps. - tof_binwidth (float): Time step size in nanoseconds. - tof_binning (int): Binning of the time-of-flight steps. - + binwidth (float): Time step size in nanoseconds. + binning (int): Binning of the time-of-flight steps. + t (float): TOF value in bin number. Returns: - Union[pd.DataFrame, dask.dataframe.DataFrame]: Dataframe with the new column. + float: Converted time in nanoseconds. """ - val = x * tof_binwidth * 2**tof_binning + val = t * binwidth * 2**binning return val @@ -2176,6 +2167,7 @@ def apply_energy_offset( """Apply an energy shift to the given column(s). # TODO: This funcion can still be improved and needs testsing + # TODO: move inside the ec class Args: df (Union[pd.DataFrame, dask.dataframe.DataFrame]): Dataframe to use. @@ -2244,6 +2236,8 @@ def align_dld_sectors( ) -> Tuple[Union[pd.DataFrame, dask.dataframe.DataFrame], dict]: """Aligns the time-of-flight axis of the different sections of a detector. + # TODO: move inside the ec class + Args: df (Union[pd.DataFrame, dask.dataframe.DataFrame]): Dataframe to use. sector_delays (Sequece[float], optional): Sector delays for the 8s time. diff --git a/sed/core/processor.py b/sed/core/processor.py index 891c2dc1..df696d23 100644 --- a/sed/core/processor.py +++ b/sed/core/processor.py @@ -1172,9 +1172,8 @@ def apply_energy_offset( duplicate_policy="append", ) - def tof_step_to_ns( + def append_tof_ns_axis( self, - tof_ns_column: str = None, **kwargs, ): """Convert time-of-flight channel steps to nanoseconds. @@ -1190,15 +1189,13 @@ def tof_step_to_ns( print("Adding time-of-flight column in nanoseconds to dataframe:") # TODO assert order of execution through metadata - self._dataframe, metadata = energy.tof_step_to_ns( + self._dataframe, metadata = self.ec.append_tof_ns_axis( df=self._dataframe, - tof_ns_column=tof_ns_column, - config=self._config, **kwargs, ) self._attributes.add( metadata, - "step_to_ns", + "tof_ns_conversion", duplicate_policy="append", ) diff --git a/tutorial/5 - hextof workflow.ipynb b/tutorial/5 - hextof workflow.ipynb index a9980191..17333475 100644 --- a/tutorial/5 - hextof workflow.ipynb +++ b/tutorial/5 - hextof workflow.ipynb @@ -60,7 +60,17 @@ "metadata": {}, "outputs": [], "source": [ - "sp.append_energy_axis()\n" + "sp.append_energy_axis()\n", + "sp.append_tof_ns_axis()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sp.dataframe[['dldTime','dldTimeSteps','energy']].head()" ] }, {