ULTRA l1c fix exposure time calculation

imap_processing/tests/ultra/unit/conftest.py

@@ -39,6 +39,11 @@
     ULTRA_RATES,
 )
 from imap_processing.ultra.l1a.ultra_l1a import ultra_l1a
+from imap_processing.ultra.l1b.ultra_l1b_culling import (
+    get_binned_energy_ranges,
+    get_energy_range_flags,
+)
+from imap_processing.ultra.l1c.l1c_lookup_utils import build_energy_bins
 from imap_processing.utils import packet_file_to_datasets
 
 
@@ -644,13 +649,27 @@ def mock_helio_pointing_lookups():
 @pytest.fixture
 def mock_goodtimes_dataset():
     """Create a mock goodtimes dataset."""
+    # Set up bit flags
+    intervals, _, _ = build_energy_bins()
+    energy_ranges = get_binned_energy_ranges(intervals)
+    energy_flags = get_energy_range_flags(energy_ranges)
+    nspins = 100
+    flags = 2 ** np.arange(9)
+    quality = np.zeros(nspins, dtype=np.uint16)
+    quality[0] = flags[0]  # Set the first flag for the first spin
+    quality[1] = flags[1]  # Set the second flag for the second
+    quality[2] = flags[2]  # Set the third flag for the third spin
     return xr.Dataset(
         {
-            "spin_number": ("epoch", np.zeros(5)),
-            "energy_range_flags": ("energy_flags", np.zeros(10, dtype=np.uint16)),
-            "quality_low_voltage": ("spin_number", np.zeros(5, dtype=np.uint16)),
-            "quality_high_energy": ("spin_number", np.zeros(5, dtype=np.uint16)),
-            "quality_statistics": ("spin_number", np.zeros(5, dtype=np.uint16)),
-            "energy_range_edges": ("energy_ranges", np.zeros(11, dtype=np.uint16)),
+            "spin_number": ("epoch", np.zeros(nspins)),
+            "energy_range_flags": ("energy_flags", energy_flags),
+            "quality_low_voltage": ("spin_number", quality),
+            "quality_high_energy": ("spin_number", np.zeros(nspins, dtype=np.uint16)),
+            "quality_statistics": ("spin_number", np.zeros(nspins, dtype=np.uint16)),
+            "energy_range_edges": ("energy_ranges", energy_ranges),
+            "spin_period": (
+                "spin_number",
+                np.full(nspins, 15),
+            ),  # nominal spin period of 15 seconds
         }
     )

imap_processing/tests/ultra/unit/test_ultra_l1b.py

@@ -63,12 +63,19 @@ def mock_data_l1b_extendedspin_dict():
     )
     spin_start_time = np.array([0, 1, 2], dtype="uint64")
     quality = np.zeros((2, 3), dtype="uint16")
+    # These should be shape: (3,)
+    energy_dep_flags = np.zeros(len(spin), dtype="uint16")
     data_dict = {
         "epoch": epoch,
         "spin_number": spin,
         "energy_bin_geometric_mean": energy,
         "spin_start_time": spin_start_time,
         "quality_ena_rates": quality,
+        "quality_low_voltage": energy_dep_flags,
+        "quality_high_energy": energy_dep_flags,
+        "quality_statistics": energy_dep_flags,
+        "energy_range_flags": np.ones(5, dtype=np.uint16),
+        "energy_range_edges": np.ones(4, dtype=np.float64),
     }
     return data_dict
 

imap_processing/tests/ultra/unit/test_ultra_l1c_pset_bins.py

@@ -1,5 +1,6 @@
 "Tests pointing sets"
 
+import logging
 from unittest import mock
 
 import astropy_healpix.healpy as hp
@@ -10,6 +11,7 @@
 from scipy import interpolate
 
 from imap_processing import imap_module_directory
+from imap_processing.ultra.constants import UltraConstants
 from imap_processing.ultra.l1c import ultra_l1c_pset_bins
 from imap_processing.ultra.l1c.spacecraft_pset import (
     calculate_fwhm_spun_scattering,
@@ -395,12 +397,14 @@ def test_get_spacecraft_exposure_times(
     ancillary_files,
     use_fake_spin_data_for_time,
     aux_dataset,
+    mock_goodtimes_dataset,
+    caplog,
 ):
     """Test get_spacecraft_exposure_times function."""
     data_start_time = 445015665.0
     data_end_time = 453070000.0
     use_fake_spin_data_for_time(data_start_time, data_end_time)
-    steps = 500  # reduced for testing
+    steps = 200  # reduced for testing
 
     pix = 786
     mock_theta = np.random.uniform(-60, 60, (steps, pix))
@@ -417,21 +421,37 @@ def test_get_spacecraft_exposure_times(
         )
     )
     boundary_sf = xr.DataArray(np.ones((steps, pix)), dims=("spin_phase_step", "pixel"))
