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 fdae38096..a0cd7930a 100644 --- a/imap_processing/tests/ultra/unit/test_ultra_l1b_extended.py +++ b/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), + ) diff --git a/imap_processing/ultra/l1b/ultra_l1b_extended.py b/imap_processing/ultra/l1b/ultra_l1b_extended.py index 79fb07985..35fa6e87a 100644 --- a/imap_processing/ultra/l1b/ultra_l1b_extended.py +++ b/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