From e0916bc50139f19e4fa22b79a0b04d76be82e124 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 26 Oct 2026 13:04:47 -0600 Subject: [PATCH 1/2] Skip IDEX L2A fits for saturated waveforms --- imap_processing/idex/idex_l2a.py | 29 +++- imap_processing/tests/idex/test_idex_l2a.py | 145 ++++++++++++++++++++ 2 files changed, 171 insertions(+), 3 deletions(-) diff --git a/imap_processing/idex/idex_l2a.py b/imap_processing/idex/idex_l2a.py index 0c3eff6b7a..803a2b6330 100644 --- a/imap_processing/idex/idex_l2a.py +++ b/imap_processing/idex/idex_l2a.py @@ -91,6 +91,9 @@ def idex_l2a(l1b_dataset: xr.Dataset, ancillary_files: dict) -> xr.Dataset: The TOF peaks are fitted to EMG curves to determine total intensity, max amplitude, and signal quality. + Saturated waveforms are not fitted. Their fit products contain NaN values; + unsaturated channels in the same event are processed independently. + Parameters ---------- l1b_dataset : xarray.Dataset @@ -133,6 +136,9 @@ def idex_l2a(l1b_dataset: xr.Dataset, ancillary_files: dict) -> xr.Dataset: # Peaks_2d is a list of variable-length arrays peaks_2d = [find_peaks(tof, prominence=0.01)[0] for tof in tof_high] kappa = calculate_kappa(mass_scales, peaks_2d) + # Empty peak lists keep saturated TOF waveforms out of the peak fitter. + for event in np.flatnonzero(l1b_dataset["tof_high_saturation_flag"].data == 1): + peaks_2d[event] = np.array([], dtype=int) # Analyze peaks for estimating dust composition peak_fits_params, area_under_fits, fit_chisqr, fit_redchi = xr.apply_ufunc( @@ -187,7 +193,7 @@ def idex_l2a(l1b_dataset: xr.Dataset, ancillary_files: dict) -> xr.Dataset: ) # Preserve event classification and saturation flags for the event-level L2A - # products. Fits remain available for every event, including non-dust events. + # products. Unsaturated waveforms are fitted even for non-dust events. for flag_name in ALL_FLAG_NAMES: if flag_name in l1b_dataset: l2a_dataset[flag_name] = l1b_dataset[flag_name].copy() @@ -199,9 +205,15 @@ def idex_l2a(l1b_dataset: xr.Dataset, ancillary_files: dict) -> xr.Dataset: estimate_dust_mass, ls_time, l1b_dataset[waveform], + False, + waveform, + l1b_dataset[f"{waveform.lower()}_saturation_flag"], input_core_dims=[ ["time_low_sample_rate_index"], ["time_low_sample_rate_index"], + [], + [], + [], ], output_core_dims=[ ["target_fit_parameter_index"], @@ -213,7 +225,6 @@ def idex_l2a(l1b_dataset: xr.Dataset, ancillary_files: dict) -> xr.Dataset: vectorize=True, output_dtypes=[np.float64] * 6, keep_attrs=True, - kwargs={"waveform_name": waveform}, ) # Calculate mass and velocity estimates if waveform == "Ion_Grid": @@ -393,7 +404,7 @@ def _mask_non_science_derived_estimates(l2a_dataset: xr.Dataset) -> None: def _mask_saturated_derived_estimates(l2a_dataset: xr.Dataset) -> None: """Mask fitted charges and derived estimates for saturated waveforms. - Fit parameters remain available for diagnostics. Impact charge, velocity, + Saturated waveforms are skipped during fitting. Impact charge, velocity, and mass estimates are not scientifically valid when their source waveform is saturated. @@ -1003,6 +1014,7 @@ def estimate_dust_mass( target_signal: xr.DataArray, remove_noise: bool = False, waveform_name: str = "", + saturation_flag: int = 0, ) -> tuple[NDArray, float, float, float, NDArray]: """ Filter and fit the target or ion grid signals to get the total dust impact charge. @@ -1018,6 +1030,9 @@ def estimate_dust_mass( signal. waveform_name : str Channel name used to select channel-specific fit bounds. + saturation_flag : int + Waveform saturation flag from L1b. A value of 1 skips fitting and returns + NaN fit parameters, charge, chi-square values, and fitted samples. Returns ------- @@ -1035,6 +1050,14 @@ def estimate_dust_mass( result : numpy.ndarray The model values evaluated at each time point. """ + if saturation_flag == 1: + return ( + np.full(5, np.nan), + np.nan, + np.nan, + np.nan, + np.full(target_signal.shape, np.nan), + ) signal = np.array(target_signal.data) time = np.array(low_sampling_time.data) # window_start = float(np.min(time)) diff --git a/imap_processing/tests/idex/test_idex_l2a.py b/imap_processing/tests/idex/test_idex_l2a.py index dba158713c..c187b3aa66 100644 --- a/imap_processing/tests/idex/test_idex_l2a.py +++ b/imap_processing/tests/idex/test_idex_l2a.py @@ -7,6 +7,7 @@ import pandas as pd import pytest import xarray as xr +from scipy.optimize import curve_fit from scipy.stats import exponnorm from imap_processing.cdf.utils import load_cdf, write_cdf @@ -26,6 +27,7 @@ chi_square, estimate_dust_mass, fit_impact, + idex_l2a, invert_rise_time_to_velocity, load_calibration_files, log_smooth_powerlaw, @@ -35,6 +37,136 @@ ) +@pytest.mark.parametrize("waveform", ["Target_Low", "Target_High", "Ion_Grid"]) +def test_estimate_dust_mass_skips_saturated_waveform(waveform): + time = xr.DataArray(np.arange(8, dtype=float)) + signal = xr.DataArray(np.full(8, 4095), name=waveform) + + with mock.patch("imap_processing.idex.idex_l2a.curve_fit") as optimizer: + parameters, charge, chisqr, redchi, fitted = estimate_dust_mass( + time, signal, waveform_name=waveform, saturation_flag=1 + ) + + optimizer.assert_not_called() + assert parameters.shape == (5,) + assert fitted.shape == signal.shape + for value in (parameters, charge, chisqr, redchi, fitted): + assert np.all(np.isnan(value)) + + +@pytest.mark.parametrize("saturation", ["mixed", "all", "none"]) +def test_l2a_fits_only_unsaturated_waveforms(saturation): + """Saturation skips the optimizer per channel without removing any events.""" + n_events = 6 + high_dim = "time_high_sample_rate_index" + low_dim = "time_low_sample_rate_index" + high_time = np.linspace(-8.0, 8.0, 33) + low_time = np.linspace(-10.0, 40.0, 64) + tof = exponnorm.pdf(high_time, 1.0, 0.0, 0.5) + 0.001 * np.cos(np.arange(33)) + low_signal = fit_impact(low_time, 0.0, 1.0, 10.0, 0.371, 37.1) + dataset = xr.Dataset( + { + "time_high_sample_rate": ( + ("epoch", high_dim), + np.tile(high_time, (n_events, 1)), + ), + "time_low_sample_rate": ( + ("epoch", low_dim), + np.tile(low_time, (n_events, 1)), + ), + }, + coords={"epoch": np.arange(n_events, dtype=np.int64)}, + attrs={"Logical_source": "imap_idex_l1b_sci-10days"}, + ) + for flag in ALL_FLAG_NAMES: + dataset[flag] = xr.DataArray(np.zeros(n_events, dtype=np.uint8), dims="epoch") + dataset["science_event_flag"][:] = 1 + flags = { + "TOF_High": [0, 0, 0, 0, 1, 1], + "Target_Low": [0, 0, 1, 0, 1, 0], + "Target_High": [0, 1, 0, 0, 1, 0], + "Ion_Grid": [0, 0, 0, 1, 1, 0], + } + for channel, mixed_flags in flags.items(): + values = ( + mixed_flags + if saturation == "mixed" + else np.full(n_events, int(saturation == "all")) + ) + dataset[f"{channel.lower()}_saturation_flag"][:] = values + is_tof = channel == "TOF_High" + samples = np.tile(tof if is_tof else low_signal, (n_events, 1)) + # Any optimizer call receiving these flagged records must fail this test. + samples[np.asarray(values) == 1, samples.shape[1] // 2] = 999.0 + dataset[channel] = xr.DataArray( + samples, dims=("epoch", high_dim if is_tof else low_dim) + ) + original = dataset.copy(deep=True) + fit_calls = [] + + def fit_unsaturated(function, time, signal, *args, **kwargs): + assert not np.any(signal == 999.0), "Saturated samples reached the optimizer" + fit_calls.append(function) + return curve_fit(function, time, signal, *args, **kwargs) + + with ( + mock.patch( + "imap_processing.idex.idex_l2a.load_calibration_files", + return_value=( + np.array([1.27, -0.2, -2.1, -0.37, 5.3, 13.3, 13.3, 0.28]), + np.array([0.06, 2.8, 5.9, 4.1, 13.0, 22.7, 8.2, 0.40]), + ), + ), + mock.patch( + "imap_processing.idex.idex_l2a.time_to_mass", + return_value=( + np.ones(n_events), + np.zeros(n_events), + np.ones_like(dataset.TOF_High), + ), + ), + mock.patch( + "imap_processing.idex.idex_l2a.curve_fit", side_effect=fit_unsaturated + ), + ): + result = idex_l2a(dataset, {}) + + expected_low_fits = 0 + for channel in ("Target_Low", "Target_High", "Ion_Grid"): + prefix = channel.lower() + saturated = dataset[f"{prefix}_saturation_flag"].values == 1 + expected_low_fits += np.count_nonzero(~saturated) + for suffix in ( + "fit_parameters", + "fit_results", + "impact_charge", + "chi_squared", + "reduced_chi_squared", + "velocity_estimate", + "dust_mass_estimate", + ): + assert np.all(np.isnan(result[f"{prefix}_{suffix}"].values[saturated])) + if np.any(~saturated): + np.testing.assert_allclose( + result[f"{prefix}_fit_parameters"].values[~saturated], + np.tile( + [0.0, 1.0, 10.0, 0.371, 37.1], (np.count_nonzero(~saturated), 1) + ), + atol=1e-8, + ) + np.testing.assert_allclose( + result[f"{prefix}_fit_results"].values[~saturated], + np.tile(low_signal, (np.count_nonzero(~saturated), 1)), + atol=1e-8, + ) + assert fit_calls.count(fit_impact) == expected_low_fits + assert fit_calls.count(exponnorm.pdf) == np.count_nonzero( + dataset.tof_high_saturation_flag.values == 0 + ) + np.testing.assert_array_equal(result.epoch, dataset.epoch) + xr.testing.assert_identical(dataset, original) + + def test_non_science_derived_estimates_are_nan() -> None: """Only Science Events retain derived velocity and mass estimates.""" estimates = { @@ -153,6 +285,19 @@ def test_l2a_logical_source_and_cdf(l2a_dataset: xr.Dataset, l1b_dataset: xr.Dat l2a_dataset[flag_name].values, l1b_dataset[flag_name].values ) + for channel in ("target_low", "target_high", "ion_grid"): + saturated = l1b_dataset[f"{channel}_saturation_flag"].values == 1 + for suffix in ( + "fit_parameters", + "fit_results", + "impact_charge", + "chi_squared", + "reduced_chi_squared", + ): + name = f"{channel}_{suffix}" + assert np.all(np.isnan(l2a_dataset[name].values[saturated])) + assert np.all(np.isnan(ds[name].values[saturated])) + with cdflib.CDF(file_name) as cdf_file: assert cdf_file.varattsget("mass_index")["CATDESC"] == ( "Index variable for mass-resolved TOF peak-fit products." From 7f55c5e94e605727183c2526e446a5dc1734284d Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 9 Nov 2026 17:12:51 -0700 Subject: [PATCH 2/2] Update IDEX saturation tests for required L1b statistics --- imap_processing/tests/idex/test_idex_l2a.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/imap_processing/tests/idex/test_idex_l2a.py b/imap_processing/tests/idex/test_idex_l2a.py index 24eac72c9d..38335fffba 100644 --- a/imap_processing/tests/idex/test_idex_l2a.py +++ b/imap_processing/tests/idex/test_idex_l2a.py @@ -58,7 +58,9 @@ def test_l2a_uses_required_l1b_statistics(l1b_dataset, missing_statistic): if missing_statistic is not None: dataset = dataset.drop_vars(missing_statistic) - def fit_waveform(time, signal, remove_noise, waveform_name, baseline_mean): + def fit_waveform( + time, signal, remove_noise, waveform_name, saturation_flag, baseline_mean + ): return ( np.array([0.0, baseline_mean, 1.0, 0.371, 37.1]), 1.0, @@ -104,7 +106,7 @@ def fit_waveform(time, signal, remove_noise, waveform_name, baseline_mean): assert fit.call_count == 6 for channel, expected in expected_baselines.items(): supplied = [ - call.args[4] for call in fit.call_args_list if call.args[3] == channel + call.args[5] for call in fit.call_args_list if call.args[3] == channel ] np.testing.assert_allclose(supplied, expected) np.testing.assert_allclose( @@ -176,6 +178,11 @@ def test_l2a_fits_only_unsaturated_waveforms(saturation): dataset[channel] = xr.DataArray( samples, dims=("epoch", high_dim if is_tof else low_dim) ) + baseline, rms = calculate_baseline_and_rms( + samples, high_time if is_tof else low_time + ) + dataset[f"{channel.lower()}_baseline"] = xr.DataArray(baseline, dims="epoch") + dataset[f"{channel.lower()}_rms"] = xr.DataArray(rms, dims="epoch") original = dataset.copy(deep=True) fit_calls = []