diff --git a/imap_processing/tests/external_test_data_config.py b/imap_processing/tests/external_test_data_config.py index 2e8bb9d3c..98a71415e 100644 --- a/imap_processing/tests/external_test_data_config.py +++ b/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", diff --git a/imap_processing/tests/ultra/unit/test_ultra_l1b_culling.py b/imap_processing/tests/ultra/unit/test_ultra_l1b_culling.py index f698d4989..2fbba6c56 100644 --- a/imap_processing/tests/ultra/unit/test_ultra_l1b_culling.py +++ b/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)) diff --git a/imap_processing/ultra/constants.py b/imap_processing/ultra/constants.py index 2090936e4..ce759f859 100644 --- a/imap_processing/ultra/constants.py +++ b/imap_processing/ultra/constants.py @@ -188,7 +188,6 @@ 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 @@ -196,3 +195,12 @@ class UltraConstants: # 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 diff --git a/imap_processing/ultra/l1b/extendedspin.py b/imap_processing/ultra/l1b/extendedspin.py index 74328e1f4..a68f0fb6a 100644 --- a/imap_processing/ultra/l1b/extendedspin.py +++ b/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) diff --git a/imap_processing/ultra/l1b/ultra_l1b_culling.py b/imap_processing/ultra/l1b/ultra_l1b_culling.py index 5320b95e2..946364923 100644 --- a/imap_processing/ultra/l1b/ultra_l1b_culling.py +++ b/imap_processing/ultra/l1b/ultra_l1b_culling.py @@ -680,6 +680,140 @@ def flag_low_voltage( return quality_flags +def flag_high_energy( + de_dataset: xr.Dataset, + spin_tbin_edges: NDArray, + energy_ranges: NDArray, + energy_range_flags: np.ndarray, + energy_thresholds: np.ndarray = UltraConstants.HIGH_ENERGY_CULL_THRESHOLDS, + sensor_id: int = 90, +) -> NDArray: + """ + Flag high energy events. + + Parameters + ---------- + de_dataset : xr.Dataset + Direct event dataset. + spin_tbin_edges : NDArray + Edges of the spin time bins. + energy_ranges : numpy.ndarray + Array of energy range edges. + energy_range_flags : numpy.ndarray + Array of quality flag values corresponding to each energy range. + energy_thresholds : numpy.ndarray + Array of count thresholds for flagging high energy events corresponding to + each energy range. + sensor_id : int + Sensor ID (e.g., 45 or 90). + + Returns + ------- + quality_flags : numpy.ndarray + Quality flags. + """ + cull_channel = UltraConstants.HIGH_ENERGY_CULL_CHANNEL + valid_events_per_energy = get_valid_events_per_energy_range( + de_dataset, energy_ranges, UltraConstants.EARTH_ANGLE_45_THRESHOLD, sensor_id + ) + # check to make sure the number of energy ranges matches the number of energy range + # flags + num_e_ranges = valid_events_per_energy.shape[0] + if num_e_ranges != len(energy_range_flags) or num_e_ranges != len( + energy_thresholds + ): + raise ValueError( + f"Number of energy ranges ({num_e_ranges}) does not match number of energy" + f" range flags ({len(energy_range_flags)}) or expected number of " + f"energy range thresholds ({len(energy_thresholds)})." + ) + if cull_channel >= num_e_ranges: + raise ValueError( + f"HIGH_ENERGY_CULL_CHANNEL ({cull_channel}) is out of bounds" + f" for {num_e_ranges} energy ranges." + ) + + # Initialize all spin bins to have no high energy flag + spin_bin_size = len(spin_tbin_edges) - 1 + quality_flags = np.full( + spin_bin_size, ImapRatesUltraFlags.NONE.value, dtype=np.uint16 + ) + # Get valid events and counts at each spin bin for the + # designated culling channel. + cull_channel_events = valid_events_per_energy[cull_channel] + # get each valid event count per spin bin for the culling channel + cull_channel_counts = np.histogram( + de_dataset["de_event_met"].values[cull_channel_events], spin_tbin_edges + )[0] + # loop through each energy range + for flag, e_threshold in zip(energy_range_flags, energy_thresholds, strict=False): + quality_flags[cull_channel_counts >= e_threshold] |= flag + + return quality_flags + + +def get_valid_events_per_energy_range( + de_dataset: xr.Dataset, energy_ranges: NDArray, earth_ang_45: float, sensor_id: int +) -> NDArray: + """ + Get valid events per energy range. + + Parameters + ---------- + de_dataset : xr.Dataset + Direct event dataset. + energy_ranges : numpy.ndarray + Array of energy range edges. + earth_ang_45 : float + Earth angle to use for culling in ULTRA 45. + sensor_id : int + Sensor ID (e.g., 45 or 90). + + Returns + ------- + valid_events_per_range : numpy.ndarray + A boolean array of shape (n_energy_ranges, n_events). + """ + event_energies = de_dataset["energy_spacecraft"].values + valid_events = np.zeros((len(energy_ranges) - 1, len(event_energies)), dtype=bool) + valid_outliers = de_dataset["quality_outliers"].values == 0 + valid_scattering = de_dataset["quality_scattering"].values == 0 + # TODO what about species non-proton? For those psets dont cull based on + # High energy? + ebin = de_dataset["ebin"].values + valid_ebin = np.isin(ebin, UltraConstants.TOFXPH_SPECIES_GROUPS["proton"]) + for i in range(len(energy_ranges) - 1): + energy_mask = (event_energies >= energy_ranges[i]) & ( + event_energies < energy_ranges[i + 1] + ) + if not np.any(energy_mask): + continue + # subset the dataset to events within the energy range + de_dataset_subset = de_dataset.isel(epoch=energy_mask) + valid_earth_angle = np.full(np.sum(energy_mask), True, dtype=bool) + # For ultra45, also apply an Earth angle cut to remove times when + # the Earth is in the field of view. ULTRA 90 does not require this since Earth + # is always outside the field of view. + if sensor_id == 45: + valid_earth_angle = get_valid_earth_angle_events( + de_dataset_subset, earth_ang_45 + ) + + # Flag events at the valid energy ranges if they meet all the criteria for + # valid events: not flagged as outliers, not flagged as scattering, + # in a valid ebin, and (for ultra45) have a valid Earth angle. + valid_events[i, energy_mask] = np.logical_and.reduce( + [ + valid_outliers[energy_mask], + valid_scattering[energy_mask], + valid_ebin[energy_mask], + valid_earth_angle, + ] + ) + + return valid_events + + def get_valid_earth_angle_events( de_dataset_subset: xr.Dataset, earth_ang_45: float = UltraConstants.EARTH_ANGLE_45_THRESHOLD, @@ -735,7 +869,7 @@ def get_valid_earth_angle_events( return sep_angle > earth_ang_45 -def get_binned_energy_range_flags(energy_ranges_edges: NDArray) -> NDArray: +def get_energy_range_flags(energy_ranges_edges: NDArray) -> NDArray: """ Get the energy bin flags for energy dependent culling.