ENH: Add lo-de-rates data product (IMAP-Science-Operations-Center#2566)

Author: greglucas

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"
             )

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.

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"