l1b extended

imap_processing/tests/ultra/unit/test_ultra_l1b_extended.py

@@ -7,6 +7,8 @@
 from imap_processing.ultra.l1b.ultra_l1b_extended import (
     StartType,
     StopType,
+    CoinType,
+    get_coincidence_positions,
     get_front_x_position,
     get_front_y_position,
     get_path_length,
@@ -30,12 +32,11 @@ def yf_fixture(de_dataset, events_fsw_comparison_theta_0):
 
 def test_get_front_x_position(
     de_dataset,
-    events_fsw_comparison_theta_0,
+    yf_fixture,
 ):
     """Tests get_front_x_position function."""
 
-    df = pd.read_csv(events_fsw_comparison_theta_0)
-    df_filt = df[df["StartType"] != -1]
+    df_filt, _, _ = yf_fixture
 
     xf = get_front_x_position(
         de_dataset["START_TYPE"].data,
@@ -50,6 +51,7 @@ def test_get_front_y_position(yf_fixture):
     df_filt, d, yf = yf_fixture
 
     assert yf == pytest.approx(df_filt["Yf"].astype("float"), abs=1e-5)
+    assert d == pytest.approx(df_filt["d"].astype("float"), abs=1e-5)
 
 
 def test_get_path_length(de_dataset, yf_fixture):
@@ -68,14 +70,13 @@ def test_get_path_length(de_dataset, yf_fixture):
 
 def test_get_ph_tof_and_back_positions(
     de_dataset,
-    events_fsw_comparison_theta_0,
+    yf_fixture,
 ):
     """Tests get_ph_tof_and_back_positions function."""
 
-    df = pd.read_csv(events_fsw_comparison_theta_0)
-    df_filt = df[df["StartType"] != -1]
+    df_filt, _, _ = yf_fixture
 
-    _, _, ph_xb, ph_yb = get_ph_tof_and_back_positions(
+    ph_tof, _, ph_xb, ph_yb = get_ph_tof_and_back_positions(
         de_dataset, df_filt.Xf.astype("float").values, "ultra45"
     )
 
@@ -87,6 +88,7 @@ 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"))
+    np.testing.assert_allclose(ph_tof, selected_rows["TOF"].astype("float"), atol=1e-5, rtol=0)
 
 
 def test_get_ssd_back_position_and_tof_offset(
@@ -127,3 +129,41 @@ def test_get_ssd_back_position_and_tof_offset(
     assert np.all(ssd_number_rt >= 0), "Values in ssd_number_rt out of range."
 
     assert np.all(ssd_number_rt <= 7), "Values in ssd_number_rt out of range."
+
+
+def test_get_coincidence_positions(de_dataset, yf_fixture):
+    """Tests get_coincidence_positions function."""
+    df_filt, _, _ = yf_fixture
+    tof = df_filt["TOF"].astype("float").values
+
+    etof, xc = get_coincidence_positions(de_dataset, tof, "ultra45")
+
+    coin_indices = np.nonzero(
+        np.isin(de_dataset["COIN_TYPE"], [CoinType.Top.value, CoinType.Bottom.value])
+    )[0]
+
+    coin_rows = df_filt.iloc[coin_indices]
+
+    np.testing.assert_allclose(xc[coin_indices],
+                               coin_rows["Xc"].astype("float"), atol=1e-4, rtol=0)
+
+    # look at etof
+    coin_top_indices = np.nonzero(
+        np.isin(de_dataset["COIN_TYPE"], [CoinType.Top.value])
+    )[0]
+    coin_bottom_indices = np.nonzero(
+        np.isin(de_dataset["COIN_TYPE"], [CoinType.Bottom.value])
+    )[0]
+
+    top_rows = df_filt.iloc[coin_top_indices]
+    bottom_rows = df_filt.iloc[coin_bottom_indices]
+
+    np.testing.assert_allclose(etof[coin_bottom_indices],
+                               bottom_rows["eTOF"].astype("float"))
+
+    # assert top
+    np.testing.assert_allclose(etof[coin_top_indices],
+                               top_rows["eTOF"].astype("float"))
+
+
+    print('hi')