From be557df5970818a11b37ef23f9fff180fc9de19f Mon Sep 17 00:00:00 2001
From: Stefan Appelhoff <stefan.appelhoff@mailbox.org>
Date: Tue, 14 Jul 2026 10:24:09 +0200
Subject: [PATCH] [MAINT] ASR: fit_eeg_distribution docstring fixes +
 asr_calibrate input guards

- fit_eeg_distribution: correct the transposed min_clean_fraction /
  max_dropout_fraction default values in the docstring, and describe X as
  the 1-D amplitude vector it actually is (not a 2-D array).
- asr_calibrate: zero non-finite samples before filtering so a NaN/Inf in
  the calibration data no longer propagates into a silent all-NaN
  threshold; and raise clear errors when the calibration data is too short
  to form at least two analysis windows, instead of an obscure downstream
  IndexError.
---
 meegkit/asr.py       | 12 ++++++++++++
 meegkit/utils/asr.py |  9 +++++----
 tests/test_asr.py    | 23 +++++++++++++++++++++++
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/meegkit/asr.py b/meegkit/asr.py
index 75ea41aa..fd40c811 100755
--- a/meegkit/asr.py
+++ b/meegkit/asr.py
@@ -479,11 +479,18 @@ def asr_calibrate(X, sfreq, cutoff=5, blocksize=100, win_len=0.5,
     # set number of channels and number of samples
     [nc, ns] = X.shape
 
+    # avoid propagating non-finite samples into the filter and thresholds
+    X = np.where(np.isfinite(X), X, 0.0)
+
     # filter the data
     X, _zf = yulewalk_filter(X, sfreq, ab=None)
 
     # window length for calculating thresholds
     N = int(np.round(win_len * sfreq))
+    if ns < N:
+        raise ValueError(
+            f"Calibration data has {ns} samples, shorter than one analysis "
+            f"window of {N} samples (win_len={win_len} at sfreq={sfreq}).")
 
     U = block_covariance(X, window=blocksize, overlap=win_overlap,
                          estimator=estimator)
@@ -502,6 +509,11 @@ def asr_calibrate(X, sfreq, cutoff=5, blocksize=100, win_len=0.5,
     # get the threshold matrix T
     xsq = np.dot(V.T, X) ** 2
     offsets = np.arange(0, ns - N, np.round(N * (1 - win_overlap))).astype(int)
+    if len(offsets) < 2:
+        raise ValueError(
+            f"Only {len(offsets)} window(s) available for threshold "
+            "estimation; at least 2 are required. Provide more calibration "
+            "data.")
 
     # root mean squared amplitude per channel (windowed sums via cumulative sum)
     csum = np.zeros((nc, ns + 1))
diff --git a/meegkit/utils/asr.py b/meegkit/utils/asr.py
index 3c5ced93..010613a0 100755
--- a/meegkit/utils/asr.py
+++ b/meegkit/utils/asr.py
@@ -38,16 +38,17 @@ def fit_eeg_distribution(X, min_clean_fraction=0.25, max_dropout_fraction=0.1,
 
     Parameters
     ----------
-    X : array, shape=(n_channels, n_samples)
-        EEG data, possibly containing artifacts.
+    X : array, shape=(n_samples,)
+        1-D vector of amplitude values (e.g., per-window RMS), possibly
+        containing artifacts.
     min_clean_fraction : float
         Minimum fraction that needs to be clean. This is the minimum fraction
         of time windows that need to contain essentially uncontaminated EEG
-        (default=0.1).
+        (default=0.25).
     max_dropout_fraction : float
         Maximum fraction that can have dropouts. This is the maximum fraction
         of time windows that may have arbitrarily low amplitude (e.g., due to
-        the sensors being unplugged) (default=0.25).
+        the sensors being unplugged) (default=0.1).
     fit_quantiles : 2-tuple
         Quantile range [lower,upper] of the truncated generalized Gaussian
         distribution that shall be fit to the EEG contents (default=[0.022
diff --git a/tests/test_asr.py b/tests/test_asr.py
index 811bcb38..04aa3a63 100644
--- a/tests/test_asr.py
+++ b/tests/test_asr.py
@@ -276,6 +276,29 @@ def test_asr_class(method, reref, show=False):
     ASR(sfreq=125)
     ASR(Sfreq=150)
 
+
+def test_asr_calibrate_nonfinite_input():
+    """Non-finite calibration samples must not silently yield NaN M/T."""
+    raw = np.load(os.path.join(THIS_FOLDER, "data", "eeg_raw.npy"))
+    sfreq = 250
+    train_idx = np.arange(5 * sfreq, 45 * sfreq, dtype=int)
+    X = raw[:, train_idx].copy()
+    X[0, 100] = np.nan
+    X[0, 200] = np.inf
+
+    M, T = asr_calibrate(X, sfreq)
+
+    assert np.all(np.isfinite(M))
+    assert np.all(np.isfinite(T))
+
+
+def test_asr_calibrate_too_short():
+    """Calibration data shorter than one analysis window should raise."""
+    X_short = rng.standard_normal((8, 50))  # N = round(0.5 * 250) = 125
+    with pytest.raises(ValueError, match="shorter than one analysis window"):
+        asr_calibrate(X_short, 250)
+
+
 if __name__ == "__main__":
     pytest.main([__file__])
     # test_yulewalk(250, True)