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
32 changes: 24 additions & 8 deletions imap_processing/hi/hi_goodtimes.py
Original file line number Diff line number Diff line change
Expand Up @@ -1433,17 +1433,26 @@ def mark_bad_tdc_cal(
logger.info(f"Dropped {n_times_removed} time(s) due to bad TDC calibration")


def _get_sweep_indices(esa_step: np.ndarray) -> np.ndarray:
def _get_sweep_indices(esa_step: np.ndarray, esa_step_met: np.ndarray) -> np.ndarray:
"""
Assign sweep indices to each MET based on ESA step transitions.

A new sweep starts when ESA step transitions from high to low
(e.g., 9 -> 1), detected using np.diff().
Consecutive entries sharing the same esa_step_met belong to the same
8-spin set (e.g., the multiple DE packets of one set). A new sweep starts
at an 8-spin set whose ESA step is less than or equal to the previous
set's ESA step (e.g., 9 -> 1, or 1 -> 1). Treating a repeated ESA step as
a sweep boundary guarantees that each ESA step appears at most once per
sweep, which downstream per-(sweep, ESA step) logic relies on. Repeats
occur, for example, when a sweep is cut short after its first 8-spin set
by a drop out of HVSCI mode, and the next sweep restarts at ESA step 1.
Comment thread
lacoak21 marked this conversation as resolved.

Parameters
----------
esa_step : numpy.ndarray
ESA step values for each MET (epoch dimension).
esa_step_met : numpy.ndarray
MET at which the ESA was stepped for each entry. Used to identify
which entries belong to the same 8-spin set.

Returns
-------
Expand All @@ -1453,10 +1462,13 @@ def _get_sweep_indices(esa_step: np.ndarray) -> np.ndarray:
if len(esa_step) == 0:
return np.array([], dtype=np.int32)

# Find sweep boundaries where ESA step transitions from high to low
# Only compare ESA steps across 8-spin set transitions, not between
# packets within the same 8-spin set.
new_set = np.diff(esa_step_met) != 0
esa_diff = np.diff(esa_step.astype(np.int32))
# Negative diff indicates high-to-low transition (e.g., 9 -> 1 = -8)
sweep_boundaries = esa_diff < 0
# A non-increasing ESA step at a new 8-spin set starts a new sweep
# (e.g., 9 -> 1 = -8, or a repeated step 1 -> 1 = 0)
sweep_boundaries = new_set & (esa_diff <= 0)

# Create sweep indices using cumsum on boundaries
# Prepend False so first MET is in sweep 0
Expand All @@ -1474,15 +1486,19 @@ def _add_sweep_indices(l1b_de: xr.Dataset) -> xr.Dataset:
Parameters
----------
l1b_de : xarray.Dataset
L1B Direct Event dataset or goodtimes dataset.
L1B Direct Event dataset (with "esa_step_met") or goodtimes dataset
(with "met").

Returns
-------
xarray.Dataset
Dataset with esa_sweep coordinate added on the time dimension
(either 'epoch' or 'met').
"""
sweep_indices = _get_sweep_indices(l1b_de["esa_step"].values)
met_name = "esa_step_met" if "esa_step_met" in l1b_de else "met"
sweep_indices = _get_sweep_indices(
l1b_de["esa_step"].values, l1b_de[met_name].values
)
# Determine which dimension to use (epoch for CDF data, met for in-memory)
time_dim = "epoch" if "epoch" in l1b_de.dims else "met"
return l1b_de.assign_coords(esa_sweep=(time_dim, sweep_indices))
Expand Down
109 changes: 94 additions & 15 deletions imap_processing/tests/hi/test_hi_goodtimes.py
Original file line number Diff line number Diff line change
Expand Up @@ -2024,35 +2024,48 @@ def test_final_event_is_last_in_list(self, mock_goodtimes, mock_config_df):
class TestGetSweepIndices:
"""Test suite for _get_sweep_indices() helper function."""

@staticmethod
def _one_set_per_entry(esa_step):
"""Return esa_step_met values giving each entry its own 8-spin set."""
return 1000.0 + 120.0 * np.arange(len(esa_step))

def test_empty_array(self):
"""Test with empty input."""
result = _get_sweep_indices(np.array([]))
result = _get_sweep_indices(np.array([]), np.array([]))
assert len(result) == 0
assert result.dtype == np.int32

def test_single_sweep(self):
"""Test with single complete ESA sweep (no transitions)."""
esa_step = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9])
result = _get_sweep_indices(esa_step)
result = _get_sweep_indices(esa_step, self._one_set_per_entry(esa_step))

