diff --git a/meegkit/asr.py b/meegkit/asr.py index 9124b1f1..4568a60c 100755 --- a/meegkit/asr.py +++ b/meegkit/asr.py @@ -417,6 +417,15 @@ def clean_windows(X, sfreq, max_bad_chans=0.2, zthresholds=[-3.5, 5], return clean, sample_mask +def _rms_window_offsets(ns, N, win_overlap): + """Window start indices for windowed RMS. + + Rounds each window START (not the step), matching ``clean_windows`` and + avoiding phase drift when ``N * (1 - win_overlap)`` is non-integer. + """ + return np.round(np.arange(0, ns - N, N * (1 - win_overlap))).astype(int) + + def asr_calibrate(X, sfreq, cutoff=5, blocksize=100, win_len=0.5, win_overlap=0.66, max_dropout_fraction=0.1, min_clean_fraction=0.25, method="euclid", estimator="scm"): @@ -524,7 +533,7 @@ 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) + offsets = _rms_window_offsets(ns, N, win_overlap) if len(offsets) < 2: raise ValueError( f"Only {len(offsets)} window(s) available for threshold " @@ -604,7 +613,7 @@ def asr_process(X, X_filt, state, cov=None, detrend=False, method="riemann", cov = geometric_median(cov.reshape((-1, nc * nc))) cov = cov.reshape((nc, nc)) - maxdims = int(np.fix(0.66 * nc)) # constant TODO make param + maxdims = int(np.floor(0.66 * nc + 0.5)) # constant TODO make param # do a PCA to find potential artifacts if method == "riemann": @@ -618,7 +627,7 @@ def asr_process(X, X_filt, state, cov=None, detrend=False, method="riemann", # determine which components to keep (variance below directional threshold # or not admissible for rejection) keep = (D < np.sum(np.dot(T, V)**2, axis=0)) - keep += (np.arange(nc) < nc - maxdims) + keep += (np.arange(nc) < nc - maxdims - 1) # update the reconstruction matrix R (reconstruct artifact components using # the mixing matrix) @@ -632,7 +641,7 @@ def asr_process(X, X_filt, state, cov=None, detrend=False, method="riemann", if state["R"] is not None: # apply the reconstruction to intermediate samples (using raised-cosine # blending) - blend = (1 - np.cos(np.pi * np.arange(ns) / ns)) / 2 + blend = (1 - np.cos(np.pi * np.arange(1, ns + 1) / ns)) / 2 clean = blend * R.dot(X) + (1 - blend) * state["R"].dot(X) else: clean = R.dot(X) diff --git a/tests/test_asr.py b/tests/test_asr.py index 635a3055..8dc280ea 100644 --- a/tests/test_asr.py +++ b/tests/test_asr.py @@ -301,7 +301,6 @@ def test_asr_class(method, reref, show=False): ASR(sfreq=125) ASR(Sfreq=150) - def test_clean_windows_uses_full_shape_range(): """clean_windows passes the full 13-point SHAPE_RANGE (incl beta=3.5) down. @@ -537,6 +536,70 @@ def test_block_covariance_no_hang_and_float_window(): assert np.isfinite(cov).all() +@pytest.mark.parametrize("nc, expected", [(8, 6), (32, 22), (64, 43), (19, 14)]) +def test_asr_process_removed_dims(nc, expected): + """Removed-dim count uses round-half-away maxdims and the fixed keep-mask index. + + T is zeroed so the variance test is uniformly False and only the index + term decides how many dimensions are removed. + """ + ns = 20 * nc + X = rng.standard_normal((nc, ns)) + cov = X @ X.T / ns + 1e-6 * np.eye(nc) # well-conditioned PSD + M = np.eye(nc) + T = np.zeros((nc, nc)) + state = dict(M=M, T=T, R=None) + + clean, state = asr_process(X, X, state, cov=cov[None], method="euclid") + + removed = nc - np.linalg.matrix_rank(state["R"]) + assert removed == expected + + +def test_asr_process_blend_reaches_one(): + """Raised-cosine blend reaches exactly 1.0 at the block-final sample. + + Two calls so the second blends against the first reconstruction; the + final sample must equal the new reconstruction with no residual seam. + """ + nc, ns = 6, 100 + X1 = rng.standard_normal((nc, ns)) + X2 = rng.standard_normal((nc, ns)) + cov1 = X1 @ X1.T / ns + 1e-6 * np.eye(nc) + cov2 = X2 @ X2.T / ns + 1e-6 * np.eye(nc) + M = np.eye(nc) + T = 0.5 * np.eye(nc) # non-trivial threshold so R is non-identity + + state = dict(M=M, T=T, R=None) + _, state = asr_process(X1, X1, state, cov=cov1[None], method="euclid") + R_old = state["R"].copy() + + clean2, state2 = asr_process(X2, X2, state, cov=cov2[None], method="euclid") + R_new = state2["R"] + + # meaningfulness guard: blend is only visible when R changes. + assert not np.allclose(R_old, R_new) + assert np.allclose(clean2[:, -1], R_new.dot(X2)[:, -1], rtol=1e-9) + + +def test_rms_window_offsets_rounds_each_start(): + """RMS window starts round each start, not the step (matches clean_windows).""" + from meegkit.asr import _rms_window_offsets + + ns, N, win_overlap = 1000, 125, 0.66 # step = 42.5 (non-integer) + + # rounding the step first gives step 42 -> [0, 42, 84, ...] + old_step = np.round(N * (1 - win_overlap)) + old_offsets = np.arange(0, ns - N, old_step).astype(int) + + offsets = _rms_window_offsets(ns, N, win_overlap) + + assert not np.array_equal(offsets, old_offsets) + # float step is 42.4999..., so 3*step rounds to 127 (not 128) + np.testing.assert_array_equal( + offsets[:7], np.array([0, 42, 85, 127, 170, 212, 255])) + + if __name__ == "__main__": pytest.main([__file__]) # test_yulewalk(250, True)