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]]: diff --git a/tests/test_metric_snr_envelope.py b/tests/test_metric_snr_envelope.py new file mode 100644 index 0000000..68209af --- /dev/null +++ b/tests/test_metric_snr_envelope.py @@ -0,0 +1,723 @@ +"""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): + """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") + + 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): + """The deprecated 'mad' alias should silently resolve to 'aind_mad'.""" + 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): + """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): + 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): + """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) + 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): + """interp_kind='linear' is a valid alternative to the default 'pchip'.""" + 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): + """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} + ).fit(trace) + 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 + + 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): + """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) + + +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): + """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") + 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): + """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") + + 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): + """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) + 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): + """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) + 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): + """.estimate(apply_correction=True) should return the corrected total, matching .fit().""" + 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): + """.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) + 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): + """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) + 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): + """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): + """.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): + """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): + """_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): + """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") + + 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()