# All should be in sweep 0
np.testing.assert_array_equal(result, np.zeros(9, dtype=np.int32))

def test_two_sweeps_standard_transition(self):
"""Test with two sweeps with standard 9->1 transition."""
esa_step = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 6, 7, 8, 9])
result = _get_sweep_indices(esa_step)
result = _get_sweep_indices(esa_step, self._one_set_per_entry(esa_step))

# First 9 should be sweep 0, next 9 should be sweep 1
expected = np.array(
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype=np.int32
)
np.testing.assert_array_equal(result, expected)

def test_ten_step_sweeps(self):
"""Test with 10-step sweeps (9 science steps + ESA step 10)."""
esa_step = np.tile(np.arange(1, 11), 2)
result = _get_sweep_indices(esa_step, self._one_set_per_entry(esa_step))

expected = np.repeat([0, 1], 10).astype(np.int32)
np.testing.assert_array_equal(result, expected)

def test_multiple_sweeps(self):
"""Test with multiple sweeps."""
esa_step = np.array([3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3])
result = _get_sweep_indices(esa_step)
result = _get_sweep_indices(esa_step, self._one_set_per_entry(esa_step))

# Transitions at index 6->7 (9->1) and 15->16 (9->1)
expected = np.array(
Expand All @@ -2063,27 +2076,55 @@ def test_multiple_sweeps(self):
def test_non_standard_transition(self):
"""Test with non-standard ESA step decrease (e.g., 5->2)."""
esa_step = np.array([5, 6, 7, 8, 9, 2, 3, 4, 5])
result = _get_sweep_indices(esa_step)
result = _get_sweep_indices(esa_step, self._one_set_per_entry(esa_step))

# Transition at index 4->5 (9->2, diff=-7, negative so boundary)
expected = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1], dtype=np.int32)
np.testing.assert_array_equal(result, expected)

def test_no_decreases_only_increases(self):
"""Test with only increasing steps (no sweep boundaries)."""
esa_step = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9])
result = _get_sweep_indices(esa_step)
def test_truncated_sweep_repeated_step(self):
"""Test a sweep cut short after its first 8-spin set (repeated ESA 1).

# All in sweep 0
np.testing.assert_array_equal(result, np.zeros(9, dtype=np.int32))
Mimics repoint 152 gain test cycling, where each HVSCI segment holds one
full sweep plus the first 8-spin set of the next sweep before HV drops
out of HVSCI; the next segment restarts the sweep at ESA step 1.
"""
sweep = list(range(1, 10))
esa_step = np.array([*sweep, 1, *sweep, 1, *sweep])
result = _get_sweep_indices(esa_step, self._one_set_per_entry(esa_step))

# The truncated [1] sweeps are their own sweeps, so no ESA step
# repeats within a sweep.
expected = np.array([0] * 9 + [1] + [2] * 9 + [3] + [4] * 9, dtype=np.int32)
np.testing.assert_array_equal(result, expected)

def test_constant_esa_step(self):
"""Test with constant ESA step (no transitions)."""
"""Test with constant ESA step in separate 8-spin sets."""
esa_step = np.array([5, 5, 5, 5, 5])
result = _get_sweep_indices(esa_step)
result = _get_sweep_indices(esa_step, self._one_set_per_entry(esa_step))

# All in sweep 0
np.testing.assert_array_equal(result, np.zeros(5, dtype=np.int32))
# Each repeat is a new 8-spin set at the same step, so a new sweep
np.testing.assert_array_equal(result, np.arange(5, dtype=np.int32))

def test_multiple_packets_per_set(self):
"""Test that packets within the same 8-spin set share one sweep."""
# Four packets per 8-spin set, two sweeps of 3 steps each
esa_step = np.repeat([1, 2, 3, 1, 2, 3], 4)
esa_step_met = np.repeat(1000.0 + 120.0 * np.arange(6), 4)
result = _get_sweep_indices(esa_step, esa_step_met)

expected = np.repeat([0, 1], 12).astype(np.int32)
np.testing.assert_array_equal(result, expected)

def test_multiple_packets_per_set_truncated_sweep(self):
"""Test a truncated sweep with multiple packets per 8-spin set."""
steps = [1, 2, 3, 1, 1, 2, 3]
esa_step = np.repeat(steps, 4)
esa_step_met = np.repeat(1000.0 + 120.0 * np.arange(len(steps)), 4)
result = _get_sweep_indices(esa_step, esa_step_met)

