added tof offset

imap_processing/tests/ultra/unit/test_ultra_l1b_extended.py

@@ -5,12 +5,13 @@
 import pytest
 
 from imap_processing.ultra.l1b.ultra_l1b_extended import (
+    StartType,
     StopType,
     get_front_x_position,
     get_front_y_position,
     get_path_length,
     get_ph_tof_and_back_positions,
-    get_ssd_back_position,
+    get_ssd_back_position_and_tof_offset,
 )
 
 
@@ -88,14 +89,33 @@ def test_get_ph_tof_and_back_positions(
     np.testing.assert_array_equal(ph_yb, selected_rows["Yb"].astype("float"))
 
 
-def test_get_ssd_back_position(
+def test_get_ssd_back_position_and_tof_offset(
     de_dataset,
     events_fsw_comparison_theta_0,
 ):
     """Tests get_ssd_back_position function."""
-    ssd_yb = get_ssd_back_position(de_dataset)
+    yb, tof_offset, ssd_number = get_ssd_back_position_and_tof_offset(de_dataset)
 
     df = pd.read_csv(events_fsw_comparison_theta_0)
     df_filt = df[(df["StartType"] != -1) & (df["StopType"] >= 8)]
 
-    np.testing.assert_array_equal(ssd_yb, df_filt["Yb"].astype("float"))
+    np.testing.assert_array_equal(yb, df_filt["Yb"].astype("float"))
+
+    tof_offset_lt = tof_offset[df_filt["StartType"] == StartType.Left.value]
+    tof_offset_rt = tof_offset[df_filt["StartType"] == StartType.Right.value]
+
+    ssd_number_lt = ssd_number[df_filt["StartType"] == StartType.Left.value]
+    ssd_number_rt = ssd_number[df_filt["StartType"] == StartType.Right.value]
+
+    np.testing.assert_array_equal(
+        tof_offset_lt[ssd_number_lt == 3],
+        np.full(len(tof_offset_lt[ssd_number_lt == 3]), -4.2),
+    )
+    np.testing.assert_array_equal(
+        tof_offset_rt[ssd_number_rt == 7],
+        np.full(len(tof_offset_rt[ssd_number_rt == 7]), -6),
+    )
+    np.testing.assert_array_equal(
+        tof_offset_rt[ssd_number_rt == 4],
+        np.full(len(tof_offset_rt[ssd_number_rt == 4]), -4),
+    )

imap_processing/ultra/l1b/ultra_l1b_extended.py

@@ -6,7 +6,6 @@
 import numpy as np
 import xarray
 from numpy import ndarray
-from numpy.typing import NDArray
 
 from imap_processing.ultra.l1b.lookup_utils import (
     get_back_position,
@@ -25,6 +24,13 @@
 # TODO: make lookup tables into config files.
 
 
+class StartType(Enum):
+    """Start Type: 1=Left, 2=Right."""
+
+    Left = 1
+    Right = 2
+
+
 class StopType(Enum):
     """Stop Type: 1=Top, 2=Bottom, SSD: 8-15."""
 
@@ -263,14 +269,11 @@ def get_path_length(front_position: tuple, back_position: tuple, d: float) -> fl
     return r
 
 
-def get_ssd_back_position(
+def get_ssd_back_position_and_tof_offset(
     de_dataset: xarray.Dataset,
-) -> NDArray[np.float64]:
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     """
-    Calculate the Y SSD positions (yb) for the specified dataset.
-
-    This function calculates the back positions in the Y direction (yb)
-    for SSDs, expressed in hundredths of a millimeter.
+    Lookup the Y SSD positions (yb), TOF Offset, and SSD number.
 
     Parameters
     ----------
@@ -280,22 +283,37 @@ def get_ssd_back_position(
     Returns
     -------
     yb : np.ndarray
-        A NumPy array containing the calculated Y SSD positions
-        in hundredths of a millimeter for each relevant epoch
-        in the dataset.
+        Y SSD positions in hundredths of a millimeter.
+    tof_offset : np.ndarray
+        TOF offset.
+    ssd_number : np.ndarray
+        SSD number.
 
     Notes
     -----
     The X back position (xb) is assumed to be 0 for SSD.
     """
     indices = np.nonzero(np.isin(de_dataset["STOP_TYPE"], StopType.SSD.value))[0]
-    yb = np.zeros(len(indices), dtype=np.float64)
     de_filtered = de_dataset.isel(epoch=indices)
 
+    yb = np.zeros(len(indices), dtype=np.float64)
+    ssd_number = np.zeros(len(indices), dtype=int)
+    tof_offset = np.zeros(len(indices), dtype=np.float64)
+
     for i in range(8):
         # Multiply ybs times 100 to convert to hundredths of a millimeter.
         yb[de_filtered[f"SSD_FLAG_{i}"].data == 1] = (
             get_image_params(f"YBKSSD{i}") * 100
         )
-
-    return np.asarray(yb, dtype=np.float64)
+        ssd_number[de_filtered[f"SSD_FLAG_{i}"].data == 1] = i
+
+        tof_offset[
+            (de_filtered["START_TYPE"] == StartType.Left.value)
+            & (de_filtered[f"SSD_FLAG_{i}"].data == 1)
+        ] = get_image_params(f"TOFSSDLTOFF{i}")
+        tof_offset[
+            (de_filtered["START_TYPE"] == StartType.Right.value)
+            & (de_filtered[f"SSD_FLAG_{i}"].data == 1)
+        ] = get_image_params(f"TOFSSDRTOFF{i}")
+
+    return yb, tof_offset, ssd_number