Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions imap_processing/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
Expand Down
192 changes: 192 additions & 0 deletions imap_processing/lo/l1b/lo_l1b.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would be helpful to have a comment explaining this equation.


# 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)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice!

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()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great use of xarray functionality.



def _get_esa_level_indices(epochs: np.ndarray, anc_dependencies: list) -> np.ndarray:
"""
Get the ESA level indices (reswept indices) for the given epochs.
Expand Down
121 changes: 121 additions & 0 deletions imap_processing/tests/lo/test_lo_l1b.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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"