[FIX] Ritz pair rotation to address review on PR #100

Author: nbara

meegkit/utils/covariances.py

@@ -503,8 +503,13 @@ def ehess(X, U):
     Xsol = solver.run(problem, initial_point=U0)
     X = Xsol.point
 
-    # Keep eigpairs consistent: compute each optimized vector's Rayleigh
-    # quotient against L instead of returning the random-init eigenvalues S0.
-    S = np.real(np.diag(X.T @ L @ X))
-
-    return S, X
+    # The Grassmann manifold only fixes the *subspace* spanned by X, so X carries
+    # an arbitrary internal rotation.  Diagonalise the k×k projected matrix to
+    # obtain proper Ritz pairs: after rotation, X.T @ L @ X is diagonal and each
+    # column of X is a genuine approximate eigenvector of L.
+    T = X.T @ L @ X
+    T = (T + T.T) / 2  # symmetrise against numerical noise
+    S, R = linalg.eigh(T)
+    X = X @ R
+
+    return np.real(S), np.real(X)

tests/test_cov.py

@@ -91,23 +91,30 @@ def test_convmtx():
 
 
 def test_nonlinear_eigenspace_consistent_eigpairs():
-    """Returned eigenvalues must match returned eigenvectors."""
+    """Each returned eigenvalue must match the Rayleigh quotient of its column."""
     pytest.importorskip("pymanopt")
     from meegkit.utils.covariances import nonlinear_eigenspace
 
     A = rng.standard_normal((6, 6))
     L = A.T @ A + np.eye(6) * 1e-6
 
     S, X = nonlinear_eigenspace(L, 6)
+    # X.T @ L @ X should be (near) diagonal; check the diagonal matches S
+    # element-wise so a permutation between S and X columns cannot slip through.
     rayleigh = np.real(np.diag(X.T @ L @ X))
 
     np.testing.assert_allclose(
-        np.sort(np.real(S)),
-        np.sort(rayleigh),
+        S,
+        rayleigh,
         rtol=1e-3,
         atol=1e-6,
     )
 
+    # Off-diagonal entries of the projected matrix should be negligible.
+    proj = X.T @ L @ X
+    off_diag = proj - np.diag(np.diag(proj))
+    assert np.max(np.abs(off_diag)) < 1e-6 * np.max(np.abs(np.diag(proj)))
+
 if __name__ == "__main__":
     import pytest
     pytest.main([__file__])