-    exposure_pointing, deadtimes = get_spacecraft_exposure_times(
+    caplog.set_level(logging.INFO)
+    (
+        exposure_pointing,
+        deadtimes,
+    ) = get_spacecraft_exposure_times(
         rates_dataset,
         pixels_below_threshold,
         boundary_sf,
         aux_dataset,
-        (
-            data_start_time,
-            data_start_time,
-        ),
-        46,  # number of energy bins
-        pix,
+        build_energy_bins()[2],
+        goodtimes_dataset=mock_goodtimes_dataset,
     )
     np.testing.assert_array_equal(exposure_pointing.shape, (46, pix))
     np.testing.assert_array_equal(deadtimes.shape, (steps,))
 
+    # Check that the number of good spins per energy bin is logged correctly
+    # The first 3 energy bins were not used in the goodtimes culling and therefore
+    # should have all 100 spins.
+    # Energy range 1,2,3 (which maps to the next UltraConstants.N_CULL_EBINS * 3 bins)
+    # should have 99 spins because there was a flag set to true for one spin
+    # in the mock_goodtimes_dataset. The rest of the energy bins should have 100 spins
+    # because the goodtimes dataset did not flag any spins for those energy bins.
+    expected_good_spins = np.full(46, 100.0)
+    expected_good_spins[
+        UltraConstants.BASE_CULL_EBIN : UltraConstants.N_CULL_EBINS * 3
+        + UltraConstants.BASE_CULL_EBIN
+    ] = 99.0
+    # The goodtimes dataset has flags set to True for energy bin 0-3 and for the
+    # first three spins.
+    assert f"Found {expected_good_spins.tolist()} valid spins" in caplog.text
+
 
 def test_get_spacecraft_background_rates(
     rates_l1_test_path, use_fake_spin_data_for_time, ancillary_files

imap_processing/ultra/l1b/ultra_l1b_culling.py

@@ -561,33 +561,30 @@ def get_energy_and_spin_dependent_rejection_mask(
         goodtimes_dataset[flag_name].values
         for flag_name in ENERGY_DEPENDENT_SPIN_QUALITY_FLAG_FILTERS
     ]
-    ebin_flags = goodtimes_dataset["energy_range_flags"].values
-    # Create a dict of spin_number to index in the goodtimes dataset
-    spin_to_idx = {
-        spin: idx for idx, spin in enumerate(goodtimes_dataset["spin_number"].values)
-    }
-
     # Initialize all events to not rejected
-    rejected = np.full(energy.shape, False, dtype=bool)
-    # loop through each energy bin and flag events that fall within an energy
-    # bin and have the corresponding energy bin flag set in the goodtimes dataset.
-    for i in range(len(energy_range_edges) - 1):
-        mask = (energy >= energy_range_edges[i]) & (energy < energy_range_edges[i + 1])
-        goodtimes_inds = [spin_to_idx[spin] for spin in spin_number[mask]]
-        # Get the flag value for the current energy bin
-        energy_bin_flag = ebin_flags[i]
-        # If the flag is set for any of the quality arrays, then reject
-        # the event.
-        flagged_at_spins = (
-            np.bitwise_or.reduce(
-                [qf[goodtimes_inds] & energy_bin_flag for qf in flag_arrays]
-            )
-            > 0
-        )
-
-        # Mark flagged events as rejected
-        mask_indices = np.where(mask)[0]
-        rejected[mask_indices[flagged_at_spins]] = True
+    rejected = np.zeros_like(energy, dtype=bool)
+    ebin_flags = goodtimes_dataset["energy_range_flags"].values
+    # Get the index of the spin number in the goodtimes dataset for each event
+    # all spin numbers should be present in the goodtimes dataset since we have already
+    # filtered any events that are not
+    spin_idx = np.searchsorted(goodtimes_dataset.spin_number, spin_number)
+    event_energy_bins: NDArray = (np.digitize(energy, energy_range_edges) - 1).astype(
+        np.intp
+    )
+    in_valid_bin = (event_energy_bins >= 0) & (event_energy_bins < len(ebin_flags))
+    # get the flags for each event
+    event_flags = np.zeros_like(energy, dtype=np.uint16)
+    event_flags[in_valid_bin] = ebin_flags[event_energy_bins[in_valid_bin]]
+    for qf_array in flag_arrays:
+        # select the quality flag for each event
+        quality_flags_at_events = qf_array[spin_idx]
+        # If that flag is "turned on" for the spin of that event, and the event is in
+        # an energy bin that is flagged for culling, then we reject that event.
+        rejected |= quality_flags_at_events & event_flags > 0
+
+    logger.info(
+        "Rejected %d events based on energy and spin dependent flags.", np.sum(rejected)
+    )
 
     return rejected
 

imap_processing/ultra/l1c/helio_pset.py

@@ -173,11 +173,11 @@ def calculate_helio_pset(
         pixels_below_scattering,
         boundary_scale_factors,
         aux_dataset,
-        pointing_range_met,
-        n_energy_bins=len(energy_bin_geometric_means),
+        energy_bins=energy_bin_geometric_means,
         sensor_id=sensor_id,
         ancillary_files=ancillary_files,
         apply_bsf=apply_bsf,
+        goodtimes_dataset=goodtimes_dataset,
     )
     logger.info("Calculating spun efficiencies and geometric function.")
     # calculate efficiency and geometric function as a function of energy

imap_processing/ultra/l1c/spacecraft_pset.py

@@ -162,11 +162,11 @@ def calculate_spacecraft_pset(
         valid_spun_pixels,
         boundary_scale_factors,
         aux_dataset,
-        pointing_range_met,
-        n_energy_bins=len(energy_bin_geometric_means),
+        energy_bins=energy_bin_geometric_means,
         sensor_id=sensor_id,
         ancillary_files=ancillary_files,
         apply_bsf=apply_bsf,
+        goodtimes_dataset=goodtimes_dataset,
     )
     logger.info("Calculating spun efficiencies and geometric function.")
     # calculate efficiency and geometric function as a function of energy

imap_processing/ultra/l1c/ultra_l1c_pset_bins.py

@@ -11,15 +11,15 @@
 from imap_processing.spice.geometry import (
     cartesian_to_spherical,
 )
-from imap_processing.spice.spin import (
-    get_spin_data,
-)
 from imap_processing.ultra.constants import UltraConstants
 from imap_processing.ultra.l1b.lookup_utils import (
     get_geometric_factor,
     get_image_params,
     load_geometric_factor_tables,
 )
+from imap_processing.ultra.l1b.quality_flag_filters import (
+    ENERGY_DEPENDENT_SPIN_QUALITY_FLAG_FILTERS,
+)
 from imap_processing.ultra.l1b.ultra_l1b_culling import (
     get_pulses_per_spin,
 )
@@ -385,6 +385,7 @@ def calculate_exposure_time(
     -------
     exposure_pointing: xarray.DataArray
         Adjusted exposure times accounting for dead time.
+        Shape: ``(energy, pixel)``.
     """
     # nominal spin phase step.
     nominal_ms_step = 15 / valid_spun_pixels.shape[0]  # time step
@@ -407,8 +408,8 @@ def get_spacecraft_exposure_times(
     valid_spun_pixels: xr.DataArray,
     boundary_scale_factors: xr.DataArray,
     aux_dataset: xr.Dataset,
-    pointing_range_met: tuple[float, float],
-    n_energy_bins: int,
+    energy_bins: np.ndarray,
+    goodtimes_dataset: xr.Dataset,
     sensor_id: int | None = None,
     ancillary_files: dict | None = None,
     apply_bsf: bool = True,
@@ -431,10 +432,13 @@ def get_spacecraft_exposure_times(
         Boundary scale factors for each pixel at each spin phase.
     aux_dataset : xarray.Dataset
         Auxiliary dataset containing spin information.
-    pointing_range_met : tuple
-        Start and stop time of the pointing period in mission elapsed time.
-    n_energy_bins : int
-        Number of energy bins.
+    energy_bins : np.ndarray
+        Array of energy bin geometric means.
+    goodtimes_dataset : xarray.Dataset
+        Dataset containing the quality-filtered spins with energy dependent quality
+        flags (quality_low_voltage, quality_high_energy, quality_statistics).
+        Exposure times are computed using only these good spins
+        and can be adjusted per energy bin based on quality flags.
     sensor_id : int, optional
         Sensor ID, either 45 or 90.
     ancillary_files : dict, optional
@@ -448,6 +452,7 @@ def get_spacecraft_exposure_times(
         Total exposure times of pixels in a
         Healpix tessellation of the sky
         in the pointing (dps) frame.
+        Shape: (n_energy_bins, n_pix).
     nominal_deadtime_ratios : np.ndarray
         Deadtime ratios at each spin phase step (1ms res).
     """
@@ -464,35 +469,60 @@ def get_spacecraft_exposure_times(
     exposure_time = calculate_exposure_time(
         nominal_deadtime_ratios, valid_spun_pixels, boundary_scale_factors, apply_bsf
     )
-    # Use the universal spin table to determine the actual number of spins
+    if exposure_time.ndim != 2:
+        raise ValueError(
+            "Exposure time must be 2D with dimensions ('energy', 'pixel'); "
+            f"got dims {exposure_time.dims} and shape {exposure_time.shape}."
+        )
     nominal_spin_seconds = 15.0
-    spin_data = get_spin_data()
-    # Filter for spins only in pointing
-    spin_data = spin_data[
-        (spin_data["spin_start_met"] >= pointing_range_met[0])
-        & (spin_data["spin_start_met"] <= pointing_range_met[1])
+    # Use filtered spins from goodtimes dataset to include only the spins that
+    # passed the quality flag filtering.
+    spin_periods = goodtimes_dataset["spin_period"].values
+    energy_flags = goodtimes_dataset["energy_range_flags"].values
+    # only get valid flags for the energy bins we are using at l1c
+    energy_flags = energy_flags[energy_flags > 0]
+    # Get the quality flag arrays "turned on" for energy dependent culling from the
+    # goodtimes dataset.
+    flag_arrays = [
+        goodtimes_dataset[flag_name].values
+        for flag_name in ENERGY_DEPENDENT_SPIN_QUALITY_FLAG_FILTERS
     ]
-    # Get only valid spin data
-    valid_mask = (spin_data["spin_phase_valid"].values == 1) & (
-        spin_data["spin_period_valid"].values == 1
+    bin_to_range = np.digitize(energy_bins, goodtimes_dataset.energy_range_edges)
+    valid_spins = (
+        np.bitwise_or.reduce(flag_arrays)[np.newaxis, :] & energy_flags[:, np.newaxis]
+    ) == 0
+    # Pad valid_spins with arrays of all true at either end to account for energy bins
+    # that fall outside the range of the goodtimes dataset energy edges. Energy bins
+    # outside these ranges were not included in the quality flag filtering, so they are
+    # considered valid (all true).
+    valid_spins_padded = np.pad(
+        valid_spins,
+        pad_width=((1, 1), (0, 0)),
+        # pad only the energy axis
+        mode="constant",
+        constant_values=True,
+    )
+    # Select the valid spins for each energy bin using the bin_to_range indices
+    good_spins_per_ebin = valid_spins_padded[bin_to_range]
+    # Broadcast spin periods to have the correct shape e.g. (energy_bins, spins)
+    spin_periods_2d = np.broadcast_to(
+        spin_periods[np.newaxis, :], good_spins_per_ebin.shape
     )
-    n_spins_in_pointing: float = np.sum(
-        spin_data[valid_mask].spin_period_sec / nominal_spin_seconds
+    # Calculate total normalized spin count using spin periods from goodtimes
+    # Shape (n_energy_bins)
+    n_spins_in_pointing = (
+        np.sum(spin_periods_2d, axis=1, where=good_spins_per_ebin)
+        / nominal_spin_seconds
     )
+
     logger.info(
-        f"Calculated total spins universal spin table. Found {n_spins_in_pointing} "
-        f"valid spins."
+        f"Calculated total spins. Found {n_spins_in_pointing.tolist()} valid spins per "
+        f"energy range."
     )
-    # Adjust exposure time by the actual number of valid spins in the pointing
-    exposure_pointing_adjusted = n_spins_in_pointing * exposure_time
-
-    if exposure_pointing_adjusted.shape[0] != n_energy_bins:
-        exposure_pointing_adjusted = np.repeat(
-            exposure_pointing_adjusted,
-            n_energy_bins,
-            axis=0,
-        )
-    return exposure_pointing_adjusted.values, nominal_deadtime_ratios.values
+    # Shape (n_energy_bins, n_pix)
+    exposure_pointing_adjusted = exposure_time.data * n_spins_in_pointing[:, np.newaxis]
+
+    return exposure_pointing_adjusted, nominal_deadtime_ratios.values
 
 
 def get_efficiencies_and_geometric_function(