ULTRA l1b High Energy (SEP) culling

imap_processing/tests/external_test_data_config.py

@@ -152,6 +152,8 @@
     ("extendedspin_test_data_repoint00047.csv", "ultra/data/l1/"),
     ("status_test_data_repoint00047.csv", "ultra/data/l1/"),
     ("voltage_culling_results_repoint00047.csv", "ultra/data/l1/"),
+    ("validate_high_energy_culling_results_repoint00047.csv", "ultra/data/l1/"),
+    ("de_test_data_repoint00047.csv", "ultra/data/l1/"),
     ("FM45_UltraFM45Extra_TV_Tests_2024-01-22T0930_20240122T093008.CCSDS", "ultra/data/l0/"),
     ("ultra45_raw_sc_rawnrgevnt_19840122_00.csv", "ultra/data/l0/"),
     ("ultra45_raw_sc_enaphxtofhnrgimg_FM45_UltraFM45Extra_TV_Tests_2024-01-22T0930_20240122T093008.csv",

imap_processing/tests/ultra/unit/test_ultra_l1b_culling.py

@@ -21,21 +21,26 @@
     count_rejected_events_per_spin,
     expand_bin_flags_to_spins,
     flag_attitude,
+    flag_high_energy,
     flag_hk,
     flag_imap_instruments,
     flag_low_voltage,
     flag_rates,
     flag_scattering,
+    get_binned_energy_ranges,
     get_binned_spins_edges,
     get_de_rejection_mask,
     get_energy_and_spin_dependent_rejection_mask,
     get_energy_histogram,
+    get_energy_range_flags,
     get_n_sigma,
     get_pulses_per_spin,
     get_spin_data,
     get_valid_earth_angle_events,
+    get_valid_events_per_energy_range,
 )
 from imap_processing.ultra.l1b.ultra_l1b_extended import get_spin_info
+from imap_processing.ultra.l1c.l1c_lookup_utils import build_energy_bins
 
 TEST_PATH = imap_module_directory / "tests" / "ultra" / "data" / "l1"
 
@@ -514,3 +519,232 @@ def test_get_valid_earth_angle_events(mock_spkezr):
 
     actual_flags = get_valid_earth_angle_events(de_dataset, earth_angle_threshold)
     np.testing.assert_array_equal(actual_flags, expected_flags)
