From 45f714116642c591ef9288fde1d58ac8c4ad9bac Mon Sep 17 00:00:00 2001
From: Laura Sandoval <Laura.Sandoval@lasp.colorado.edu>
Date: Tue, 27 Aug 2024 13:16:36 -0600
Subject: [PATCH] added back position

---
 .../ultra/unit/test_ultra_l1b_extended.py     | 14 ++++++
 .../ultra/l1b/ultra_l1b_extended.py           | 44 ++++++++++++++++++-
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/imap_processing/tests/ultra/unit/test_ultra_l1b_extended.py b/imap_processing/tests/ultra/unit/test_ultra_l1b_extended.py
index 84b027666..fdae38096 100644
--- a/imap_processing/tests/ultra/unit/test_ultra_l1b_extended.py
+++ b/imap_processing/tests/ultra/unit/test_ultra_l1b_extended.py
@@ -10,6 +10,7 @@
     get_front_y_position,
     get_path_length,
     get_ph_tof_and_back_positions,
+    get_ssd_back_position,
 )
 
 
@@ -85,3 +86,16 @@ def test_get_ph_tof_and_back_positions(
 
     np.testing.assert_array_equal(ph_xb, selected_rows["Xb"].astype("float"))
     np.testing.assert_array_equal(ph_yb, selected_rows["Yb"].astype("float"))
+
+
+def test_get_ssd_back_position(
+    de_dataset,
+    events_fsw_comparison_theta_0,
+):
+    """Tests get_ssd_back_position function."""
+    ssd_yb = get_ssd_back_position(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"))
diff --git a/imap_processing/ultra/l1b/ultra_l1b_extended.py b/imap_processing/ultra/l1b/ultra_l1b_extended.py
index 47c25b653..319077b5e 100644
--- a/imap_processing/ultra/l1b/ultra_l1b_extended.py
+++ b/imap_processing/ultra/l1b/ultra_l1b_extended.py
@@ -1,10 +1,12 @@
 """Calculates Extended Raw Events for ULTRA L1b."""
 
 from enum import Enum
+from typing import ClassVar
 
 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,
@@ -24,10 +26,11 @@
 
 
 class StopType(Enum):
-    """Stop Type: 1=Top, 2=Bottom."""
+    """Stop Type: 1=Top, 2=Bottom, SSD: 8-15."""
 
     Top = 1
     Bottom = 2
+    SSD: ClassVar[list[int]] = [8, 9, 10, 11, 12, 13, 14, 15]
 
 
 def get_front_x_position(start_type: ndarray, start_position_tdc: ndarray) -> ndarray:
@@ -258,3 +261,42 @@ def get_path_length(front_position: tuple, back_position: tuple, d: float) -> fl
     )
 
     return r
+
+
+def get_ssd_back_position(
+    de_dataset: xarray.Dataset,
+) -> NDArray[np.float64]:
+    """
+    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.
+
+    Parameters
+    ----------
+    de_dataset : xarray.Dataset
+        The input dataset containing STOP_TYPE and SSD_FLAG data.
+
+    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. The length of this array matches the number
+        of epochs where the STOP_TYPE is SSD.
+
+    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)
+
+    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)