From e25b4cd3ff2f87119c22a8456085faf9254dea28 Mon Sep 17 00:00:00 2001 From: Kyle Severson <12833749+ksseverson57@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:31:27 -0700 Subject: [PATCH 1/3] Improve SNR robustness to artifacts Updated the SNR envelope RR method to enhance noise estimation and tonic range calculation. Added robust methods for tonic signal estimation, improving the overall SNR estimation accuracy, especially in edge cases where large artifacts occur. --- .../metrics/snr_envelope_rr.py | 685 ++++++++++++++++-- 1 file changed, 623 insertions(+), 62 deletions(-) diff --git a/src/aind_dynamic_foraging_basic_analysis/metrics/snr_envelope_rr.py b/src/aind_dynamic_foraging_basic_analysis/metrics/snr_envelope_rr.py index 6646b33..8b31fb0 100644 --- a/src/aind_dynamic_foraging_basic_analysis/metrics/snr_envelope_rr.py +++ b/src/aind_dynamic_foraging_basic_analysis/metrics/snr_envelope_rr.py @@ -37,27 +37,43 @@ - Default ``fps`` is 20 Hz; NaNs are filled with the trace median. - Needs only numpy/scipy -- no other dependencies for any ``noise_method``. -Example +Example: random noise with sinusoidal tonic baseline (SNR ~ 10) +and no real phasic transients (SNR ~ 0) ------- + >>> import numpy as np >>> from snr_envelope_rr import EnvelopeRRSNR ->>> rng = np.random.default_rng(0) +>>> rng = np.random.default_rng(2) >>> t = np.arange(1200) / 20.0 ->>> y = 0.05 * np.sin(2 * np.pi * t / 40.0) + 0.01 * rng.standard_normal(1200) ->>> result = EnvelopeRRSNR(fps=20.0).fit(y) ->>> isinstance(result.snr, float) and isinstance(result.noise, float) +>>> noise_floor = 0.01 +>>> tonic_amp = 5 * noise_floor # 10 * noise floor +>>> y = tonic_amp * np.sin(2 * np.pi * t / 40.0) + noise_floor * rng.standard_normal(1200) +>>> estimator = EnvelopeRRSNR(fps=20.0) +>>> result = estimator.fit(y) +>>> snr, noise, peaks = estimator.estimate(y) # one-shot form +>>> print(f"total SNR: {snr}") +total SNR: 12.965802243331636 +>>> print(f"noise estimate (true=0.01): {noise}") +noise estimate (true=0.01): 0.009333765351685406 +>>> print(f"detected peaks: {peaks}") +detected peaks: [ 56 105 291 348 604 724] +>>> print(f"tonic SNR (true=10.0): {result.snr_tonic:.2f}") +tonic SNR (true=10.0): 10.17 +>>> print(f"phasic SNR (true=0.0): {result.snr_phasic:.2f}") +phasic SNR (true=0.0): 2.80 True + """ from __future__ import annotations import warnings -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Dict, Optional, Tuple, Union import numpy as np from numpy.typing import NDArray -from scipy import interpolate, ndimage, signal, stats +from scipy import interpolate, ndimage, signal __all__ = ["EnvelopeRRSNR", "EnvelopeRRResult"] @@ -88,6 +104,12 @@ # fit_bias_correction_from_benchmark for a different configuration. _TUNED_BIAS_CORRECTION: Tuple[float, float] = (1.8494, -0.5508) +# Only used if tonic_range_method='robust' -- samples at the reference +# fps above (~5.0s); not tuned/validated like _TUNED_DEFAULTS, just a +# reasonable starting point (see _robust_tonic_range's docstring). Scaled +# to the instance's actual fps in __init__, same as _TUNED_DEFAULTS. +_DEFAULT_TONIC_ROBUST_MIN_DISTANCE = 100 + class _UnsetType: """Sentinel distinguishing "bias_correction not specified" (auto- @@ -211,6 +233,14 @@ def _mad_noise_std(residual: NDArray[np.floating], filter_length: int = 31) -> f takes what's left, trims positive-peak outliers, then trims any remaining outliers on either side, and returns the scaled MAD of that twice-trimmed remainder. + + Falls back to a less-trimmed robust std (rather than NaN) if a + trimming step empties out completely -- this happens on short or + heavily-anomalous windows (e.g. one chunk of a chunked/windowed + estimate landing squarely on a large artifact) where the first + trim's own scale estimate collapses to 0. Silently returning NaN + there would poison any downstream aggregation across windows + (``np.median`` propagates a single NaN to the whole result). """ residual = np.asarray(residual, dtype=np.float64) if np.any(np.isnan(residual)): @@ -221,8 +251,14 @@ def _mad_noise_std(residual: NDArray[np.floating], filter_length: int = 31) -> f return float("nan") filtered_0 = noise[noise < 1.5 * np.abs(noise.min())] rstd = _robust_std(filtered_0) - filtered_1 = filtered_0[np.abs(filtered_0) < 2.5 * rstd] - return _robust_std(filtered_1) + filtered_1 = filtered_0[np.abs(filtered_0) < 2.5 * rstd] if rstd > 0 else np.array([]) + if filtered_1.size > 0: + result = _robust_std(filtered_1) + if result > 0: + return result + if rstd > 0: + return rstd + return _robust_std(noise) def _moving_average( @@ -333,7 +369,24 @@ def _arpls_baseline(y, lam: float = 1e7, n_iter: int = 15, ratio: float = 1e-6): def _detect_peaks_rise_rate(residual, candidates, sigma, rise_window: int = 3): """Reject candidate peaks whose approach isn't fast enough to be a - real transient onset (vs. a slow drift crossing threshold).""" + real transient onset (vs. a slow drift crossing threshold). + + Calibrates its slope threshold from a robust (MAD-based) scale + estimate of the below-median ("noise-like") candidate slopes, not + an ordinary standard deviation (``scipy.stats.norm.fit``, what this + used before). Ordinary std is itself not robust: a single moderate + anomaly elsewhere in the trace can spawn a few large-magnitude + candidate slopes from its own decay tail, and even just one or two + of those landing in the "below-median" bucket can inflate an + ordinary std by 2x+ -- which pushes slope_thresh above genuine + events' own slopes and silently suppresses their detection + (measured: one moderate, localized anomaly cut detected event count + from 73 to 19 out of ~80 genuine events, purely through this + threshold-calibration side effect, though the anomaly itself was + nowhere near most of the suppressed events). The MAD-based estimate + used here left detection completely unchanged (73 to 73) on that + same case. + """ if len(candidates) < 5: return candidates @@ -349,7 +402,7 @@ def _detect_peaks_rise_rate(residual, candidates, sigma, rise_window: int = 3): lower = slopes[slopes <= np.median(slopes)] if len(lower) >= 5: - _, sigma_slope = stats.norm.fit(lower) + sigma_slope = 1.4826 * np.median(np.abs(lower - np.median(lower))) slope_thresh = max(sigma_slope * 3.0, sigma * 0.1) else: slope_thresh = np.percentile(slopes, 50) @@ -357,79 +410,288 @@ def _detect_peaks_rise_rate(residual, candidates, sigma, rise_window: int = 3): return candidates[slopes > slope_thresh] -def _decompose_envelope_rr(x: NDArray[np.floating], config: Dict) -> Dict: - """Tonic/phasic decomposition + rise-rate gated peak detection.""" - cfg = config +def _despike_interpolate( + x: NDArray[np.floating], window: int, k: float +) -> Tuple[NDArray[np.floating], NDArray[np.bool_], float]: + """Flag samples far from a local (rolling-median) baseline and + replace them via linear interpolation from the nearest un-flagged + samples on either side -- run *before* tonic fitting, not as a + replacement for it. + + Why this and not a fixed-ceiling clip: clipping a decaying artifact + to a hard ceiling turns it into a flat plateau, which a smoothness- + penalized fit (arPLS/ALS) can end up tracking *more* readily than + the original decay, since a sustained flat elevation looks more + like genuine slow drift than a decaying transient does -- clipping + can make tonic contamination worse, not better (measured directly: + a naive global clip at 50x a robust noise estimate roughly + tripled downstream tonic contamination in one test case). Removing + the flagged span and interpolating over it avoids introducing that + new, more-trackable shape. - midpoint = _moving_average( - x, - window=cfg["midpoint_window"], - polyorder=cfg["midpoint_polyorder"], + Parameters + ---------- + x : ndarray + window : int + Rolling-median window (samples) for the local baseline and + local MAD estimate. + k : float + Flag samples where ``|x - rolling_median| > k * local_MAD``. + + Returns + ------- + x_clean : ndarray + Copy of ``x`` with flagged samples replaced. + flagged : ndarray of bool + local_noise_est : float + Median of the local MAD estimate across the trace (diagnostic). + """ + rolling_med = ndimage.median_filter(x, size=window, mode="reflect") + dev = x - rolling_med + local_mad = 1.4826 * ndimage.median_filter(np.abs(dev), size=window, mode="reflect") + local_mad = np.maximum(local_mad, 1e-12) + flagged = np.abs(dev) > k * local_mad + + x_clean = x.copy() + if flagged.any() and not flagged.all(): + idx = np.arange(len(x)) + x_clean[flagged] = np.interp(idx[flagged], idx[~flagged], x[~flagged]) + + return x_clean, flagged, float(np.median(local_mad)) + + +def _robust_tonic_range(tonic: NDArray[np.floating], min_distance: int) -> Tuple[float, int]: + """Median peak-to-valley swing amplitude of the tonic curve -- + mirrors how phasic amplitude is estimated (detect individual + events, take the median across them) instead of ``ptp(tonic)``'s + single global max-minus-min. + + Detects local maxima and minima in ``tonic`` + (``scipy.signal.argrelextrema``, ``order=min_distance`` -- the + number of samples on each side a point must exceed to count as a + local extremum), takes the absolute difference between each pair of + temporally-adjacent extrema (each one a single peak-to-valley or + valley-to-peak "swing"), and returns the median swing amplitude. + + A one-off contaminated region typically produces only one or two + anomalous swings among many genuine ones, so the median has a real + breakdown point -- unlike ``ptp(tonic)``, which *is* the + contaminated region's own extreme value whenever that region + happens to contain the global max or min. Unlike a chunked/windowed + estimate (:meth:`EnvelopeRRSNR.fit_chunked`), this doesn't impose an + arbitrary fixed-duration window that can slice a genuine slow cycle + in half and underestimate it by construction -- it finds extrema + wherever the curve's own structure actually puts them, the same + principle ``find_peaks`` already uses for phasic events. + + The trade-off: robustness here scales with how many genuine swings + get detected, same as phasic's median scales with how many genuine + events get detected. A short recording relative to the tonic's own + drift period (few genuine cycles) gives the median little to work + with, same limitation phasic would have with only a handful of true + transients. + + Parameters + ---------- + tonic : ndarray + Fitted tonic curve. + min_distance : int + Minimum spacing (samples) between detected extrema -- should be + well below the tonic's own characteristic drift period, but + well above any remaining fast wiggle in the fitted curve. + + Returns + ------- + tonic_range : float + Median swing amplitude, or ``ptp(tonic)`` if fewer than 2 + extrema are detected (not enough structure for a meaningful + median -- e.g. a very short trace or an almost perfectly flat + tonic). + n_swings : int + Number of swings the median was computed over (0 if it fell + back to ``ptp``). + """ + peaks = signal.argrelextrema(tonic, np.greater, order=min_distance)[0] + valleys = signal.argrelextrema(tonic, np.less, order=min_distance)[0] + extrema_idx = np.sort(np.concatenate([peaks, valleys])) + if len(extrema_idx) < 2: + return float(np.ptp(tonic)), 0 + swings = np.abs(np.diff(tonic[extrema_idx])) + return float(np.median(swings)), len(swings) + + +def _apply_pre_despike( + x: NDArray[np.floating], cfg: Dict +) -> Tuple[NDArray[np.floating], int, float]: + """Optionally despike ``x`` before it's used for tonic fitting. + + Off by default (``cfg['pre_despike_window']`` is ``None``), + preserving prior behavior exactly unless opted in. Scoped to + protecting the tonic fit only -- callers should run residual/peak + detection on the ORIGINAL ``x``, not this function's output, so a + genuine large artifact still surfaces as an inspectable outlier + rather than silently vanishing. Even with despiking on, a large + enough outlier relative to the true tonic's own dynamic range can + still leak through a smoothness-penalized fit; + ``n_extreme_samples``/``frac_extreme_samples`` are diagnostics for + exactly that residual risk, not a guarantee despiking fully + removed it. + + Returns + ------- + x_for_tonic : ndarray + ``x`` unchanged if pre-despiking is off, else the despiked copy. + n_extreme_samples : int + frac_extreme_samples : float + """ + pre_despike_window = cfg.get("pre_despike_window", None) + if pre_despike_window is None: + return x, 0, 0.0 + x_for_tonic, flagged, _local_noise = _despike_interpolate( + x, window=pre_despike_window, k=cfg.get("pre_despike_k", 5.0) ) + n_extreme_samples = int(np.sum(flagged)) + frac_extreme_samples = float(n_extreme_samples / len(x)) + return x_for_tonic, n_extreme_samples, frac_extreme_samples - tonic_method = cfg.get("tonic_method", "envelope") + +def _fit_tonic_curve( + x_for_tonic: NDArray[np.floating], + midpoint: NDArray[np.floating], + tonic_method: str, + cfg: Dict, +) -> Tuple[NDArray[np.floating], NDArray[np.intp]]: + """Fit the tonic (slow baseline) curve with the requested tracker. + + Parameters + ---------- + x_for_tonic : ndarray + Trace to fit (already despiked if pre-despiking was applied). + midpoint : ndarray + Smoothed reference curve; only used by ``tonic_method='envelope'``. + tonic_method : {'als', 'arpls', 'envelope'} + cfg : dict + Method-specific tuning knobs (``als_lam``/``als_p``/``als_n_iter``, + ``arpls_lam``/``arpls_n_iter``/``arpls_ratio``, or + ``lower_smooth_window``/``lower_order``/``lower_min_distance``/ + ``interp_kind``). + + Returns + ------- + tonic : ndarray + tonic_minima : ndarray of int + Valley indices tracked by the envelope method; empty for + ``'als'``/``'arpls'`` (they have no discrete "minima" concept). + + Raises + ------ + ValueError + If ``tonic_method`` isn't one of the three supported values. + """ if tonic_method == "als": tonic = _als_baseline( - x, + x_for_tonic, lam=cfg.get("als_lam", 1e7), p=cfg.get("als_p", 0.01), n_iter=cfg.get("als_n_iter", 10), ) - tonic_minima = np.array([], dtype=int) + return tonic, np.array([], dtype=int) elif tonic_method == "arpls": tonic = _arpls_baseline( - x, + x_for_tonic, lam=cfg.get("arpls_lam", 1e7), n_iter=cfg.get("arpls_n_iter", 15), ratio=cfg.get("arpls_ratio", 1e-6), ) - tonic_minima = np.array([], dtype=int) + return tonic, np.array([], dtype=int) elif tonic_method == "envelope": - tonic, tonic_minima = _lower_envelope( - x, + return _lower_envelope( + x_for_tonic, midpoint, cfg["lower_smooth_window"], cfg["lower_order"], cfg["lower_min_distance"], interp_kind=cfg["interp_kind"], ) - else: - raise ValueError( - f"tonic_method must be 'envelope', 'als', or 'arpls'; got {tonic_method!r}." - ) + raise ValueError(f"tonic_method must be 'envelope', 'als', or 'arpls'; got {tonic_method!r}.") - residual = x - tonic - noise_method = cfg.get("noise_method", _DEFAULT_NOISE_METHOD) - if noise_method in _DEPRECATED_NOISE_METHOD_ALIASES: - noise_method = _DEPRECATED_NOISE_METHOD_ALIASES[noise_method] +def _resolve_noise_method(noise_method: str) -> str: + """Resolve a deprecated ``noise_method`` alias (e.g. ``'mad'``) to + its canonical name (``'aind_mad'``); returns unrecognized names + unchanged so the caller's own validation can reject them.""" + return _DEPRECATED_NOISE_METHOD_ALIASES.get(noise_method, noise_method) + +def _estimate_residual_noise(residual: NDArray[np.floating], noise_method: str) -> float: + """Dispatch to the requested noise-floor estimator on the tonic- + subtracted residual. + + Parameters + ---------- + residual : ndarray + noise_method : {'aind_mad', 'folded_iqr', 'mad_iqr_avg'} + Deprecated aliases (e.g. ``'mad'``) are resolved first via + :func:`_resolve_noise_method`. + + Returns + ------- + float + + Raises + ------ + ValueError + If ``noise_method`` (after alias resolution) isn't one of the + three supported values. + """ + noise_method = _resolve_noise_method(noise_method) if noise_method == "folded_iqr": - sigma_iqr = _folded_iqr_noise_std(residual) + return _folded_iqr_noise_std(residual) elif noise_method == "aind_mad": - sigma_iqr = _mad_noise_std(residual) + return _mad_noise_std(residual) elif noise_method == "mad_iqr_avg": - sigma_iqr = 0.5 * (_folded_iqr_noise_std(residual) + _mad_noise_std(residual)) - else: - raise ValueError( - f"noise_method must be one of {_VALID_NOISE_METHODS}; " - f"got {noise_method!r}." - ) - - if len(tonic_minima) >= 5: - n = len(residual) - mid_idx = ((tonic_minima[:-1] + tonic_minima[1:]) // 2).astype(int) - mid_idx = mid_idx[(mid_idx >= 0) & (mid_idx < n)] - sigma_minima = float(np.std(residual[mid_idx])) if len(mid_idx) >= 5 else sigma_iqr - else: - sigma_minima = sigma_iqr - - sigma_fit = sigma_iqr - thresh = cfg["peak_threshold_sd"] * sigma_fit - + return 0.5 * (_folded_iqr_noise_std(residual) + _mad_noise_std(residual)) + raise ValueError(f"noise_method must be one of {_VALID_NOISE_METHODS}; got {noise_method!r}.") + + +def _estimate_sigma_minima( + residual: NDArray[np.floating], + tonic_minima: NDArray[np.intp], + fallback_sigma: float, +) -> float: + """Noise estimate from the residual at valley-to-valley midpoints -- + only meaningful for ``tonic_method='envelope'``, which tracks an + explicit list of minima. Falls back to ``fallback_sigma`` if there + aren't enough minima (or valid midpoints) for a stable estimate. + """ + if len(tonic_minima) < 5: + return fallback_sigma + n = len(residual) + mid_idx = ((tonic_minima[:-1] + tonic_minima[1:]) // 2).astype(int) + mid_idx = mid_idx[(mid_idx >= 0) & (mid_idx < n)] + if len(mid_idx) < 5: + return fallback_sigma + return float(np.std(residual[mid_idx])) + + +def _detect_phasic_events( + residual: NDArray[np.floating], sigma_fit: float, cfg: Dict +) -> Tuple[NDArray[np.intp], float, NDArray[np.floating]]: + """Detect suprathreshold phasic peaks and gate them by rise rate + (see :func:`_detect_peaks_rise_rate`). + + Returns + ------- + event_maxima : ndarray of int + threshold : float + ``peak_threshold_sd * sigma_fit`` -- the height cutoff used. + peak_amps : ndarray + ``residual[event_maxima]``; empty if no events were detected. + """ + threshold = cfg["peak_threshold_sd"] * sigma_fit raw_peaks, _ = signal.find_peaks( residual, - height=thresh, + height=threshold, distance=cfg.get("upper_min_distance", 3), ) event_maxima = _detect_peaks_rise_rate( @@ -441,13 +703,119 @@ def _decompose_envelope_rr(x: NDArray[np.floating], config: Dict) -> Dict: event_maxima = np.asarray(event_maxima, dtype=int) event_maxima = event_maxima[(event_maxima >= 0) & (event_maxima < len(residual))] peak_amps = residual[event_maxima] if len(event_maxima) > 0 else np.array([]) + return event_maxima, threshold, peak_amps - if len(peak_amps) > 0: - phasic_p95 = float(np.percentile(peak_amps, 95)) - phasic_median = float(np.median(peak_amps)) - phasic_sd = float(np.std(peak_amps)) - else: - phasic_p95 = phasic_median = phasic_sd = 0.0 + +def _summarize_phasic_amplitudes( + peak_amps: NDArray[np.floating], +) -> Tuple[float, float, float]: + """95th percentile / median / std of detected peak amplitudes; all + zero if no events were detected. + + Returns + ------- + (phasic_p95, phasic_median, phasic_sd) + """ + if len(peak_amps) == 0: + return 0.0, 0.0, 0.0 + return ( + float(np.percentile(peak_amps, 95)), + float(np.median(peak_amps)), + float(np.std(peak_amps)), + ) + + +def _compute_tonic_range( + tonic: NDArray[np.floating], tonic_range_method: str, cfg: Dict +) -> Tuple[float, int]: + """Dispatch to the requested tonic-amplitude summary statistic. + + ``ptp(tonic)`` (``max - min``) has a breakdown point of exactly one + sample -- a single large excursion the tonic fit only partially + absorbs (e.g. from a sustained artifact arPLS's reweighting doesn't + fully reject) inflates it directly, with nothing to average it out. + ``'percentile'`` trims ``tonic_range_trim_pct`` from each tail + before taking the range, at the cost of also clipping any genuine + tonic dynamic range that happens to live in that trimmed fraction -- + choose ``tonic_range_trim_pct`` to comfortably exceed the fraction + of the trace you expect a real artifact to occupy (e.g. a 5s glitch + in a 150s trace is ~3.3% one-sided; ``trim_pct=5`` gives a 5% + one-sided margin above that). ``'robust'`` (the default) instead + takes the median peak-to-valley swing amplitude across detected + local extrema in the tonic curve (see :func:`_robust_tonic_range`), + mirroring how phasic amplitude is estimated (detect individual + events, take the median across them) rather than reading off a + single global extreme value. + + Returns + ------- + tonic_range : float + n_tonic_swings : int + Only nonzero for ``tonic_range_method='robust'``. + + Raises + ------ + ValueError + If ``tonic_range_method`` isn't one of the three supported + values. + """ + if tonic_range_method == "ptp": + return float(np.ptp(tonic)), 0 + elif tonic_range_method == "percentile": + trim_pct = cfg.get("tonic_range_trim_pct", 5.0) + tonic_range = float(np.percentile(tonic, 100 - trim_pct) - np.percentile(tonic, trim_pct)) + return tonic_range, 0 + elif tonic_range_method == "robust": + min_distance = cfg.get("tonic_robust_min_distance", _DEFAULT_TONIC_ROBUST_MIN_DISTANCE) + return _robust_tonic_range(tonic, min_distance) + raise ValueError( + f"tonic_range_method must be 'ptp', 'percentile', or 'robust'; " + f"got {tonic_range_method!r}." + ) + + +def _decompose_envelope_rr(x: NDArray[np.floating], config: Dict) -> Dict: + """Tonic/phasic decomposition + rise-rate gated peak detection. + + A thin orchestrator: optional pre-despiking + (:func:`_apply_pre_despike`) -> tonic fit (:func:`_fit_tonic_curve`) + -> noise-floor estimate (:func:`_estimate_residual_noise`, + :func:`_estimate_sigma_minima`) -> phasic peak detection + (:func:`_detect_phasic_events`, :func:`_summarize_phasic_amplitudes`) + -> tonic-amplitude summary statistic (:func:`_compute_tonic_range`). + See each helper's own docstring for the mechanism and trade-offs of + the strategy it dispatches between; this function itself makes no + strategy decisions of its own. + """ + cfg = config + + x_for_tonic, n_extreme_samples, frac_extreme_samples = _apply_pre_despike(x, cfg) + + midpoint = _moving_average( + x_for_tonic, + window=cfg["midpoint_window"], + polyorder=cfg["midpoint_polyorder"], + ) + + tonic_method = cfg.get("tonic_method", "envelope") + tonic, tonic_minima = _fit_tonic_curve(x_for_tonic, midpoint, tonic_method, cfg) + + # residual/peak-detection intentionally use the ORIGINAL x, not + # x_for_tonic -- despiking is scoped to protecting the tonic fit; a + # genuine large artifact should still surface in the residual/phasic + # stats as an inspectable outlier rather than being silently erased. + residual = x - tonic + + noise_method = cfg.get("noise_method", _DEFAULT_NOISE_METHOD) + sigma_iqr = _estimate_residual_noise(residual, noise_method) + sigma_minima = _estimate_sigma_minima(residual, tonic_minima, sigma_iqr) + sigma_fit = sigma_iqr + + event_maxima, thresh, peak_amps = _detect_phasic_events(residual, sigma_fit, cfg) + phasic_p95, phasic_median, phasic_sd = _summarize_phasic_amplitudes(peak_amps) + + tonic_range_method = cfg.get("tonic_range_method", "robust") + tonic_range, n_tonic_swings = _compute_tonic_range(tonic, tonic_range_method, cfg) return { "tonic": tonic, @@ -462,14 +830,17 @@ def _decompose_envelope_rr(x: NDArray[np.floating], config: Dict) -> Dict: "peak_threshold": float(thresh), "detection_method": "rise_rate", "tonic_method": tonic_method, - "tonic_range": float(np.ptp(tonic)), + "tonic_range": tonic_range, "phasic_p95": phasic_p95, "phasic_median": phasic_median, "phasic_amplitude": phasic_sd, "phasic_snr_p95": float(phasic_p95 / sigma_fit) if sigma_fit > 0 else float("nan"), "phasic_snr_median": float(phasic_median / sigma_fit) if sigma_fit > 0 else float("nan"), "phasic_snr_sd": float(phasic_sd / sigma_fit) if sigma_fit > 0 else float("nan"), - "tonic_snr_sd": float(np.ptp(tonic) / sigma_fit) if sigma_fit > 0 else float("nan"), + "tonic_snr_sd": float(tonic_range / sigma_fit) if sigma_fit > 0 else float("nan"), + "n_extreme_samples": n_extreme_samples, + "n_tonic_swings": n_tonic_swings, + "frac_extreme_samples": frac_extreme_samples, "config": cfg, } @@ -519,6 +890,30 @@ class EnvelopeRRResult: ``signal_statistic``). config : dict The resolved configuration used for this fit. + n_extreme_samples : int + Count of samples flagged and interpolated over by pre-despiking + before the tonic fit (0 if ``pre_despike_window`` wasn't set in + ``config``). A nonzero count doesn't mean the tonic fit is now + clean -- see the module docstring's note on ``tonic_range`` and + despiking's actual, partial effectiveness against large, + sustained artifacts. + frac_extreme_samples : float + ``n_extreme_samples / len(trace)``. + n_chunks : int or None + Number of chunks actually used to compute ``snr_tonic``, if + this result came from :meth:`EnvelopeRRSNR.fit_chunked` rather + than :meth:`EnvelopeRRSNR.fit`. ``None`` after a plain ``fit``. + chunk_duration_s : float or None + Chunk length (seconds) used by :meth:`fit_chunked`, if + applicable. + per_chunk_tonic_snr : list of float or None + Each chunk's own ``snr_tonic`` before aggregation, if this + result came from :meth:`fit_chunked`. + n_tonic_swings : int + Number of peak-to-valley swings the median was computed over, + if ``config['tonic_range_method'] == 'robust'`` (0 otherwise, + including the fallback-to-``ptp`` case when fewer than 2 + extrema were detected -- see :func:`_robust_tonic_range`). """ snr: float @@ -532,6 +927,12 @@ class EnvelopeRRResult: config: Dict = field(repr=False) snr_corrected: Optional[float] = None snr_phasic_corrected: Optional[float] = None + n_extreme_samples: int = 0 + frac_extreme_samples: float = 0.0 + n_chunks: Optional[int] = None + chunk_duration_s: Optional[float] = None + per_chunk_tonic_snr: Optional[list] = None + n_tonic_swings: int = 0 class EnvelopeRRSNR: @@ -594,6 +995,23 @@ class EnvelopeRRSNR: ``config={'noise_method': ...}``, which takes precedence over the ``noise_method`` argument. + ``tonic_range_method`` (default ``'robust'``) controls how + ``snr_tonic`` is computed from the fitted tonic curve: + ``'robust'`` -- median peak-to-valley swing amplitude across + detected local extrema, mirroring how ``snr_phasic`` is + estimated (detect individual events, take the median across + them), spacing controlled by ``tonic_robust_min_distance`` + (default 100 samples at 20fps, fps-scaled) -- or ``'ptp'`` + (``max(tonic) - min(tonic)``, the previous default; a fragile, + single-sample-breakdown-point statistic -- kept for backward + compatibility, and still useful as a plain, unadorned baseline + to compare against), or ``'percentile'`` (trims + ``tonic_range_trim_pct`` from each tail before taking the + range). See :func:`_robust_tonic_range` for the full mechanism + and its own trade-off (robustness scales with how many genuine + tonic swings actually get detected, same as ``snr_phasic``'s + robustness scales with detected event count). + Notes ----- - Noise is estimated from the tonic-subtracted residual, not the @@ -663,6 +1081,7 @@ def __init__( _TUNED_DEFAULTS["lower_smooth_window"], fps, make_odd=True ), "rise_window": self.scale_window(_TUNED_DEFAULTS["rise_window"], fps), + "tonic_robust_min_distance": self.scale_window(_DEFAULT_TONIC_ROBUST_MIN_DISTANCE, fps), } self.config = { **_DEFAULT_CONFIG, @@ -761,9 +1180,151 @@ def fit(self, trace: NDArray[np.floating]) -> EnvelopeRRResult: residual=raw["residual"], signal=signal, config=raw["config"], + n_extreme_samples=raw["n_extreme_samples"], + frac_extreme_samples=raw["frac_extreme_samples"], + n_tonic_swings=raw["n_tonic_swings"], ) return self.result_ + def fit_chunked( + self, + trace: NDArray[np.floating], + chunk_duration_s: Optional[float] = None, + min_chunk_duration_s: float = 30.0, + chunk_fraction: float = 0.20, + aggregate: str = "median", + ) -> EnvelopeRRResult: + """Like :meth:`fit`, but computes ``snr_tonic`` from a chunked, + per-window median instead of the single global fit's + ``ptp(tonic)``. + + Why: ``ptp(tonic)`` has a breakdown point of exactly one sample + -- a single large excursion the tonic fit only partially + absorbs (e.g. a sustained artifact arPLS's reweighting doesn't + fully reject) inflates it directly, with nothing to average it + out. Splitting the trace into independent chunks and fitting + each one separately means a one-off artifact can only poison + the (hopefully minority of) chunks it actually overlaps; the + median across chunks then has a real breakdown point, rather + than relying on one global fit's shape being trustworthy in + the first place. Measured on synthetic one-off sustained + artifacts: brings tonic_snr inflation from ~6x (a flat true + baseline, 5000x-noise_sigma artifact) down to ~1.04x, versus + ~15-20% reductions from pre-despiking or percentile-trimmed + ``tonic_range`` alone -- this changes the failure mode instead + of just softening it. + + The trade-off: a chunk shorter than the recording's own genuine + tonic drift period will only see a fraction of that drift, + biasing ``snr_tonic`` down even on a clean trace (measured: + ~3x underestimate using 15s chunks against a 120s-period + drift). Default chunk sizing (``chunk_duration_s=None``) + balances this with ``max(min_chunk_duration_s, chunk_fraction * + recording_duration)`` -- a fixed floor for short recordings, + scaling up for longer ones so chunk size tracks a long + recording's own timescale rather than staying fixed at the + floor. Pass ``chunk_duration_s`` explicitly to match your own + recordings' known drift timescale instead of relying on this + heuristic -- it's a reasonable default, not a substitute for + knowing your own data. + + Everything else (phasic peak detection, noise floor, residual, + the returned ``tonic`` curve) is unchanged from :meth:`fit` -- + only ``snr_tonic`` (and the ``snr``/``snr_corrected`` totals + derived from it) differ. New fields not populated by a plain + :meth:`fit` call: ``n_chunks``, ``chunk_duration_s``, + ``per_chunk_tonic_snr``. + + Parameters + ---------- + trace : ndarray + chunk_duration_s : float, optional + Chunk length in seconds. If None (default), uses + ``max(min_chunk_duration_s, chunk_fraction * len(trace)/fps)``. + min_chunk_duration_s : float + Floor on the default chunk duration (seconds). Ignored if + ``chunk_duration_s`` is given explicitly. + chunk_fraction : float + Fraction of the recording's total duration used for the + default chunk duration, before the floor is applied. + Ignored if ``chunk_duration_s`` is given explicitly. + aggregate : {'median', 'mean'} + How to combine per-chunk ``snr_tonic`` values. ``'median'`` + (default) is what gives this its outlier robustness; + ``'mean'`` has no breakdown point and defeats the purpose -- + provided mainly for comparison/diagnostics. + + Returns + ------- + EnvelopeRRResult + Same structure as :meth:`fit`'s return value, with + ``snr_tonic``/``snr``/``snr_corrected`` computed from the + chunked-median tonic estimate, plus ``n_chunks``, + ``chunk_duration_s``, ``per_chunk_tonic_snr``. Also stored + on ``self.result_``. + """ + trace = np.nan_to_num(trace, nan=float(np.nanmedian(trace))) + + # Global fit first: phasic peaks/signal/noise/residual/the + # returned tonic curve all come from here, unchanged from a + # plain fit() -- only snr_tonic (and totals derived from it) + # get overridden below. + result = self.fit(trace) + + n = len(trace) + if chunk_duration_s is None: + chunk_duration_s = max(min_chunk_duration_s, chunk_fraction * (n / self.fps)) + chunk_len = max(1, int(round(chunk_duration_s * self.fps))) + n_chunks_requested = max(1, n // chunk_len) + chunks = np.array_split(trace, n_chunks_requested) + + min_samples = max(50, int(self.config.get("midpoint_window", 101))) + per_chunk_tonic_snr = [] + for c in chunks: + if len(c) < min_samples: + continue + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r = _decompose_envelope_rr(c, self.config) + per_chunk_tonic_snr.append(float(r["tonic_snr_sd"])) + + if aggregate == "median": + agg_fn = np.nanmedian + elif aggregate == "mean": + agg_fn = np.nanmean + else: + raise ValueError(f"aggregate must be 'median' or 'mean'; got {aggregate!r}.") + + if per_chunk_tonic_snr: + snr_tonic_chunked = float(agg_fn(per_chunk_tonic_snr)) + else: + warnings.warn( + "No chunk was long enough to fit (trace too short for " + "chunk_duration_s/min_chunk_duration_s) -- falling back " + "to the global (non-chunked) snr_tonic.", + RuntimeWarning, + stacklevel=2, + ) + snr_tonic_chunked = result.snr_tonic + + snr_total_chunked = snr_tonic_chunked + result.snr_phasic # NaN-propagates + + snr_corrected_chunked = None + if self.bias_correction is not None and result.snr_phasic_corrected is not None: + snr_corrected_chunked = snr_tonic_chunked + result.snr_phasic_corrected + + chunked_result = replace( + result, + snr=snr_total_chunked, + snr_tonic=snr_tonic_chunked, + snr_corrected=snr_corrected_chunked, + n_chunks=len(per_chunk_tonic_snr), + chunk_duration_s=float(chunk_duration_s), + per_chunk_tonic_snr=per_chunk_tonic_snr, + ) + self.result_ = chunked_result + return chunked_result + def estimate( self, trace: NDArray[np.floating], apply_correction: bool = False ) -> Tuple[float, float, NDArray[np.intp]]: From 19f86600f34f5537ef511067e6b6c4745882a6cf Mon Sep 17 00:00:00 2001 From: Kyle Severson <12833749+ksseverson57@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:25:02 -0700 Subject: [PATCH 2/3] added tests for envelope RR SNR estimation metric --- tests/test_metric_snr_envelope.py | 691 ++++++++++++++++++++++++++++++ 1 file changed, 691 insertions(+) create mode 100644 tests/test_metric_snr_envelope.py diff --git a/tests/test_metric_snr_envelope.py b/tests/test_metric_snr_envelope.py new file mode 100644 index 0000000..5a63e1e --- /dev/null +++ b/tests/test_metric_snr_envelope.py @@ -0,0 +1,691 @@ +"""Test the envelope + rise-rate (Env+RR) SNR estimator. + +To run the test, execute "python -m unittest tests/test_metric_snr_envelope.py". +""" + +import unittest +import warnings +from unittest.mock import patch + +import numpy as np + +from aind_dynamic_foraging_basic_analysis.metrics.snr_envelope_rr import (_UNSET, EnvelopeRRSNR, + _detect_peaks_rise_rate, + _estimate_residual_noise, + _estimate_sigma_minima, + _folded_iqr_noise_std, + _half_sample_mode, + _interpolate_envelope, + _lower_envelope, + _mad_noise_std, + _robust_std, _UnsetType) + +# Shared simulation parameters, reused across every test so intent stays visible at +# each call site (e.g. `event_amplitude=0.2` reads as "20x NOISE_SIGMA") without +# repeating magic numbers. +FPS = 20.0 +N_SAMPLES = 3000 # 150 s at FPS +NOISE_SIGMA = 0.01 + + +def make_tonic_only_trace(seed=0, tonic_amp=0.05, tonic_period_s=40.0): + """Build a trace with a known sinusoidal tonic component and NO phasic events. + + Ground truth is exact: the tonic component is returned separately, so + ``np.ptp(tonic_true) / NOISE_SIGMA`` is the true tonic SNR to compare estimates + against, and the true phasic amplitude is exactly zero. + + Returns + ------- + trace : np.ndarray + tonic_true : np.ndarray + The trace's only signal component (noise aside). + """ + rng = np.random.default_rng(seed) + t = np.arange(N_SAMPLES) / FPS + tonic_true = tonic_amp * np.sin(2 * np.pi * t / tonic_period_s) + noise = NOISE_SIGMA * rng.standard_normal(N_SAMPLES) + return tonic_true + noise, tonic_true + + +def make_phasic_only_trace(seed=1, event_amplitude=0.2, n_events=20, tau_rise=2, tau_decay=8): + """Build a flat-tonic trace with a known number of phasic transients of a known + amplitude, evenly spaced (no overlap), plus Gaussian noise. + + Returns + ------- + trace : np.ndarray + event_indices : np.ndarray + True sample index of each transient's onset. + """ + rng = np.random.default_rng(seed) + trace = NOISE_SIGMA * rng.standard_normal(N_SAMPLES) + + spacing = N_SAMPLES // (n_events + 1) + event_indices = np.arange(1, n_events + 1) * spacing + + kernel_len = 60 # 3 s -- well under `spacing`, so events never overlap + kernel_t = np.arange(kernel_len) + kernel = (1 - np.exp(-kernel_t / tau_rise)) * np.exp(-kernel_t / tau_decay) + kernel = kernel / kernel.max() + + for idx in event_indices: + end = min(N_SAMPLES, idx + kernel_len) + trace[idx:end] += event_amplitude * kernel[: end - idx] + + return trace, event_indices + + +def inject_artifact(trace, amplitude, onset_s, width_s=2.0): + """Add one large, brief artifact to ``trace`` and return the modified trace plus + the sample-index range it occupies. + + Uses a raised-cosine bump: exactly zero outside its window, and both its value + and slope reach zero at each edge, so it has no long tail to distort a tonic fit + with -- a genuinely transient artifact (seconds long, decays rapidly to + baseline), not an unbounded one. + + Returns + ------- + trace_with_artifact : np.ndarray + onset_idx, end_idx : int + Sample-index range the artifact occupies (for checking it gets detected). + """ + n = len(trace) + t = np.arange(n) / FPS + onset_idx = int(round(onset_s * FPS)) + end_idx = onset_idx + int(round(width_s * FPS)) + rel_t = t - onset_s + in_window = (rel_t >= 0) & (rel_t <= width_s) + bump = np.where(in_window, amplitude * 0.5 * (1 - np.cos(2 * np.pi * rel_t / width_s)), 0.0) + return trace + bump, onset_idx, end_idx + + +class TestEnvelopeRRSNRTonic(unittest.TestCase): + """Test snr_tonic / noise estimation on traces with a known sinusoidal tonic and + no phasic events.""" + + def test_tonic_snr_matches_ground_truth(self): + """snr_tonic should track ptp(tonic_true) / NOISE_SIGMA for a clean sinusoid.""" + trace, tonic_true = make_tonic_only_trace() + true_tonic_snr = np.ptp(tonic_true) / NOISE_SIGMA + + result = EnvelopeRRSNR(fps=FPS).fit(trace) + + self.assertAlmostEqual(result.snr_tonic / true_tonic_snr, 1.0, delta=0.3) + + def test_noise_estimate_close_to_true_sigma(self): + """noise should track the true injected NOISE_SIGMA.""" + trace, _ = make_tonic_only_trace() + result = EnvelopeRRSNR(fps=FPS).fit(trace) + self.assertAlmostEqual(result.noise / NOISE_SIGMA, 1.0, delta=0.3) + + def test_no_phasic_events_returns_nan_and_warns(self): + """A strict enough threshold on a no-phasic-signal trace should detect zero + peaks, return NaN for snr_phasic (and the total snr, via propagation), warn, + and leave snr_tonic unaffected.""" + trace, _ = make_tonic_only_trace() + estimator = EnvelopeRRSNR(fps=FPS, peak_threshold_sd=50.0) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + result = estimator.fit(trace) + self.assertTrue(any(issubclass(w.category, RuntimeWarning) for w in caught)) + + self.assertEqual(len(result.peaks), 0) + self.assertTrue(np.isnan(result.snr_phasic)) + self.assertTrue(np.isnan(result.snr)) + self.assertFalse(np.isnan(result.snr_tonic)) + + +class TestEnvelopeRRSNRPhasic(unittest.TestCase): + """Test snr_phasic / peak detection on traces with a flat tonic and known phasic + events.""" + + def test_detects_expected_number_of_events(self): + """Detected event count should be close to the true count (rise-rate gating + or boundary effects can occasionally add or drop one).""" + trace, event_indices = make_phasic_only_trace() + result = EnvelopeRRSNR(fps=FPS).fit(trace) + self.assertAlmostEqual(len(result.peaks), len(event_indices), delta=2) + + def test_detected_peaks_land_near_true_event_indices(self): + """Every true event should have a detected peak within 0.5 s of it.""" + trace, event_indices = make_phasic_only_trace() + result = EnvelopeRRSNR(fps=FPS).fit(trace) + self.assertGreater(len(result.peaks), 0) + for true_idx in event_indices: + nearest_distance = np.min(np.abs(result.peaks - true_idx)) + self.assertLess(nearest_distance, 0.5 * FPS) + + def test_snr_phasic_scales_linearly_with_true_amplitude(self): + """Tripling the true event amplitude should roughly triple snr_phasic.""" + trace_small, _ = make_phasic_only_trace(event_amplitude=0.1) + trace_large, _ = make_phasic_only_trace(event_amplitude=0.3) + + snr_small = EnvelopeRRSNR(fps=FPS).fit(trace_small).snr_phasic + snr_large = EnvelopeRRSNR(fps=FPS).fit(trace_large).snr_phasic + + expected_ratio = 0.3 / 0.1 + actual_ratio = snr_large / snr_small + self.assertAlmostEqual(actual_ratio / expected_ratio, 1.0, delta=0.3) + + +class TestEnvelopeRRSNRArtifact(unittest.TestCase): + """Test robustness of noise/SNR estimates to one large-amplitude artifact, added + on top of a known sinusoidal-tonic trace.""" + + ARTIFACT_AMPLITUDE = 500 * NOISE_SIGMA # "hundreds of SD", per real problem recordings + ARTIFACT_ONSET_S = 40.0 + ARTIFACT_WIDTH_S = 2.0 + + def _clean_and_artifact_results(self, seed): + """Fit the same trace with and without one injected artifact; returns + (result_clean, result_artifact, onset_idx, end_idx).""" + trace_clean, _ = make_tonic_only_trace(seed=seed) + trace_artifact, onset_idx, end_idx = inject_artifact( + trace_clean.copy(), + amplitude=self.ARTIFACT_AMPLITUDE, + onset_s=self.ARTIFACT_ONSET_S, + width_s=self.ARTIFACT_WIDTH_S, + ) + result_clean = EnvelopeRRSNR(fps=FPS).fit(trace_clean) + result_artifact = EnvelopeRRSNR(fps=FPS).fit(trace_artifact) + return result_clean, result_artifact, onset_idx, end_idx + + def test_noise_floor_robust_to_large_artifact(self): + """The noise floor should stay close to its own no-artifact value despite an + artifact ~500x the noise floor being added.""" + result_clean, result_artifact, _, _ = self._clean_and_artifact_results(seed=2) + ratio = result_artifact.noise / result_clean.noise + self.assertAlmostEqual(ratio, 1.0, delta=0.5) + + def test_tonic_snr_robust_to_large_artifact(self): + """snr_tonic (default tonic_range_method='robust') should stay close to its + own no-artifact value despite the same large artifact.""" + result_clean, result_artifact, _, _ = self._clean_and_artifact_results(seed=3) + ratio = result_artifact.snr_tonic / result_clean.snr_tonic + self.assertAlmostEqual(ratio, 1.0, delta=0.5) + + def test_artifact_itself_is_detected_as_a_phasic_peak(self): + """The artifact should surface as an inspectable detected peak in its own + right, not be silently absorbed or ignored.""" + _, result_artifact, onset_idx, end_idx = self._clean_and_artifact_results(seed=4) + in_artifact_window = (result_artifact.peaks >= onset_idx) & ( + result_artifact.peaks <= end_idx + ) + self.assertTrue(np.any(in_artifact_window)) + + +class TestEnvelopeRRSNRConvenienceAPI(unittest.TestCase): + """Test that the one-shot convenience methods agree with an equivalent .fit() + call, since they're thin wrappers around it.""" + + def test_estimate_matches_fit(self): + """.estimate() returns (snr, noise, peaks) matching a .fit() result.""" + trace, _ = make_phasic_only_trace() + estimator = EnvelopeRRSNR(fps=FPS) + + result = estimator.fit(trace) + snr, noise, peaks = estimator.estimate(trace) + + self.assertEqual(snr, result.snr) + self.assertEqual(noise, result.noise) + np.testing.assert_array_equal(peaks, result.peaks) + + def test_estimate_components_matches_fit(self): + """.estimate_components() returns (snr_total, snr_tonic, snr_phasic) + matching a .fit() result.""" + trace, _ = make_phasic_only_trace() + estimator = EnvelopeRRSNR(fps=FPS) + + result = estimator.fit(trace) + snr_total, snr_tonic, snr_phasic = estimator.estimate_components(trace) + + self.assertEqual(snr_total, result.snr) + self.assertEqual(snr_tonic, result.snr_tonic) + self.assertEqual(snr_phasic, result.snr_phasic) + + def test_decompose_matches_fit(self): + """.decompose() returns (tonic, residual) matching a .fit() result.""" + trace, tonic_true = make_tonic_only_trace() + estimator = EnvelopeRRSNR(fps=FPS) + + result = estimator.fit(trace) + tonic, residual = estimator.decompose(trace) + + np.testing.assert_array_equal(tonic, result.tonic) + np.testing.assert_array_equal(residual, result.residual) + np.testing.assert_allclose(residual, trace - tonic) + + def test_decompose_reuses_cached_fit_when_trace_omitted(self): + """.decompose() with no argument should reuse the most recent .fit(), + not silently do nothing or re-fit.""" + trace, _ = make_tonic_only_trace() + estimator = EnvelopeRRSNR(fps=FPS) + result = estimator.fit(trace) + + tonic, residual = estimator.decompose() + + np.testing.assert_array_equal(tonic, result.tonic) + np.testing.assert_array_equal(residual, result.residual) + + +class TestEnvelopeRRSNRConstructorValidation(unittest.TestCase): + """Test constructor argument validation and the not-yet-fitted guard.""" + + def test_invalid_signal_statistic_raises(self): + with self.assertRaises(ValueError): + EnvelopeRRSNR(fps=FPS, signal_statistic="bogus") + + def test_invalid_noise_method_via_argument_raises(self): + with self.assertRaises(ValueError): + EnvelopeRRSNR(fps=FPS, noise_method="bogus") + + def test_invalid_noise_method_via_config_raises(self): + """config['noise_method'] takes precedence over the constructor argument, + so an invalid value there must be validated too.""" + with self.assertRaises(ValueError): + EnvelopeRRSNR(fps=FPS, config={"noise_method": "bogus"}) + + def test_deprecated_mad_alias_resolves_to_aind_mad(self): + estimator = EnvelopeRRSNR(fps=FPS, noise_method="mad") + self.assertEqual(estimator.noise_method, "aind_mad") + + def test_invalid_tonic_method_raises_on_fit(self): + """tonic_method isn't validated until .fit() actually dispatches on it.""" + trace, _ = make_tonic_only_trace() + estimator = EnvelopeRRSNR(fps=FPS, config={"tonic_method": "bogus"}) + with self.assertRaises(ValueError): + estimator.fit(trace) + + def test_invalid_tonic_range_method_raises_on_fit(self): + trace, _ = make_tonic_only_trace() + estimator = EnvelopeRRSNR(fps=FPS, config={"tonic_range_method": "bogus"}) + with self.assertRaises(ValueError): + estimator.fit(trace) + + def test_accessing_result_before_fit_raises(self): + """Every convenience property (and .decompose() with no argument) should + raise a clear RuntimeError before .fit()/.estimate() has ever been called.""" + estimator = EnvelopeRRSNR(fps=FPS) + for accessor in ( + lambda: estimator.snr_, + lambda: estimator.snr_corrected_, + lambda: estimator.snr_tonic_, + lambda: estimator.snr_phasic_, + lambda: estimator.noise_, + lambda: estimator.peaks_, + lambda: estimator.tonic_, + lambda: estimator.residual_, + lambda: estimator.decompose(), + ): + with self.assertRaises(RuntimeError): + accessor() + + def test_convenience_properties_match_fit_result(self): + """Every `*_` property should mirror the corresponding field on the most + recent .fit() result.""" + trace, _ = make_phasic_only_trace() + estimator = EnvelopeRRSNR(fps=FPS) + result = estimator.fit(trace) + + self.assertEqual(estimator.snr_, result.snr) + self.assertEqual(estimator.snr_corrected_, result.snr_corrected) + self.assertEqual(estimator.snr_tonic_, result.snr_tonic) + self.assertEqual(estimator.snr_phasic_, result.snr_phasic) + self.assertEqual(estimator.noise_, result.noise) + np.testing.assert_array_equal(estimator.peaks_, result.peaks) + np.testing.assert_array_equal(estimator.tonic_, result.tonic) + np.testing.assert_array_equal(estimator.residual_, result.residual) + + +class TestEnvelopeRRSNRAlternateConfigs(unittest.TestCase): + """Test the less-common (non-default) noise_method/tonic_method/ + tonic_range_method options and related config knobs.""" + + def test_folded_iqr_noise_method_runs(self): + trace, _ = make_tonic_only_trace() + result = EnvelopeRRSNR(fps=FPS, noise_method="folded_iqr").fit(trace) + self.assertTrue(np.isfinite(result.noise)) + self.assertGreater(result.noise, 0) + + def test_mad_iqr_avg_noise_method_runs(self): + trace, _ = make_tonic_only_trace() + result = EnvelopeRRSNR(fps=FPS, noise_method="mad_iqr_avg").fit(trace) + self.assertTrue(np.isfinite(result.noise)) + self.assertGreater(result.noise, 0) + + def test_als_tonic_method_runs(self): + trace, _ = make_tonic_only_trace() + result = EnvelopeRRSNR(fps=FPS, config={"tonic_method": "als"}).fit(trace) + self.assertEqual(result.tonic.shape, trace.shape) + self.assertTrue(np.all(np.isfinite(result.tonic))) + + def test_envelope_tonic_method_runs(self): + """tonic_method='envelope' needs its own window-size config keys, unlike + 'als'/'arpls' (which only need their own lam/p/n_iter).""" + trace, _ = make_tonic_only_trace() + result = EnvelopeRRSNR( + fps=FPS, + config={ + "tonic_method": "envelope", + "lower_smooth_window": 31, + "lower_min_distance": 20, + }, + ).fit(trace) + self.assertEqual(result.tonic.shape, trace.shape) + self.assertTrue(np.all(np.isfinite(result.tonic))) + + def test_envelope_tonic_method_with_linear_interpolation(self): + trace, _ = make_tonic_only_trace() + result = EnvelopeRRSNR( + fps=FPS, + config={ + "tonic_method": "envelope", + "lower_smooth_window": 31, + "lower_min_distance": 20, + "interp_kind": "linear", + }, + ).fit(trace) + self.assertTrue(np.all(np.isfinite(result.tonic))) + + def test_ptp_tonic_range_method_runs(self): + trace, _ = make_tonic_only_trace() + result = EnvelopeRRSNR(fps=FPS, config={"tonic_range_method": "ptp"}).fit(trace) + self.assertEqual(result.tonic_range if hasattr(result, "tonic_range") else True, True) + self.assertTrue(np.isfinite(result.snr_tonic)) + + def test_percentile_tonic_range_method_runs(self): + trace, _ = make_tonic_only_trace() + result = EnvelopeRRSNR( + fps=FPS, config={"tonic_range_method": "percentile", "tonic_range_trim_pct": 5.0} + ).fit(trace) + self.assertTrue(np.isfinite(result.snr_tonic)) + + def test_pre_despike_flags_and_removes_a_spike_before_tonic_fitting(self): + trace, _ = make_tonic_only_trace() + trace_with_spike = trace.copy() + trace_with_spike[500:505] += 5.0 # one big spike + + result = EnvelopeRRSNR( + fps=FPS, config={"pre_despike_window": 101, "pre_despike_k": 5.0} + ).fit(trace_with_spike) + + self.assertGreater(result.n_extreme_samples, 0) + self.assertGreater(result.frac_extreme_samples, 0.0) + + def test_scale_window_basic(self): + """scale_window preserves real-world duration across a different fps.""" + self.assertEqual(EnvelopeRRSNR.scale_window(20, fps=40.0), 40) + self.assertEqual(EnvelopeRRSNR.scale_window(20, fps=20.0), 20) + + def test_scale_window_make_odd_bumps_an_even_result(self): + # 30 samples at 20 fps (reference) stays 30 (even) unless make_odd=True. + self.assertEqual(EnvelopeRRSNR.scale_window(30, fps=20.0, make_odd=False), 30) + self.assertEqual(EnvelopeRRSNR.scale_window(30, fps=20.0, make_odd=True), 31) + + +class TestEnvelopeRRSNRFitChunked(unittest.TestCase): + """Test fit_chunked's chunk-sizing options, aggregation modes, and the + too-short-trace fallback.""" + + def test_default_chunking_matches_expected_sizing(self): + """Default sizing is max(min_chunk_duration_s, chunk_fraction * duration).""" + trace, _ = make_phasic_only_trace() # 150 s at 20 fps + result = EnvelopeRRSNR(fps=FPS).fit_chunked(trace) + self.assertEqual(result.chunk_duration_s, 30.0) # max(30, 0.2 * 150) + self.assertEqual(result.n_chunks, 5) + + def test_explicit_chunk_duration_is_respected(self): + trace, _ = make_phasic_only_trace() + result = EnvelopeRRSNR(fps=FPS).fit_chunked(trace, chunk_duration_s=30.0) + self.assertEqual(result.chunk_duration_s, 30.0) + + def test_aggregate_mean_runs_and_differs_from_median_in_general(self): + trace, _ = make_phasic_only_trace() + estimator = EnvelopeRRSNR(fps=FPS) + result_median = estimator.fit_chunked(trace, aggregate="median") + result_mean = estimator.fit_chunked(trace, aggregate="mean") + self.assertTrue(np.isfinite(result_median.snr_tonic)) + self.assertTrue(np.isfinite(result_mean.snr_tonic)) + + def test_invalid_aggregate_raises(self): + trace, _ = make_phasic_only_trace() + with self.assertRaises(ValueError): + EnvelopeRRSNR(fps=FPS).fit_chunked(trace, aggregate="bogus") + + def test_too_short_trace_falls_back_to_global_snr_tonic_and_warns(self): + """A trace with no chunk long enough to fit should warn and fall back to + the global (non-chunked) snr_tonic, not raise or silently return garbage.""" + rng = np.random.default_rng(9) + short_trace = NOISE_SIGMA * rng.standard_normal(80) # 4 s at 20 fps + estimator = EnvelopeRRSNR(fps=FPS) + result_global = estimator.fit(short_trace) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + result_chunked = estimator.fit_chunked(short_trace) + self.assertTrue(any(issubclass(w.category, RuntimeWarning) for w in caught)) + + self.assertEqual(result_chunked.n_chunks, 0) + self.assertEqual(result_chunked.snr_tonic, result_global.snr_tonic) + + def test_bias_correction_propagates_through_chunked_result(self): + """When bias_correction is active and phasic events are detected, the + chunked result's snr_corrected should also be populated (derived from the + chunked snr_tonic + the same snr_phasic_corrected as a plain .fit()).""" + trace, _ = make_phasic_only_trace() + estimator = EnvelopeRRSNR(fps=FPS) # default config -> bias_correction auto-applies + result_global = estimator.fit(trace) + self.assertIsNotNone(result_global.snr_corrected) + + result_chunked = estimator.fit_chunked(trace) + self.assertIsNotNone(result_chunked.snr_corrected) + expected = result_chunked.snr_tonic + result_global.snr_phasic_corrected + self.assertAlmostEqual(result_chunked.snr_corrected, expected) + + +class TestEnvelopeRRSNRBiasCorrection(unittest.TestCase): + """Test bias_correction auto-apply/override logic and apply_correction=True + on the one-shot convenience methods.""" + + def test_default_config_auto_applies_tuned_correction(self): + trace, _ = make_phasic_only_trace() + result = EnvelopeRRSNR(fps=FPS).fit(trace) + self.assertIsNotNone(result.snr_corrected) + self.assertIsNotNone(result.snr_phasic_corrected) + + def test_non_default_noise_method_disables_auto_correction(self): + """The tuned correction was fit against noise_method='aind_mad' specifically + -- switching it should silently resolve bias_correction to None rather than + misapplying a correction fit for a different configuration.""" + trace, _ = make_phasic_only_trace() + result = EnvelopeRRSNR(fps=FPS, noise_method="folded_iqr").fit(trace) + self.assertIsNone(result.snr_corrected) + self.assertIsNone(result.snr_phasic_corrected) + + def test_explicit_none_disables_correction_even_at_tuned_defaults(self): + trace, _ = make_phasic_only_trace() + result = EnvelopeRRSNR(fps=FPS, bias_correction=None).fit(trace) + self.assertIsNone(result.snr_corrected) + + def test_custom_bias_correction_tuple_is_used(self): + trace, _ = make_phasic_only_trace() + slope, intercept = 2.0, 0.5 + result = EnvelopeRRSNR(fps=FPS, bias_correction=(slope, intercept)).fit(trace) + expected_phasic_corrected = (result.snr_phasic - intercept) / slope + self.assertAlmostEqual(result.snr_phasic_corrected, expected_phasic_corrected) + + def test_estimate_apply_correction_returns_corrected_total(self): + trace, _ = make_phasic_only_trace() + estimator = EnvelopeRRSNR(fps=FPS) + result = estimator.fit(trace) + snr_corrected, _, _ = estimator.estimate(trace, apply_correction=True) + self.assertEqual(snr_corrected, result.snr_corrected) + + def test_estimate_components_apply_correction_returns_corrected_values(self): + trace, _ = make_phasic_only_trace() + estimator = EnvelopeRRSNR(fps=FPS) + result = estimator.fit(trace) + total, tonic, phasic = estimator.estimate_components(trace, apply_correction=True) + self.assertEqual(total, result.snr_corrected) + self.assertEqual(tonic, result.snr_tonic) + self.assertEqual(phasic, result.snr_phasic_corrected) + + def test_fit_bias_correction_from_benchmark_recovers_known_linear_bias(self): + true_snr = np.array([5.0, 10.0, 20.0, 40.0]) + snr_est = 0.9 * true_snr + 1.5 # simulated linear bias + slope, intercept = EnvelopeRRSNR.fit_bias_correction_from_benchmark(true_snr, snr_est) + self.assertAlmostEqual(slope, 0.9, places=6) + self.assertAlmostEqual(intercept, 1.5, places=6) + + +class TestPrivateHelperEdgeCases(unittest.TestCase): + """Test a handful of defensive edge cases in private helper functions that + aren't reachable through EnvelopeRRSNR's public API in normal operation (the + public API's own validation/preprocessing prevents these inputs from ever + reaching them), but are still worth pinning down directly since they're real + safety nets against degenerate data.""" + + def test_robust_std_on_empty_array_returns_nan(self): + self.assertTrue(np.isnan(_robust_std(np.array([])))) + + def test_mad_noise_std_returns_nan_on_nan_input(self): + """.fit() always replaces NaNs in the input trace before this is ever + reached, so this guards a case the public API itself prevents.""" + self.assertTrue(np.isnan(_mad_noise_std(np.array([1.0, 2.0, np.nan, 3.0])))) + + def test_mad_noise_std_does_not_collapse_to_nan_on_heavy_contamination(self): + """Regression test: a short window heavily dominated by one huge outlier + used to make the second trim step empty out completely and return NaN, + which would silently poison any downstream aggregation across windows.""" + rng = np.random.default_rng(1) + x = NOISE_SIGMA * rng.standard_normal(600) + x[300:350] += 50.0 # huge, mostly-dominant spike + result = _mad_noise_std(x) + self.assertTrue(np.isfinite(result)) + self.assertGreater(result, 0) + + def test_detect_peaks_rise_rate_handles_a_candidate_at_the_trace_start(self): + """A candidate at index 0 has no samples before it to compute a rise slope + from -- should be handled gracefully (treated as a zero slope), not raise.""" + residual = np.concatenate([[0.0], np.linspace(-1, 1, 25)]) + candidates = np.array([0, 5, 10, 15, 20]) # >=5, so gating actually runs + kept = _detect_peaks_rise_rate(residual, candidates, sigma=0.1, rise_window=3) + self.assertIsInstance(kept, np.ndarray) + + def test_detect_peaks_rise_rate_handles_few_candidates(self): + """With exactly 5 (still >=5, so gating runs) candidates and distinct + slopes, the below-median half naturally has fewer than 5 elements -- + exercises the percentile fallback for the slope threshold.""" + residual = np.linspace(-1, 1, 20) + candidates = np.array([5, 8, 11, 14, 17]) + kept = _detect_peaks_rise_rate(residual, candidates, sigma=0.1, rise_window=3) + self.assertIsInstance(kept, np.ndarray) + + def test_mad_noise_std_returns_rstd_when_second_trim_empties_completely(self): + """If the second trim step's own scale estimate collapses to (near) zero + rather than genuinely emptying the array, the function should fall back + to the first-pass scale estimate rather than trusting a degenerate + second-pass one. Scripted via mock, in call order: the first + _robust_std call (on the first trim) returns a genuine small positive + value, the second (on the second trim) returns exactly 0 -- forcing + `if result > 0` to fail and fall through to the first call's value. + Naturally constructing this exact internal state from a residual array + alone is impractical, since real data essentially never gives an + exactly-zero robust scale on a non-trivial sample.""" + residual = 0.01 * np.random.default_rng(5).standard_normal(300) + call_values = iter([0.001, 0.0]) + with patch( + "aind_dynamic_foraging_basic_analysis.metrics.snr_envelope_rr._robust_std", + side_effect=lambda x: next(call_values), + ): + result = _mad_noise_std(residual) + self.assertEqual(result, 0.001) + + def test_mad_noise_std_falls_back_to_full_residual_when_first_pass_collapses(self): + """If even the FIRST-pass scale estimate collapses to (non-positive) + zero, the second trim is skipped entirely (there's no positive scale to + trim against) and the function falls all the way back to a robust + scale of the whole (untrimmed) detrended residual. Scripted via mock: + the first _robust_std call returns exactly 0, so only one more call + happens (on the untrimmed residual, not a second trim).""" + residual = 0.01 * np.random.default_rng(5).standard_normal(300) + call_values = iter([0.0, 0.007]) + with patch( + "aind_dynamic_foraging_basic_analysis.metrics.snr_envelope_rr._robust_std", + side_effect=lambda x: next(call_values), + ): + result = _mad_noise_std(residual) + self.assertEqual(result, 0.007) + + def test_mad_noise_std_returns_nan_on_empty_input(self): + self.assertTrue(np.isnan(_mad_noise_std(np.array([])))) + + def test_half_sample_mode_base_cases(self): + """_half_sample_mode's recursive narrowing bottoms out at <=3 elements; + exercise each of those base cases (1, 2, and a 3-element tie) directly, + since a large, generic residual essentially never narrows down to + exactly one of these by chance.""" + self.assertEqual(_half_sample_mode(np.array([5.0])), 5.0) + self.assertEqual(_half_sample_mode(np.array([1.0, 2.0])), 1.5) + # 3 elements, evenly spaced -> the two gaps tie, hits the tie branch. + self.assertEqual(_half_sample_mode(np.array([1.0, 2.0, 3.0])), 2.0) + + def test_folded_iqr_noise_std_falls_back_on_few_below_anchor_samples(self): + """With too few samples below the anchor for a stable IQR (<20), falls + back to a plain std of the below-median half.""" + short_residual = 0.01 * np.random.default_rng(0).standard_normal(20) + result = _folded_iqr_noise_std(short_residual) + self.assertTrue(np.isfinite(result)) + self.assertGreater(result, 0) + + def test_estimate_residual_noise_invalid_method_raises(self): + """Dead code from EnvelopeRRSNR's own public API (its noise_method is + already validated at construction time), but still worth pinning as a + safety net for any other internal caller.""" + residual = 0.01 * np.random.default_rng(0).standard_normal(50) + with self.assertRaises(ValueError): + _estimate_residual_noise(residual, "bogus") + + def test_estimate_sigma_minima_falls_back_with_too_few_midpoints(self): + """tonic_minima with exactly 5 elements (the minimum to pass the first + length check) always yields only 4 valley-to-valley midpoints -- one + fewer than needed for a stable estimate -- so this always falls back, + not just in some edge case.""" + residual = np.zeros(100) + tonic_minima = np.array([0, 20, 40, 60, 80]) + result = _estimate_sigma_minima(residual, tonic_minima, fallback_sigma=0.5) + self.assertEqual(result, 0.5) + + def test_interpolate_envelope_invalid_interp_kind_raises(self): + with self.assertRaises(ValueError): + _interpolate_envelope(np.array([0, 5, 10]), np.array([1.0, 2.0, 1.5]), 11, "bogus") + + def test_lower_envelope_handles_an_even_smooth_window(self): + """smooth_window must be odd for Savitzky-Golay filtering -- an even + value should be silently bumped to odd, not raise.""" + x = 0.01 * np.random.default_rng(1).standard_normal(500) + x += 0.05 * np.sin(np.linspace(0, 6, 500)) + tonic, minima = _lower_envelope(x, x.copy(), smooth_window=30, order=2, min_distance=20) + self.assertEqual(tonic.shape, x.shape) + + def test_lower_envelope_falls_back_with_fewer_than_two_minima(self): + """A flat (constant) input has no local minima at all -- should return + a flat curve at the input's own minimum, not raise or return garbage.""" + x = np.full(200, 0.5) + tonic, minima = _lower_envelope(x, x.copy(), smooth_window=31, order=2, min_distance=20) + self.assertLess(len(minima), 2) + np.testing.assert_allclose(tonic, 0.5) + + def test_unset_sentinel_repr(self): + """_UnsetType's repr should be short and unambiguous for debugging -- + never actually shown to a normal caller (it's a default-argument + sentinel), so nothing else exercises it.""" + self.assertEqual(repr(_UNSET), "") + self.assertIsInstance(_UNSET, _UnsetType) + + +if __name__ == "__main__": + unittest.main() From 6a439e82cfe4c9a73591bcf7a01c88a1370e904e Mon Sep 17 00:00:00 2001 From: Kyle Severson <12833749+ksseverson57@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:34:20 -0700 Subject: [PATCH 3/3] test SNR - linting --- tests/test_metric_snr_envelope.py | 52 +++++++++++++++++++++++++------ 1 file changed, 42 insertions(+), 10 deletions(-) diff --git a/tests/test_metric_snr_envelope.py b/tests/test_metric_snr_envelope.py index 5a63e1e..68209af 100644 --- a/tests/test_metric_snr_envelope.py +++ b/tests/test_metric_snr_envelope.py @@ -9,16 +9,20 @@ import numpy as np -from aind_dynamic_foraging_basic_analysis.metrics.snr_envelope_rr import (_UNSET, EnvelopeRRSNR, - _detect_peaks_rise_rate, - _estimate_residual_noise, - _estimate_sigma_minima, - _folded_iqr_noise_std, - _half_sample_mode, - _interpolate_envelope, - _lower_envelope, - _mad_noise_std, - _robust_std, _UnsetType) +from aind_dynamic_foraging_basic_analysis.metrics.snr_envelope_rr import ( + _UNSET, + EnvelopeRRSNR, + _detect_peaks_rise_rate, + _estimate_residual_noise, + _estimate_sigma_minima, + _folded_iqr_noise_std, + _half_sample_mode, + _interpolate_envelope, + _lower_envelope, + _mad_noise_std, + _robust_std, + _UnsetType, +) # Shared simulation parameters, reused across every test so intent stays visible at # each call site (e.g. `event_amplitude=0.2` reads as "20x NOISE_SIGMA") without @@ -275,10 +279,12 @@ class TestEnvelopeRRSNRConstructorValidation(unittest.TestCase): """Test constructor argument validation and the not-yet-fitted guard.""" def test_invalid_signal_statistic_raises(self): + """An unrecognized signal_statistic should raise, not silently fall back.""" with self.assertRaises(ValueError): EnvelopeRRSNR(fps=FPS, signal_statistic="bogus") def test_invalid_noise_method_via_argument_raises(self): + """An unrecognized noise_method passed as a constructor argument should raise.""" with self.assertRaises(ValueError): EnvelopeRRSNR(fps=FPS, noise_method="bogus") @@ -289,6 +295,7 @@ def test_invalid_noise_method_via_config_raises(self): EnvelopeRRSNR(fps=FPS, config={"noise_method": "bogus"}) def test_deprecated_mad_alias_resolves_to_aind_mad(self): + """The deprecated 'mad' alias should silently resolve to 'aind_mad'.""" estimator = EnvelopeRRSNR(fps=FPS, noise_method="mad") self.assertEqual(estimator.noise_method, "aind_mad") @@ -300,6 +307,7 @@ def test_invalid_tonic_method_raises_on_fit(self): estimator.fit(trace) def test_invalid_tonic_range_method_raises_on_fit(self): + """tonic_range_method isn't validated until .fit() dispatches on it either.""" trace, _ = make_tonic_only_trace() estimator = EnvelopeRRSNR(fps=FPS, config={"tonic_range_method": "bogus"}) with self.assertRaises(ValueError): @@ -345,18 +353,21 @@ class TestEnvelopeRRSNRAlternateConfigs(unittest.TestCase): tonic_range_method options and related config knobs.""" def test_folded_iqr_noise_method_runs(self): + """noise_method='folded_iqr' should run and give a sane, positive noise estimate.""" trace, _ = make_tonic_only_trace() result = EnvelopeRRSNR(fps=FPS, noise_method="folded_iqr").fit(trace) self.assertTrue(np.isfinite(result.noise)) self.assertGreater(result.noise, 0) def test_mad_iqr_avg_noise_method_runs(self): + """noise_method='mad_iqr_avg' should run and give a sane, positive noise estimate.""" trace, _ = make_tonic_only_trace() result = EnvelopeRRSNR(fps=FPS, noise_method="mad_iqr_avg").fit(trace) self.assertTrue(np.isfinite(result.noise)) self.assertGreater(result.noise, 0) def test_als_tonic_method_runs(self): + """tonic_method='als' should run and produce a finite, full-length tonic curve.""" trace, _ = make_tonic_only_trace() result = EnvelopeRRSNR(fps=FPS, config={"tonic_method": "als"}).fit(trace) self.assertEqual(result.tonic.shape, trace.shape) @@ -378,6 +389,7 @@ def test_envelope_tonic_method_runs(self): self.assertTrue(np.all(np.isfinite(result.tonic))) def test_envelope_tonic_method_with_linear_interpolation(self): + """interp_kind='linear' is a valid alternative to the default 'pchip'.""" trace, _ = make_tonic_only_trace() result = EnvelopeRRSNR( fps=FPS, @@ -391,12 +403,14 @@ def test_envelope_tonic_method_with_linear_interpolation(self): self.assertTrue(np.all(np.isfinite(result.tonic))) def test_ptp_tonic_range_method_runs(self): + """tonic_range_method='ptp' (the pre-'robust' default) should still run cleanly.""" trace, _ = make_tonic_only_trace() result = EnvelopeRRSNR(fps=FPS, config={"tonic_range_method": "ptp"}).fit(trace) self.assertEqual(result.tonic_range if hasattr(result, "tonic_range") else True, True) self.assertTrue(np.isfinite(result.snr_tonic)) def test_percentile_tonic_range_method_runs(self): + """tonic_range_method='percentile' should run cleanly with a trim_pct set.""" trace, _ = make_tonic_only_trace() result = EnvelopeRRSNR( fps=FPS, config={"tonic_range_method": "percentile", "tonic_range_trim_pct": 5.0} @@ -404,6 +418,7 @@ def test_percentile_tonic_range_method_runs(self): self.assertTrue(np.isfinite(result.snr_tonic)) def test_pre_despike_flags_and_removes_a_spike_before_tonic_fitting(self): + """Enabling pre_despike_window should flag and count a large injected spike.""" trace, _ = make_tonic_only_trace() trace_with_spike = trace.copy() trace_with_spike[500:505] += 5.0 # one big spike @@ -421,6 +436,7 @@ def test_scale_window_basic(self): self.assertEqual(EnvelopeRRSNR.scale_window(20, fps=20.0), 20) def test_scale_window_make_odd_bumps_an_even_result(self): + """make_odd=True should bump an otherwise-even scaled result up by one.""" # 30 samples at 20 fps (reference) stays 30 (even) unless make_odd=True. self.assertEqual(EnvelopeRRSNR.scale_window(30, fps=20.0, make_odd=False), 30) self.assertEqual(EnvelopeRRSNR.scale_window(30, fps=20.0, make_odd=True), 31) @@ -438,11 +454,13 @@ def test_default_chunking_matches_expected_sizing(self): self.assertEqual(result.n_chunks, 5) def test_explicit_chunk_duration_is_respected(self): + """An explicit chunk_duration_s should be used as-is, not the default sizing.""" trace, _ = make_phasic_only_trace() result = EnvelopeRRSNR(fps=FPS).fit_chunked(trace, chunk_duration_s=30.0) self.assertEqual(result.chunk_duration_s, 30.0) def test_aggregate_mean_runs_and_differs_from_median_in_general(self): + """aggregate='mean' is a valid alternative aggregation mode to the default 'median'.""" trace, _ = make_phasic_only_trace() estimator = EnvelopeRRSNR(fps=FPS) result_median = estimator.fit_chunked(trace, aggregate="median") @@ -451,6 +469,7 @@ def test_aggregate_mean_runs_and_differs_from_median_in_general(self): self.assertTrue(np.isfinite(result_mean.snr_tonic)) def test_invalid_aggregate_raises(self): + """An unrecognized aggregate value should raise, not silently default.""" trace, _ = make_phasic_only_trace() with self.assertRaises(ValueError): EnvelopeRRSNR(fps=FPS).fit_chunked(trace, aggregate="bogus") @@ -491,6 +510,7 @@ class TestEnvelopeRRSNRBiasCorrection(unittest.TestCase): on the one-shot convenience methods.""" def test_default_config_auto_applies_tuned_correction(self): + """Constructing with all-default settings should auto-apply the tuned bias correction.""" trace, _ = make_phasic_only_trace() result = EnvelopeRRSNR(fps=FPS).fit(trace) self.assertIsNotNone(result.snr_corrected) @@ -506,11 +526,14 @@ def test_non_default_noise_method_disables_auto_correction(self): self.assertIsNone(result.snr_phasic_corrected) def test_explicit_none_disables_correction_even_at_tuned_defaults(self): + """bias_correction=None should disable the correction even when the rest of the + config matches the tuned defaults exactly.""" trace, _ = make_phasic_only_trace() result = EnvelopeRRSNR(fps=FPS, bias_correction=None).fit(trace) self.assertIsNone(result.snr_corrected) def test_custom_bias_correction_tuple_is_used(self): + """A custom (slope, intercept) tuple should be applied exactly as given.""" trace, _ = make_phasic_only_trace() slope, intercept = 2.0, 0.5 result = EnvelopeRRSNR(fps=FPS, bias_correction=(slope, intercept)).fit(trace) @@ -518,6 +541,7 @@ def test_custom_bias_correction_tuple_is_used(self): self.assertAlmostEqual(result.snr_phasic_corrected, expected_phasic_corrected) def test_estimate_apply_correction_returns_corrected_total(self): + """.estimate(apply_correction=True) should return the corrected total, matching .fit().""" trace, _ = make_phasic_only_trace() estimator = EnvelopeRRSNR(fps=FPS) result = estimator.fit(trace) @@ -525,6 +549,8 @@ def test_estimate_apply_correction_returns_corrected_total(self): self.assertEqual(snr_corrected, result.snr_corrected) def test_estimate_components_apply_correction_returns_corrected_values(self): + """.estimate_components(apply_correction=True) should return the corrected + total/tonic/phasic triple, matching .fit().""" trace, _ = make_phasic_only_trace() estimator = EnvelopeRRSNR(fps=FPS) result = estimator.fit(trace) @@ -534,6 +560,8 @@ def test_estimate_components_apply_correction_returns_corrected_values(self): self.assertEqual(phasic, result.snr_phasic_corrected) def test_fit_bias_correction_from_benchmark_recovers_known_linear_bias(self): + """Fitting against a synthetic, exactly-linear bias should recover its + true slope/intercept.""" true_snr = np.array([5.0, 10.0, 20.0, 40.0]) snr_est = 0.9 * true_snr + 1.5 # simulated linear bias slope, intercept = EnvelopeRRSNR.fit_bias_correction_from_benchmark(true_snr, snr_est) @@ -549,6 +577,7 @@ class TestPrivateHelperEdgeCases(unittest.TestCase): safety nets against degenerate data.""" def test_robust_std_on_empty_array_returns_nan(self): + """An empty input has no meaningful spread to estimate -- should return NaN.""" self.assertTrue(np.isnan(_robust_std(np.array([])))) def test_mad_noise_std_returns_nan_on_nan_input(self): @@ -621,6 +650,8 @@ def test_mad_noise_std_falls_back_to_full_residual_when_first_pass_collapses(sel self.assertEqual(result, 0.007) def test_mad_noise_std_returns_nan_on_empty_input(self): + """An empty residual has nothing to estimate a noise floor from -- should + return NaN rather than raise.""" self.assertTrue(np.isnan(_mad_noise_std(np.array([])))) def test_half_sample_mode_base_cases(self): @@ -660,6 +691,7 @@ def test_estimate_sigma_minima_falls_back_with_too_few_midpoints(self): self.assertEqual(result, 0.5) def test_interpolate_envelope_invalid_interp_kind_raises(self): + """An unrecognized interp_kind should raise, not silently fall back.""" with self.assertRaises(ValueError): _interpolate_envelope(np.array([0, 5, 10]), np.array([1.0, 2.0, 1.5]), 11, "bogus")