Change the ordering of eigen-/singular-vectors in eig.py (#235)

* Change ordering of the returned vecs in eig.eig_partial

* Update the eig.eigh_partial similarly

* Delete an incorrect docstring

* Update eig.svd_partial similarly

matfree/backend/linalg.py

@@ -21,7 +21,14 @@ def eigh(x, /):
 
 
 def eig(x, /):
-    return jnp.linalg.eig(x)
+    vals, vecs = jnp.linalg.eig(x)
+
+    # Order the values and vectors by magnitude
+    # (jax.numpy.linalg.eig does not do this)
+    ordered = jnp.argsort(vals)[::-1]
+    vals = vals[ordered]
+    vecs = vecs[:, ordered]
+    return vals, vecs
 
 
 def cholesky(x, /):

matfree/eig.py

@@ -1,6 +1,6 @@
 """Matrix-free eigenvalue and singular-value analysis."""
 
-from matfree.backend import linalg, np
+from matfree.backend import linalg
 from matfree.backend.typing import Array, Callable
 
 
@@ -28,7 +28,7 @@ def svd(Av: Callable, v0: Array, *parameters):
         U, S, Vt = linalg.svd(B, full_matrices=False)
 
         # Combine orthogonal transformations
-        return u @ U, S, Vt @ v.T
+        return (u @ U).T, S, Vt @ v.T
 
     return svd
 
@@ -55,7 +55,7 @@ def eigh(Av: Callable, v0: Array, *parameters):
         # Compute SVD of factorisation
         vals, vecs = linalg.eigh(H)
         vecs = Q @ vecs
-        return vals, vecs
+        return vals, vecs.T
 
     return eigh
 
@@ -82,12 +82,6 @@ def eig(Av: Callable, v0: Array, *parameters):
         # Compute SVD of factorisation
         vals, vecs = linalg.eig(H)
         vecs = Q @ vecs
-
-        # Return in descending order
-        # (jnp.linalg.eig doesn't do this by default)
-        args = np.argsort(vals)[::-1]
-        vals = vals[args]
-        vecs = vecs[:, args]
-        return vals, vecs
+        return vals, vecs.T
 
     return eig

tests/test_eig/test_eig_partial.py

@@ -11,28 +11,16 @@ def test_equal_to_linalg_eig(nrows=7):
     and the orthogonal matrices should be orthogonal. They are not unique.
     """
     A = np.triu(np.arange(1.0, 1.0 + nrows**2).reshape((nrows, nrows)))
-
-    nrows, ncols = np.shape(A)
-    num_matvecs = min(nrows, ncols)
-
-    v0 = np.ones((ncols,))
-    v0 /= linalg.vector_norm(v0)
+    v0 = np.ones((nrows,))
+    num_matvecs = nrows
 
     hessenberg = decomp.hessenberg(num_matvecs, reortho="full")
     alg = eig.eig_partial(hessenberg)
-    S, U = alg(lambda v, p: p @ v, v0, A)
-
-    # Ensure the reference is ordered
-    S_, U_ = linalg.eig(A)
-    ordered = np.argsort(S_)[::-1]
-    S_ = S_[ordered]
-    U_ = U_[:, ordered]
-
-    assert np.allclose(S, S_)
-    assert np.shape(U) == np.shape(U_)
+    vals, vecs = alg(lambda v, p: p @ v, v0, A)
 
-    tols_decomp = {"atol": 1e-5, "rtol": 1e-5}
-    assert np.allclose(U @ U.T, U_ @ U_.T, **tols_decomp)
+    S, U = linalg.eig(A)
+    assert np.allclose(vecs.T @ vecs, U @ U.T, atol=1e-5, rtol=1e-5)
+    assert np.allclose(vals, S)
 
 
 @testing.parametrize("nrows", [8])
@@ -47,4 +35,4 @@ def test_shapes_as_expected(nrows, num_matvecs):
     alg = eig.eig_partial(hessenberg)
     S, U = alg(lambda v, p: p @ v, v0, A)
     assert S.shape == (num_matvecs,)
-    assert U.shape == (nrows, num_matvecs)
+    assert U.shape == (num_matvecs, nrows)

tests/test_eig/test_eigh_partial.py

@@ -4,51 +4,30 @@
 from matfree.backend import linalg, np, testing
 
 
-@testing.fixture()
-@testing.parametrize("nrows", [10])
-def A(nrows):
-    """Make a positive definite matrix with certain spectrum."""
-    # 'Invent' a spectrum. Use the number of pre-defined eigenvalues.
-    d = np.arange(1.0, 1.0 + nrows)
-    return test_util.asymmetric_matrix_from_singular_values(d, nrows=nrows, ncols=nrows)
-
-
-def test_equal_to_linalg_eigh(A):
-    """The output of full-depth decomposition should be equal (*) to linalg.svd().
-
-    (*) Note: The singular values should be identical,
-    and the orthogonal matrices should be orthogonal. They are not unique.
-    """
-    nrows, ncols = np.shape(A)
-    num_matvecs = min(nrows, ncols)
-
-    v0 = np.ones((ncols,))
-    v0 /= linalg.vector_norm(v0)
+def test_equal_to_linalg_eigh(nrows=10):
+    eigvals = np.arange(1.0, 1.0 + nrows)
+    A = test_util.symmetric_matrix_from_eigenvalues(eigvals)
+    v0 = np.ones((nrows,))
+    num_matvecs = nrows
 
     hessenberg = decomp.hessenberg(num_matvecs, reortho="full")
     alg = eig.eigh_partial(hessenberg)
-    S, U = alg(lambda v, p: p @ v, v0, A)
-
-    S_, U_ = linalg.eigh(A)
+    vals, vecs = alg(lambda v, p: p @ v, v0, A)
 
-    assert np.allclose(S, S_)
-    assert np.shape(U) == np.shape(U_)
-
-    tols_decomp = {"atol": 1e-5, "rtol": 1e-5}
-    assert np.allclose(U @ U.T, U_ @ U_.T, **tols_decomp)
+    S, U = linalg.eigh(A)
+    assert np.allclose(vecs.T @ vecs, U @ U.T, atol=1e-5, rtol=1e-5)
+    assert np.allclose(vals, S)
 
 
 @testing.parametrize("nrows", [8])
 @testing.parametrize("num_matvecs", [8, 4, 0])
 def test_shapes_as_expected(nrows, num_matvecs):
-    A = np.arange(1.0, 1.0 + nrows**2).reshape((nrows, nrows))
-    A = A + A.T
-
+    eigvals = np.arange(1.0, 1.0 + nrows)
+    A = test_util.symmetric_matrix_from_eigenvalues(eigvals)
     v0 = np.ones((nrows,))
-    v0 /= linalg.vector_norm(v0)
 
     tridiag_sym = decomp.tridiag_sym(num_matvecs, reortho="full")
     alg = eig.eigh_partial(tridiag_sym)
     S, U = alg(lambda v, p: p @ v, v0, A)
     assert S.shape == (num_matvecs,)
-    assert U.shape == (nrows, num_matvecs)
+    assert U.shape == (num_matvecs, nrows)

tests/test_eig/test_svd_partial.py

@@ -4,62 +4,48 @@
 from matfree.backend import linalg, np, testing
 
 
-@testing.fixture()
 @testing.parametrize("nrows", [10])
 @testing.parametrize("ncols", [3])
-@testing.parametrize("num_significant_singular_vals", [3])
-def A(nrows, ncols, num_significant_singular_vals):
-    """Make a positive definite matrix with certain spectrum."""
-    # 'Invent' a spectrum. Use the number of pre-defined eigenvalues.
-    n = min(nrows, ncols)
-    d = np.arange(n) + 10.0
-    d = d.at[num_significant_singular_vals:].set(0.001)
-    return test_util.asymmetric_matrix_from_singular_values(d, nrows=nrows, ncols=ncols)
-
-
-def test_equal_to_linalg_svd(A):
+def test_equal_to_linalg_svd(nrows, ncols):
     """The output of full-depth SVD should be equal (*) to linalg.svd().
 
     (*) Note: The singular values should be identical,
     and the orthogonal matrices should be orthogonal. They are not unique.
     """
-    nrows, ncols = np.shape(A)
-    num_matvecs = min(nrows, ncols)
-
+    d = np.arange(1.0, 1.0 + min(nrows, ncols))
+    A = test_util.asymmetric_matrix_from_singular_values(d, nrows=nrows, ncols=ncols)
     v0 = np.ones((ncols,))
-    v0 /= linalg.vector_norm(v0)
+    num_matvecs = min(nrows, ncols)
 
     bidiag = decomp.bidiag(num_matvecs)
     svd = eig.svd_partial(bidiag)
-    U, S, Vt = svd(lambda v, p: p @ v, v0, A)
+    Ut, S, Vt = svd(lambda v, p: p @ v, v0, A)
 
     U_, S_, Vt_ = linalg.svd(A, full_matrices=False)
 
     tols_decomp = {"atol": 1e-5, "rtol": 1e-5}
-    assert np.allclose(S, S_)  # strict "allclose"
-
-    assert np.shape(U) == np.shape(U_)
-    assert np.shape(Vt) == np.shape(Vt_)
-    assert np.allclose(U @ U.T, U_ @ U_.T, **tols_decomp)
+    assert np.allclose(S, S_, **tols_decomp)
+    assert np.allclose(Ut.T @ Ut, U_ @ U_.T, **tols_decomp)
     assert np.allclose(Vt @ Vt.T, Vt_ @ Vt_.T, **tols_decomp)
 
 
-def test_shapes_as_expected(A):
+@testing.parametrize("nrows", [10])
+@testing.parametrize("ncols", [3])
+@testing.parametrize("num_matvecs", [0, 2, 3])
+def test_shapes_as_expected(nrows, ncols, num_matvecs):
     """The output of full-depth SVD should be equal (*) to linalg.svd().
 
     (*) Note: The singular values should be identical,
     and the orthogonal matrices should be orthogonal. They are not unique.
     """
-    nrows, ncols = np.shape(A)
-    num_matvecs = min(nrows, ncols) // 2
-
+    d = np.arange(1.0, 1.0 + min(nrows, ncols))
+    A = test_util.asymmetric_matrix_from_singular_values(d, nrows=nrows, ncols=ncols)
     v0 = np.ones((ncols,))
-    v0 /= linalg.vector_norm(v0)
 
     bidiag = decomp.bidiag(num_matvecs)
     svd = eig.svd_partial(bidiag)
-    U, S, Vt = svd(lambda v, p: p @ v, v0, A)
+    Ut, S, Vt = svd(lambda v, p: p @ v, v0, A)
 
-    assert U.shape == (nrows, num_matvecs)
+    assert Ut.shape == (num_matvecs, nrows)
     assert S.shape == (num_matvecs,)
     assert Vt.shape == (num_matvecs, ncols)