diff --git a/imap_processing/tests/ultra/unit/conftest.py b/imap_processing/tests/ultra/unit/conftest.py index e3b4f79fb..1a3fe386b 100644 --- a/imap_processing/tests/ultra/unit/conftest.py +++ b/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 } ) diff --git a/imap_processing/tests/ultra/unit/test_ultra_l1b.py b/imap_processing/tests/ultra/unit/test_ultra_l1b.py index 690070648..3132fe0ab 100644 --- a/imap_processing/tests/ultra/unit/test_ultra_l1b.py +++ b/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 diff --git a/imap_processing/tests/ultra/unit/test_ultra_l1c_pset_bins.py b/imap_processing/tests/ultra/unit/test_ultra_l1c_pset_bins.py index 2d6f24957..655af6306 100644 --- a/imap_processing/tests/ultra/unit/test_ultra_l1c_pset_bins.py +++ b/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 diff --git a/imap_processing/ultra/l1b/ultra_l1b_culling.py b/imap_processing/ultra/l1b/ultra_l1b_culling.py index 0f22af6b8..e4c7afcf5 100644 --- a/imap_processing/ultra/l1b/ultra_l1b_culling.py +++ b/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 diff --git a/imap_processing/ultra/l1c/helio_pset.py b/imap_processing/ultra/l1c/helio_pset.py index 7375d71fc..58e9e02a3 100644 --- a/imap_processing/ultra/l1c/helio_pset.py +++ b/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 diff --git a/imap_processing/ultra/l1c/spacecraft_pset.py b/imap_processing/ultra/l1c/spacecraft_pset.py index 3506c0e26..59b1b5706 100644 --- a/imap_processing/ultra/l1c/spacecraft_pset.py +++ b/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 diff --git a/imap_processing/ultra/l1c/ultra_l1c_pset_bins.py b/imap_processing/ultra/l1c/ultra_l1c_pset_bins.py index ece5f0f0e..507ca75c5 100644 --- a/imap_processing/ultra/l1c/ultra_l1c_pset_bins.py +++ b/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(