expected = np.repeat([0, 0, 0, 1, 2, 2, 2], 4).astype(np.int32)
np.testing.assert_array_equal(result, expected)


class TestAddSweepIndices:
Expand All @@ -2094,6 +2135,7 @@ def test_adds_coordinate(self):
ds = xr.Dataset(
{
"ccsds_met": (["epoch"], np.array([1000.0, 1060.0, 1120.0])),
"esa_step_met": (["epoch"], np.array([1000.0, 1060.0, 1120.0])),
"esa_step": (["epoch"], np.array([1, 2, 3], dtype=np.uint8)),
},
coords={"epoch": np.arange(3)},
Expand All @@ -2109,6 +2151,7 @@ def test_coordinate_values(self):
ds = xr.Dataset(
{
"ccsds_met": (["epoch"], np.arange(1000.0, 1000.0 + 18 * 60, 60)),
"esa_step_met": (["epoch"], np.arange(1000.0, 1000.0 + 18 * 60, 60)),
"esa_step": (
["epoch"],
np.tile([1, 2, 3, 4, 5, 6, 7, 8, 9], 2).astype(np.uint8),
Expand All @@ -2130,6 +2173,7 @@ def test_preserves_original_data(self):
ds = xr.Dataset(
{
"ccsds_met": (["epoch"], np.array([1000.0, 1060.0, 1120.0])),
"esa_step_met": (["epoch"], np.array([1000.0, 1060.0, 1120.0])),
"esa_step": (["epoch"], np.array([1, 2, 1], dtype=np.uint8)),
"other_var": (["epoch"], np.array([10, 20, 30])),
},
Expand Down Expand Up @@ -2186,6 +2230,11 @@ def _create_test_dataset(
ds = xr.Dataset(
{
"ccsds_met": (["epoch"], ccsds_met),
# Packets of the same 8-spin set share the same esa_step_met
"esa_step_met": (
["epoch"],
np.repeat(ccsds_met[::packets_per_esa_step], packets_per_esa_step),
),
"esa_step": (["epoch"], esa_step),
"esa_energy_step": (["epoch"], esa_energy_step),
"tof_ab": (["event_met"], tof_ab),
Expand Down Expand Up @@ -2284,6 +2333,7 @@ def test_raises_without_esa_sweep_coordinate(self):
ds = xr.Dataset(
{
"ccsds_met": (["epoch"], np.array([1000.0, 1060.0])),
"esa_step_met": (["epoch"], np.array([1000.0, 1060.0])),
"esa_step": (["epoch"], np.array([1, 2], dtype=np.uint8)),
},
coords={"epoch": np.arange(2)},
Expand Down Expand Up @@ -2378,6 +2428,7 @@ def _create_l1b_de_dataset(
"coincidence_type": (["event_met"], coincidence_type),
"ccsds_index": (["event_met"], ccsds_index),
"ccsds_met": (["epoch"], ccsds_met),
"esa_step_met": (["epoch"], ccsds_met),
"esa_step": (["epoch"], esa_step, {"FILLVAL": 255}),
"esa_energy_step": (["epoch"], esa_energy_step, {"FILLVAL": 255}),
},
Expand Down Expand Up @@ -2481,6 +2532,7 @@ def test_partial_sweep_culling(self, goodtimes_for_filter):
"coincidence_type": (["event_met"], coincidence_type),
"ccsds_index": (["event_met"], ccsds_index),
"ccsds_met": (["epoch"], ccsds_met),
"esa_step_met": (["epoch"], ccsds_met),
"esa_step": (["epoch"], esa_step, {"FILLVAL": 255}),
"esa_energy_step": (["epoch"], esa_energy_step, {"FILLVAL": 255}),
},
Expand Down Expand Up @@ -2814,6 +2866,13 @@ def _create_test_dataset(
["epoch"],
np.arange(1000.0, 1000.0 + n_packets * 60, 60),
),
# Two packets per 8-spin set share the same esa_step_met
"esa_step_met": (
["epoch"],
np.repeat(np.arange(1000.0, 1000.0 + n_packets * 60, 120), 2)[
:n_packets
],
),
"esa_step": (["epoch"], esa_step),
"esa_energy_step": (["epoch"], esa_energy_step),
},
Expand Down Expand Up @@ -2857,6 +2916,7 @@ def test_raises_without_coordinate(self):
"coincidence_type": (["event_met"], np.array([12, 4], dtype=np.uint8)),
"ccsds_index": (["event_met"], np.array([0, 0], dtype=np.uint16)),
"ccsds_met": (["epoch"], np.array([1000.0])),
"esa_step_met": (["epoch"], np.array([1000.0])),
"esa_step": (["epoch"], np.array([1], dtype=np.uint8)),
},
coords={"event_met": np.arange(2), "epoch": np.arange(1)},
Expand Down Expand Up @@ -2898,6 +2958,11 @@ def _create_test_dataset(
["epoch"],
np.arange(base_met, base_met + n_packets * 60, 60),
),
# Two packets per 8-spin set share the same esa_step_met
"esa_step_met": (
["epoch"],
np.repeat(np.arange(base_met, base_met + n_packets * 60, 120), 2),
),
"esa_step": (["epoch"], esa_step),
"esa_energy_step": (["epoch"], esa_energy_step),
},
Expand Down Expand Up @@ -3119,6 +3184,10 @@ def test_basic_calculation(self):
["esa_sweep", "esa_energy_step"],
np.full_like(counts_2d, 1000.0),
),
"esa_step_met": (
["esa_sweep", "esa_energy_step"],
np.full_like(counts_2d, 1000.0),
),
},
coords={
"esa_sweep": np.arange(n_sweeps),
Expand Down Expand Up @@ -3154,6 +3223,10 @@ def test_zero_median_excluded(self):
["esa_sweep", "esa_energy_step"],
np.full_like(counts_2d, 1000.0),
),
"esa_step_met": (
["esa_sweep", "esa_energy_step"],
np.full_like(counts_2d, 1000.0),
),
},
coords={
"esa_sweep": np.arange(n_sweeps),
Expand Down Expand Up @@ -3235,6 +3308,7 @@ def _create_l1b_de_dataset(
"coincidence_type": (["event_met"], coincidence_types),
"ccsds_index": (["event_met"], ccsds_index),
"ccsds_met": (["epoch"], ccsds_met),
"esa_step_met": (["epoch"], ccsds_met),
"esa_step": (["epoch"], esa_step),
"esa_energy_step": (["epoch"], esa_energy_step),
},
Expand Down Expand Up @@ -3310,6 +3384,7 @@ def test_fails_extreme_outlier(self, goodtimes_for_filter1):
"coincidence_type": (["event_met"], new_coincidence),
"ccsds_index": (["event_met"], new_ccsds_index),
"ccsds_met": current_ds["ccsds_met"],
"esa_step_met": current_ds["esa_step_met"],
"esa_step": current_ds["esa_step"],
"esa_energy_step": current_ds["esa_energy_step"],
},
Expand Down Expand Up @@ -3417,6 +3492,7 @@ def test_prefilter_marks_all_specified_esas(self, goodtimes_for_filter1):
"coincidence_type": (["event_met"], new_coincidence),
"ccsds_index": (["event_met"], new_ccsds_index),
"ccsds_met": current_ds["ccsds_met"],
"esa_step_met": current_ds["esa_step_met"],
"esa_step": current_ds["esa_step"],
"esa_energy_step": current_ds["esa_energy_step"],
},
Expand Down Expand Up @@ -3766,6 +3842,7 @@ def _create_l1b_de_for_filter2(
{
# Packet-level variables (epoch dimension)
"ccsds_met": (["epoch"], packet_mets),
"esa_step_met": (["epoch"], packet_mets),
"esa_step": (["epoch"], packet_esa_steps),
# Event-level variables (event dimension)
"ccsds_index": (["event_met"], ccsds_index),
Expand Down Expand Up @@ -3835,6 +3912,7 @@ def test_no_clusters(self, goodtimes_for_filter2):
l1b_de = xr.Dataset(
{
"ccsds_met": (["epoch"], packet_mets),
"esa_step_met": (["epoch"], packet_mets),
"esa_step": (["epoch"], packet_esa_steps),
"ccsds_index": (["event_met"], ccsds_index),
"coincidence_type": (["event_met"], coincidence_type),
Expand Down Expand Up @@ -4104,6 +4182,7 @@ def test_only_qualified_events_contribute_to_clusters(self, goodtimes_for_filter
"coincidence_type": (["event_met"], coincidence_type),
"nominal_bin": (["event_met"], nominal_bin),
"ccsds_met": (["epoch"], np.array([1000.0])),
"esa_step_met": (["epoch"], np.array([1000.0])),
"esa_step": (["epoch"], np.array([1], dtype=np.uint8)),
},
coords={
Expand Down
Loading