diff --git a/imap_processing/cli.py b/imap_processing/cli.py index 1c2d21528..2c714b16c 100644 --- a/imap_processing/cli.py +++ b/imap_processing/cli.py @@ -1019,6 +1019,8 @@ def do_processing( elif self.data_level == "l1b": data_dict = {} science_files = dependencies.get_file_paths(source="lo", data_type="l1a") + science_files += dependencies.get_file_paths(source="lo", data_type="l1b") + ancillary_files = dependencies.get_file_paths( source="lo", data_type="ancillary" ) diff --git a/imap_processing/lo/l1b/lo_l1b.py b/imap_processing/lo/l1b/lo_l1b.py index 6e7d71fde..053cd7425 100644 --- a/imap_processing/lo/l1b/lo_l1b.py +++ b/imap_processing/lo/l1b/lo_l1b.py @@ -85,6 +85,11 @@ def lo_l1b( ds = l1b_histrates(sci_dependencies, anc_dependencies, attr_mgr_l1b) datasets_to_return.append(ds) + if descriptor == "derates": + logger.info("\nProcessing IMAP-Lo L1B DE Rates...") + ds = calculate_de_rates(sci_dependencies, anc_dependencies, attr_mgr_l1b) + datasets_to_return.append(ds) + return datasets_to_return @@ -1617,6 +1622,193 @@ def calculate_histogram_rates( return l1b_histrates +def calculate_de_rates( + sci_dependencies: dict, + anc_dependencies: list, + attr_mgr_l1b: ImapCdfAttributes, +) -> xr.Dataset: + """ + Calculate direct event rates histograms. + + The histograms are per ASC (28 spins), so we need to + regroup the individual DEs from the l1b_de dataset into + their associated ASC and then bin them by ESA / spin bin. + + Parameters + ---------- + sci_dependencies : dict + The science dependencies for the derates product. + anc_dependencies : list + List of ancillary file paths. + attr_mgr_l1b : ImapCdfAttributes + Attribute manager used to get the L1B derates dataset attributes. + + Returns + ------- + l1b_derates : xr.Dataset + Dataset containing DE rates histograms. + """ + l1b_de = sci_dependencies["imap_lo_l1b_de"] + l1a_spin = sci_dependencies["imap_lo_l1a_spin"] + l1b_nhk = sci_dependencies["imap_lo_l1b_nhk"] + # Set the asc_start for each DE by removing the average spin cycle + # which is a function of esa_step (see set_spin_cycle function) + # spin_cycle is an average over esa steps and spins per asc, so finding + # the "average" spin that an esa step occurred at. + asc_start = l1b_de["spin_cycle"] - (7 + (l1b_de["esa_step"] - 1) * 2) + + # Get unique ASC values and create a mapping from asc_start to index + unique_asc, unique_idx, asc_idx = np.unique( + asc_start.values, return_index=True, return_inverse=True + ) + num_asc = len(unique_asc) + + # Pre-extract arrays for faster access (avoid repeated xarray indexing) + esa_step_idx = l1b_de["esa_step"].values - 1 # Convert to 0-based index + # Convert spin_bin from 0.1 degree bins to 6 degree bins for coarse histograms + spin_bin = l1b_de["spin_bin"].values // 60 + species = l1b_de["species"].values + coincidence_type = l1b_de["coincidence_type"].values + + if len(anc_dependencies) == 0: + logger.warning("No ancillary dependencies provided, using linear stepping.") + energy_step_mapping = np.arange(7) + else: + # An array mapping esa step index to esa level for resweeping + energy_step_mapping = _get_esa_level_indices(asc_start, anc_dependencies) + + # exposure time shape: (num_asc, num_esa_steps) + exposure_time = np.zeros((num_asc, 7), dtype=float) + # exposure_time_6deg = 4 * avg_spin_per_asc / 60 + # 4 sweeps per ASC (28 / 7) in 60 bins + asc_avg_spin_durations = 4 * l1b_de["avg_spin_durations"].data[unique_idx] / 60 + np.add.at( + exposure_time, + (slice(None), energy_step_mapping), + asc_avg_spin_durations[:, np.newaxis], + ) + + # Create output arrays + output_shape = (num_asc, 7, 60) + h_counts = np.zeros(output_shape) + o_counts = np.zeros(output_shape) + triple_counts = np.zeros(output_shape) + double_counts = np.zeros(output_shape) + + # Species masks + h_mask = species == "H" + o_mask = species == "O" + + # Coincidence type masks + triple_types = ["111111", "111100", "111000"] + double_types = [ + "110100", + "110000", + "101101", + "101100", + "101000", + "100100", + "100101", + "100000", + "011100", + "011000", + "010100", + "010101", + "010000", + "001100", + "001101", + "001000", + ] + triple_mask = np.isin(coincidence_type, triple_types) + double_mask = np.isin(coincidence_type, double_types) + + # Vectorized histogramming using np.add.at with full index arrays + np.add.at(h_counts, (asc_idx[h_mask], esa_step_idx[h_mask], spin_bin[h_mask]), 1) + np.add.at(o_counts, (asc_idx[o_mask], esa_step_idx[o_mask], spin_bin[o_mask]), 1) + np.add.at( + triple_counts, + (asc_idx[triple_mask], esa_step_idx[triple_mask], spin_bin[triple_mask]), + 1, + ) + np.add.at( + double_counts, + (asc_idx[double_mask], esa_step_idx[double_mask], spin_bin[double_mask]), + 1, + ) + + ds = xr.Dataset( + coords={ + # ASC start time in TTJ2000ns + "epoch": l1a_spin["epoch"], + "esa_step": np.arange(7), + "spin_bin": np.arange(60), + }, + ) + ds["h_counts"] = xr.DataArray( + h_counts, + dims=["epoch", "esa_step", "spin_bin"], + ) + ds["o_counts"] = xr.DataArray( + o_counts, + dims=["epoch", "esa_step", "spin_bin"], + ) + ds["triple_counts"] = xr.DataArray( + triple_counts, + dims=["epoch", "esa_step", "spin_bin"], + ) + ds["double_counts"] = xr.DataArray( + double_counts, + dims=["epoch", "esa_step", "spin_bin"], + ) + ds["exposure_time"] = xr.DataArray( + exposure_time, + dims=["epoch", "esa_step"], + ) + ds["h_rates"] = ds["h_counts"] / ds["exposure_time"] + ds["o_rates"] = ds["o_counts"] / ds["exposure_time"] + ds["triple_rates"] = ds["triple_counts"] / ds["exposure_time"] + ds["double_rates"] = ds["double_counts"] / ds["exposure_time"] + + # (N, 7) + unique_asc = xr.DataArray(unique_asc, dims=["epoch"]) + ds["spin_cycle"] = unique_asc + 7 + (ds["esa_step"] - 1) * 2 + + # TODO: Add badtimes + ds["badtime"] = xr.zeros_like(ds["epoch"], dtype=int) + + pivot_angle = _get_nearest_pivot_angle(ds["epoch"].values[0], l1b_nhk) + ds["pivot_angle"] = xr.DataArray([pivot_angle], dims=["pivot_angle"]) + + pointing_start_met, pointing_end_met = get_pointing_times( + ttj2000ns_to_met(ds["epoch"].values[0].item()) + ) + ds = set_esa_mode(pointing_start_met, pointing_end_met, anc_dependencies, ds) + + ds.attrs = attr_mgr_l1b.get_global_attributes("imap_lo_l1b_derates") + ds["epoch"].attrs = attr_mgr_l1b.get_variable_attributes("epoch") + + return ds + + +def _get_nearest_pivot_angle(epoch: int, ds_nhk: xr.Dataset) -> float: + """ + Get the nearest pivot angle for the given epoch from the NHK dataset. + + Parameters + ---------- + epoch : int + The epoch in TTJ2000ns format. + ds_nhk : xr.Dataset + The NHK dataset containing pivot angle information. + + Returns + ------- + pivot_angle : float + The nearest pivot angle for the given epoch. + """ + return ds_nhk["pcc_cumulative_cnt_pri"].sel(epoch=epoch, method="nearest").item() + + def _get_esa_level_indices(epochs: np.ndarray, anc_dependencies: list) -> np.ndarray: """ Get the ESA level indices (reswept indices) for the given epochs. diff --git a/imap_processing/tests/lo/test_lo_l1b.py b/imap_processing/tests/lo/test_lo_l1b.py index 7080e1ddf..b6879ad8a 100644 --- a/imap_processing/tests/lo/test_lo_l1b.py +++ b/imap_processing/tests/lo/test_lo_l1b.py @@ -9,6 +9,7 @@ from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes from imap_processing.cdf.utils import load_cdf from imap_processing.lo.l1b.lo_l1b import ( + calculate_de_rates, calculate_histogram_rates, calculate_tof1_for_golden_triples, convert_start_end_acq_times, @@ -1275,3 +1276,123 @@ def test_set_spin_cycle_from_spin_data_insufficient_spins(): # Verify spin_cycle shape matches filtered data assert result["spin_cycle"].shape == (1, 7) + + +@patch( + "imap_processing.lo.l1b.lo_l1b.get_pointing_times", + return_value=(473389199, 473472001), +) +@patch( + "imap_processing.lo.l1b.lo_l1b._get_esa_level_indices", + return_value=np.arange(7), +) +def test_calculate_de_rates( + mock_get_esa_level_indices, mock_get_pointing_times, attr_mgr_l1b, anc_dependencies +): + """Test the calculate_de_rates function.""" + # Use MET times from the test sweep table (2025-01-01) + met_start = 473389200 + epoch_time = met_to_ttj2000ns([met_start, met_start + 15 * 28]) + + # Create individual epochs for each direct event in TTJ2000ns + de_epochs = met_to_ttj2000ns( + [ + met_start + 10, + met_start + 20, + met_start + 30, + met_start + 15 * 28 + 10, + met_start + 15 * 28 + 20, + ] + ) + + # Create a simple l1b_de dataset with a few direct events + l1b_de = xr.Dataset( + { + "spin_cycle": ("epoch", [7, 9, 11, 35, 37]), + "esa_step": ("epoch", [1, 2, 3, 1, 2]), + "spin_bin": ("epoch", [0, 120, 240, 60, 180]), + "species": ("epoch", ["H", "O", "H", "H", "O"]), + "coincidence_type": ( + "epoch", + ["111111", "110100", "111000", "101000", "100100"], + ), + "avg_spin_durations": ("epoch", [15.0, 15.0, 15.0, 15.0, 15.0]), + }, + coords={"epoch": de_epochs}, + ) + + # Create l1a_spin dataset + l1a_spin = xr.Dataset( + { + "shcoarse": ("epoch", [met_start, met_start + 15 * 28]), + "num_completed": ("epoch", [28, 28]), + "acq_start_sec": ("epoch", [met_start, met_start + 15 * 28]), + "acq_start_subsec": ("epoch", [0, 0]), + "acq_end_sec": ("epoch", [met_start + 15 * 28, met_start + 2 * 15 * 28]), + "acq_end_subsec": ("epoch", [0, 0]), + }, + coords={"epoch": epoch_time}, + ) + + # Create l1b_nhk dataset with pivot angle information + l1b_nhk = xr.Dataset( + {"pcc_cumulative_cnt_pri": ("epoch", [45.0])}, + coords={"epoch": epoch_time[:1]}, + ) + + sci_dependencies = { + "imap_lo_l1b_de": l1b_de, + "imap_lo_l1a_spin": l1a_spin, + "imap_lo_l1b_nhk": l1b_nhk, + } + + result = calculate_de_rates(sci_dependencies, anc_dependencies, attr_mgr_l1b) + + assert result.attrs["Logical_source"] == "imap_lo_l1b_derates" + assert "epoch" in result.coords + assert "esa_step" in result.coords + assert "spin_bin" in result.coords + + # Check that all expected data variables are present + expected_vars = [ + "h_counts", + "o_counts", + "triple_counts", + "double_counts", + "h_rates", + "o_rates", + "triple_rates", + "double_rates", + "exposure_time", + "spin_cycle", + "esa_mode", + ] + for var in expected_vars: + assert var in result.data_vars + + # Check shapes + assert result["h_counts"].shape == (2, 7, 60) # (num_asc, num_esa_steps, num_bins) + assert result["o_counts"].shape == (2, 7, 60) + assert result["exposure_time"].shape == (2, 7) + + # Verify some counts are correct based on our test data + # First ASC (spin_cycle 0) has 3 events at esa_step 1, 2, 3 + # Second ASC (spin_cycle 28) has 2 events at esa_step 1, 2 + # H species: indices 0, 2, 3 + # ASC 0 has 2 H (esa_step 1, 3), ASC 1 has 1 H (esa_step 1) + # O species: indices 1, 4 + # ASC 0 has 1 O (esa_step 2), ASC 1 has 1 O (esa_step 2) + # First ASC, esa_step 1, spin_bin 0 + assert result["h_counts"][0, 0, 0] == 1 + # First ASC, esa_step 3, spin_bin 4 (240//60) + assert result["h_counts"][0, 2, 4] == 1 + # First ASC, esa_step 2, spin_bin 2 (120//60) + assert result["o_counts"][0, 1, 2] == 1 + + # Check that pivot angle was set + assert result["pivot_angle"].values[0] == 45.0 + + # Test that lo_l1b() with descriptor="derates" produces the correct output + output_datasets = lo_l1b(sci_dependencies, anc_dependencies, descriptor="derates") + assert len(output_datasets) == 1 + assert output_datasets[0].attrs["Logical_source"] == "imap_lo_l1b_derates"