+
+
+def test_get_valid_events_per_energy_range():
+    """Tests get_valid_events_per_energy_range function."""
+    np.random.seed(0)
+    energy_range_edges = np.array([3, 5, 7, 18])  # 3 example energy bins
+    # example energy values that fall into different bins
+    # - Events 1-3 and 8 fall into the second bin (5-7)
+    # - Events 4-7 fall into the third bin (7-18)
+    # - The rest of the event energies don't fall into any bins and should be
+    # invalid for that energy range
+    energy = np.array([5, 5, 6, 9, 10, 11, 12, 6, 7, 1, 20, 63])
+    # Mark event 2 (energy bin 2) and 5 (energy bin 3) as outliers
+    quality_outliers = np.array([0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0])
+    # Mark event 1 (energy bin 2), 9 (energy bin 3), and 11 and 12 (No energy bin)
+    quality_scattering = np.array([1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1])
+    ebin = np.full(len(energy), 10)
+    # mark event 6 as having an invalid ebin
+    ebin[5] = -1
+    de_dps_velocity = np.random.random((len(energy), 3))
+
+    de_dataset = xr.Dataset(
+        {
+            "de_dps_velocity": (("epoch", "component"), de_dps_velocity),
+            "event_times": ("epoch", np.arange(len(energy))),
+            "energy_spacecraft": ("epoch", energy),
+            "quality_outliers": ("epoch", quality_outliers),
+            "quality_scattering": ("epoch", quality_scattering),
+            "ebin": ("epoch", ebin),
+        }
+    )
+    keepout_angle = np.radians(180)
+    valid_events = get_valid_events_per_energy_range(
+        de_dataset, energy_range_edges, keepout_angle, 90
+    )
+
+    # Assert that for the first energy bin (3-5), all are false
+    assert np.array_equal(valid_events[0], np.full(len(valid_events[0]), False))
+    # Assert that for the second energy bin (5-7), all are false except
+    # events 3 and 8 (event 1 had an outlier flag, event 2 had a scattering flag)
+    expected_flags_ebin2 = np.array(
+        [
+            False,
+            False,
+            True,
+            False,
+            False,
+            False,
+            False,
+            True,
+            False,
+            False,
+            False,
+            False,
+        ]
+    )
+    assert np.array_equal(valid_events[1], expected_flags_ebin2)
+    # Assert that for the third energy bin (7-18), all are false except events 4 and 7
+    # (event 5 was marked as an outlier and event 6 has an invalid ebin)
+    expected_flags_ebin3 = np.array(
+        [
+            False,
+            False,
+            False,
+            True,
+            False,
+            False,
+            True,
+            False,
+            False,
+            False,
+            False,
+            False,
+        ]
+    )
+    assert np.array_equal(valid_events[2], expected_flags_ebin3)
+
+
+@mock.patch("imap_processing.ultra.l1b.ultra_l1b_culling.sp.spkezr")
+def test_get_valid_events_per_energy_range_ultra45(mock_spkezr):
+    """Tests get_valid_events_per_energy_range function."""
+    np.random.seed(0)
+    mock_imap_state = np.random.random(6)  # Mock IMAP state for testing
+    mock_spkezr.return_value = (mock_imap_state, None)
+    energy_range_edges = np.array([3, 5, 7, 18])  # 3 example energy bins
+    energy = np.arange(18)
+
+    # mark all events with valid outlier and scattering flags and valid ebins.
+    de_dps_velocity = np.random.random((len(energy), 3))
+
+    de_dataset = xr.Dataset(
+        {
+            "de_dps_velocity": (("epoch", "component"), de_dps_velocity),
+            "event_times": ("epoch", np.full(len(energy), 798033671)),
+            "energy_spacecraft": ("epoch", energy),
+            "quality_outliers": ("epoch", np.full(len(energy), 0)),
+            "quality_scattering": ("epoch", np.full(len(energy), 0)),
+            "ebin": ("epoch", np.full(len(energy), 10)),
+        }
+    )
+    # ensure that all events fail the earth angle check by setting a very large
+    # keepout angle
+    keepout_angle = np.radians(360)
+    valid_events = get_valid_events_per_energy_range(
+        de_dataset, energy_range_edges, keepout_angle, 45
+    )
+
+    # although all events were valid for outliers, scattering, and ebin, all events
+    # failed the earth angle check for ultra45
+    assert not np.any(valid_events)
+
+
+@mock.patch(
+    "imap_processing.ultra.l1b.ultra_l1b_culling.UltraConstants.HIGH_ENERGY_CULL_CHANNEL",
+    2,
+)
+def test_flag_high_energy():
+    """Tests flag_high_energy function."""
+    # 7-18 is the culling energy bin and shown in the mock.patch above
+    # Flag energy ranges will compare the counts from this energy range at each spin bin
+    # to the culling threshold to determine if the spin bin should be flagged for high
+    # energy.
+    energy_range_edges = np.array([3, 5, 7, 18, 25])  # Example energy bin edges
+    # Spin bin 1 (events 0-3) 4 events fall within the culling energy bin
+    #   - This is above all the energy thresholds except the second one (10),
+    #   so it should be flagged for all energy ranges except flag #2
+    # Spin bin 2 (events 4-7) only 1 event falls within the culling energy bin
+    #   - This is above the lowest energy threshold (1) but below the rest, so
+    #   it should only be flagged with the last energy range flag (#3)
+    # Spin bin 3 (events 8-11) No events fall within the culling energy bin,
+    #   so it should not be flagged for any energy range
+    energy = np.array([17, 16, 12, 15, 5, 8, 4, 6, 4, 1, 22, 20])
+    cull_thresholds = np.array([3, 10, 2, 1])
+    de_dataset = xr.Dataset(
+        {
+            "de_event_met": ("epoch", np.arange(len(energy))),
+            "energy_spacecraft": ("epoch", energy),
+            "quality_outliers": ("epoch", np.full(len(energy), 0)),
+            "quality_scattering": ("epoch", np.full(len(energy), 0)),
+            "ebin": ("epoch", np.full(len(energy), 10)),
+        }
+    )
+    # make one ebin invalid to make sure the valid events filtering is working
+    de_dataset["ebin"].data[-1] = -1
+    spin_tbin_edges = np.arange(
+        start=0, stop=len(energy) + 1, step=4
+    )  # create spin bins of 4 seconds
+    energy_range_flags = get_energy_range_flags(energy_range_edges)
+    quality_flags = flag_high_energy(
+        de_dataset,
+        spin_tbin_edges,
+        energy_range_edges,
+        energy_range_flags,
+        cull_thresholds,
+        90,
+    )
+
+    # check shape
+    assert len(quality_flags) == len(spin_tbin_edges) - 1
+    # Assert that the first spin bin is flagged for high energy for all energy ranges
+    # except the second one
+    assert quality_flags[0] == (2**0 | 2**2 | 2**3)
+    # Assert that the second spin bin is only flagged for high energy for the last
+    # energy range
+    assert quality_flags[1] == 2**3
+    # Assert that the third spin bin is not flagged for any energy range
+    assert quality_flags[2] == 0
+
+
+@pytest.mark.external_test_data
+def test_validate_high_energy_cull():
+    """Validate that high energy spins are correctly flagged"""
+    # Mock thresholds to match the test data (I used fake ones to create more
+    # complexity)
+    mock_thresholds = np.array([0.05, 1.5, 0.6, 119.2, 0.2, 0.25]) * 20
+    # read test data from csv files
+    xspin = pd.read_csv(TEST_PATH / "extendedspin_test_data_repoint00047.csv")
+    expected_qf = pd.read_csv(
+        TEST_PATH / "validate_high_energy_culling_results_repoint00047.csv"
+    ).to_numpy()
+    de_df = pd.read_csv(TEST_PATH / "de_test_data_repoint00047.csv")
+    de_ds = xr.Dataset(
+        {
+            "de_event_met": ("epoch", de_df.event_times.values),
+            "energy_spacecraft": ("epoch", de_df.energy_spacecraft.values),
+            "quality_outliers": ("epoch", de_df.quality_outliers.values),
+            "quality_scattering": ("epoch", de_df.quality_scattering.values),
+            "ebin": ("epoch", de_df.ebin.values),
+        }
+    )
+    # Use constants from the code to ensure consistency with the actual culling code
+    spin_bin_size = UltraConstants.SPIN_BIN_SIZE
+    spin_tbin_edges = get_binned_spins_edges(
+        xspin.spin_number.values,
+        xspin.spin_period.values,
+        xspin.spin_start_time.values,
+        spin_bin_size,
+    )
+    intervals, _, _ = build_energy_bins()
+    # Get the energy ranges
+    energy_ranges = get_binned_energy_ranges(intervals)
+    flags = get_energy_range_flags(energy_ranges)
+    e_flags = flag_high_energy(
+        de_ds, spin_tbin_edges, energy_ranges, flags, mock_thresholds
+    )
+    # The ULTRA IT flags are shaped n_energy_ranges, spin_bin while the SDC
+    # is only spin_bin but different flags are set for different energy ranges, so we
+    # need to check each energy separately
+    # We also need to invert the expected flags since the ULTRA IT mask is True
+    # for good spins (counts below threshold) while the SDC quality flags are set
+    # for bad spins (counts exceed threshold).
+    for i in range(expected_qf.shape[0]):
+        np.testing.assert_array_equal(
+            (e_flags & 2**i) > 0,
+            ~expected_qf[i, :].astype(bool),
+            err_msg=f"High energy flag mismatch for energy range {i} with edges"
+            f" {energy_ranges[i]}-{energy_ranges[i + 1]}",
+        )
+
+
+def test_get_energy_range_flags():
+    """Tests get_binned_energy_range_flags function."""
+    # Get energy bins used at l1c
+    intervals, _, _ = build_energy_bins()
+    # Get the energy ranges
+    energy_ranges = get_binned_energy_ranges(intervals)
+    flags = get_energy_range_flags(energy_ranges)
+
+    np.testing.assert_array_equal(flags, 2 ** np.arange(6))

