[FIX] ASR: tolerate flat/degenerate segments in fit_eeg_distribution

meegkit/utils/asr.py

@@ -125,11 +125,16 @@ def fit_eeg_distribution(X, min_clean_fraction=0.25, max_dropout_fraction=0.1,
         H = newX[:m] * cols
 
         # per-column histogram over integer bins [0, 1, ..., nbins]
+        # a tied/flat run (e.g. a dead channel) can make a column's baseline
+        # value 0, so H may contain inf/nan; such entries are counted in no
+        # bin instead of crashing the cast to int
         ncol = H.shape[1]
-        binidx = np.minimum(H.astype(int), nbins - 1)
+        finite = np.isfinite(H)
+        binidx = np.minimum(np.where(finite, H, 0).astype(int), nbins - 1)
         flat = binidx + np.arange(ncol) * nbins
         hist_all = np.bincount(
-            flat.ravel(), minlength=nbins * ncol).reshape(ncol, nbins).T
+            flat.ravel(), weights=finite.ravel(),
+            minlength=nbins * ncol).reshape(ncol, nbins).T
         hist_all = np.vstack((hist_all, np.zeros(ncol, dtype=int)))
         logq = np.log(hist_all + 0.01)
 

tests/test_asr.py

@@ -170,6 +170,50 @@ def test_round_half_away():
     np.testing.assert_array_equal(_round_half_away(ties), expected)
 
 
+def test_fit_eeg_distribution_flat_run():
+    """fit_eeg_distribution must not crash on a run of tied values."""
+    vec = np.r_[np.zeros(700), np.linspace(0, 1, 300)]
+    mu, sig, alpha, beta = fit_eeg_distribution(vec)
+    assert np.isfinite(mu)
+    assert np.isfinite(sig)
+    assert np.isfinite(alpha)
+    assert np.isfinite(beta)
+
+
+def test_fit_eeg_distribution_clean_data():
+    """fit_eeg_distribution result on clean data is unchanged."""
+    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]
+
+    win_len, win_overlap = 0.5, 0.66
+    N = int(win_len * sfreq)
+    ns = X.shape[1]
+    offsets = np.round(
+        np.arange(0, ns - N, (N * (1 - win_overlap)))).astype(int)
+    csum = np.zeros((1, ns + 1))
+    np.cumsum(X[0:1] ** 2, axis=1, out=csum[:, 1:])
+    rms = np.sqrt((csum[:, offsets + N] - csum[:, offsets]) / N)[0]
+
+    mu, sig, alpha, beta = fit_eeg_distribution(rms)
+    assert np.isfinite(mu) and np.isfinite(sig)
+    expected = (8.001450519084575, 0.5339755761243605, 0.8942468424981882, 3.35)
+    np.testing.assert_allclose((mu, sig, alpha, beta), expected)
+
+
+def test_asr_fit_dead_channel():
+    """ASR.fit must not crash when a channel is flatlined."""
+    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] = 0.0  # dead channel
+
+    asr = ASR(sfreq=sfreq)
+    asr.fit(X)
+
+
 def test_asr_functions(show=False, method="riemann"):
     """Test ASR functions (offline use).