[FIX] ASR: correct reconstruction dim count, blend seam, and RMS window rounding

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)

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)