imap_processing/ultra/constants.py

@@ -188,11 +188,19 @@ class UltraConstants:
     # L1b extended spin culling parameters
     LOW_VOLTAGE_CULL_THRESHOLD = 3400.0
     SPIN_BIN_SIZE = 20
-    # TODO add thresholds for energies (different for 45 and 90)
     # Number of energy bins to use in energy dependent culling
     N_CULL_EBINS = 8
     # Bin to start culling at
     BASE_CULL_EBIN = 4
     # Angle threshold in radians for ULTRA 45 degree culling.
     # This is only needed for ULTRA 45 since earth may be in the FOV.
     EARTH_ANGLE_45_THRESHOLD = np.radians(15)
+    # An array of energy thresholds to use for culling. Each one corresponds to
+    # the number of energy bins used.
+    # n_bins=len(PSET_ENERGY_BIN_EDGES)[BASE_CULL_EBIN:] // N_CULL_EBINS
+    # an error will be raised if this does not match n_bins
+    HIGH_ENERGY_CULL_THRESHOLDS = (
+        np.array([2.0, 1.5, 0.6, 0.25, 0.25, 0.25]) * SPIN_BIN_SIZE
+    )
+    # Use the channel defined below to determine which spins are contaminated
+    HIGH_ENERGY_CULL_CHANNEL = 4

