move tof_to_ns inside ec class

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.

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",
             )
 

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()"
    ]
   },
   {