[FIX] ASR: use a real PSD square root in asr_calibrate

meegkit/asr.py

@@ -444,6 +444,20 @@ def _rms_window_offsets(ns, N, win_overlap):
     return np.round(np.arange(0, ns - N, N * (1 - win_overlap))).astype(int)
 
 
+def _psd_sqrtm(A):
+    """Real symmetric positive-semidefinite matrix square root.
+
+    Symmetrizes ``A`` and clamps negative eigenvalues to zero before taking
+    the square root, so the result stays real even for a slightly-indefinite
+    input (where ``scipy.linalg.sqrtm`` would return a complex matrix).
+    """
+    A = np.real(A)
+    A = (A + A.T) / 2
+    evals, evecs = linalg.eigh(A)
+    evals = np.clip(evals, 0, None)
+    return (evecs * np.sqrt(evals)) @ evecs.T
+
+
 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"):
@@ -544,7 +558,7 @@ def asr_calibrate(X, sfreq, cutoff=5, blocksize=100, win_len=0.5,
         Uavg = mean_covariance(U, metric="riemann")
 
     # get the mixing matrix M
-    M = linalg.sqrtm(np.real(Uavg))
+    M = _psd_sqrtm(Uavg)
     D, Vtmp = linalg.eigh(M)
     # D, Vtmp = nonlinear_eigenspace(M, nc)  TODO
     V = Vtmp[:, np.argsort(D)]

tests/test_asr.py

@@ -214,6 +214,34 @@ def test_clean_windows_one_sided_zthresholds(zthresholds):
     assert sample_mask.dtype == bool
 
 
+def test_psd_sqrtm_real_on_indefinite():
+    """_psd_sqrtm stays real on a slightly indefinite matrix (unlike scipy sqrtm)."""
+    from scipy import linalg
+
+    from meegkit.asr import _psd_sqrtm
+
+    # Build a symmetric matrix with one tiny negative eigenvalue.
+    B = rng.standard_normal((5, 5))
+    A = B @ B.T  # PSD
+    evals, evecs = linalg.eigh(A)
+    evals[0] = -1e-8
+    A_indef = (evecs * evals) @ evecs.T
+
+    M = _psd_sqrtm(A_indef)
+
+    # The result must stay real...
+    assert np.isrealobj(M)
+
+    # ...and square back to the clamped PSD matrix.
+    evals_clamped = np.clip(evals, 0, None)
+    A_psd = (evecs * evals_clamped) @ evecs.T
+    assert np.allclose(M @ M, A_psd, atol=1e-8)
+
+    # plain sqrtm(real(A)) on the indefinite matrix still returns a complex root
+    M_bad = linalg.sqrtm(np.real(A_indef))
+    assert np.iscomplexobj(M_bad) and not np.allclose(M_bad.imag, 0)
+
+
 @pytest.mark.parametrize(argnames="method", argvalues=("riemann", "euclid"))
 @pytest.mark.parametrize(argnames="reref", argvalues=(False, True))
 def test_asr_class(method, reref, show=False):