diff --git a/meegkit/utils/asr.py b/meegkit/utils/asr.py index 0df8d154..d2524259 100755 --- a/meegkit/utils/asr.py +++ b/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) diff --git a/tests/test_asr.py b/tests/test_asr.py index 842e6586..ec66ff56 100644 --- a/tests/test_asr.py +++ b/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).