imap_processing/ultra/l1b/extendedspin.py

@@ -10,14 +10,15 @@
     count_rejected_events_per_spin,
     expand_bin_flags_to_spins,
     flag_attitude,
+    flag_high_energy,
     flag_hk,
     flag_imap_instruments,
     flag_low_voltage,
     flag_rates,
-    get_binned_energy_range_flags,
     get_binned_energy_ranges,
     get_binned_spins_edges,
     get_energy_histogram,
+    get_energy_range_flags,
     get_pulses_per_spin,
 )
 from imap_processing.ultra.l1c.l1c_lookup_utils import build_energy_bins
@@ -74,13 +75,27 @@ def calculate_extendedspin(
         spin, spin_period, spin_starttime, spin_bin_size
     )
     voltage_qf = flag_low_voltage(spin_tbin_edges, status_dataset)
+    voltage_qf = expand_bin_flags_to_spins(len(spin), voltage_qf, spin_bin_size)
     # Get energy bins used at l1c
     intervals, _, _ = build_energy_bins()
     # Get the energy ranges
     energy_ranges = get_binned_energy_ranges(intervals)
-    energy_bin_flags = get_binned_energy_range_flags(energy_ranges)
-    # TODO do not include low voltage spins in the calculation of the rest of the flag!!
-
+    energy_bin_flags = get_energy_range_flags(energy_ranges)
+    # Calculate the high energy quality flags using the de dataset with low voltage
+    # events removed. Use the same spin and energy bins that
+    # were used for low voltage flags to maintain consistency in the flags.
+    valid_voltage_spins = spin[np.where(voltage_qf == 0)]
+    valid_de_spins = np.isin(de_dataset["spin"].values, valid_voltage_spins)
+    de_dataset_filtered = de_dataset.isel(epoch=valid_de_spins)
+    energy_thresholds = UltraConstants.HIGH_ENERGY_CULL_THRESHOLDS
+    high_energy_qf = flag_high_energy(
+        de_dataset_filtered,
+        spin_tbin_edges,
+        energy_ranges,
+        energy_bin_flags,
+        energy_thresholds,
+        instrument_id,
+    )
     # Get the number of pulses per spin.
     pulses = get_pulses_per_spin(aux_dataset, rates_dataset)
 
@@ -118,7 +133,7 @@ def calculate_extendedspin(
     coin_per_spin[valid] = pulses.coin_per_spin[idx[valid]]
 
     # Expand binned quality flags to individual spins.
-    voltage_qf = expand_bin_flags_to_spins(len(spin), voltage_qf, spin_bin_size)
+    high_energy_qf = expand_bin_flags_to_spins(len(spin), high_energy_qf, spin_bin_size)
     # account for rates spins which are not in the direct event spins
     extendedspin_dict["start_pulses_per_spin"] = start_per_spin
     extendedspin_dict["stop_pulses_per_spin"] = stop_per_spin
@@ -134,9 +149,7 @@ def calculate_extendedspin(
     extendedspin_dict["quality_statistics"] = np.full_like(
         voltage_qf, ImapRatesUltraFlags.NONE.value, np.uint16
     )  # shape (nspin,)
-    extendedspin_dict["quality_high_energy"] = np.full_like(
-        voltage_qf, ImapRatesUltraFlags.NONE.value, np.uint16
-    )  # shape (nspin,)
+    extendedspin_dict["quality_high_energy"] = high_energy_qf  # shape (nspin,)
     # Add an array of flags for each energy bin. Shape: (n_energy_bins)
     extendedspin_dict["energy_range_flags"] = energy_bin_flags
     # Add energy ranges  Shape: (n_energy_bins + 1)