Implement eig.eigh_partial and allow parametric matvecs in partial SVD & eig

matfree/backend/linalg.py

@@ -20,6 +20,10 @@ def eigh(x, /):
     return jnp.linalg.eigh(x)
 
 
+def eig(x, /):
+    return jnp.linalg.eig(x)
+
+
 def cholesky(x, /):
     return jnp.linalg.cholesky(x)
 

matfree/eig.py

@@ -20,9 +20,9 @@ def svd_partial(bidiag: Callable):
 
     """
 
-    def svd(Av: Callable, v0: Array):
+    def svd(Av: Callable, v0: Array, *parameters):
         # Factorise the matrix
-        (u, v), B, *_ = bidiag(Av, v0)
+        (u, v), B, *_ = bidiag(Av, v0, *parameters)
 
         # Compute SVD of factorisation
         U, S, Vt = linalg.svd(B, full_matrices=False)
@@ -31,3 +31,30 @@ def svd(Av: Callable, v0: Array):
         return u @ U, S, Vt @ v.T
 
     return svd
+
+
+def eigh_partial(tridiag_sym: Callable):
+    """Partial symmetric/Hermitian eigenvalue decomposition.
+
+    Combines tridiagonalization with a decomposition
+    of the (small) tridiagonal matrix.
+
+    Parameters
+    ----------
+    tridiag_sym:
+        An implementation of tridiagonalization.
+        For example, the output of
+        [decomp.tridiag_sym][matfree.decomp.tridiag_sym].
+
+    """
+
+    def eigh(Av: Callable, v0: Array, *parameters):
+        # Factorise the matrix
+        Q, H, *_ = tridiag_sym(Av, v0, *parameters)
+
+        # Compute SVD of factorisation
+        vals, vecs = linalg.eigh(H)
+        vecs = Q @ vecs
+        return vals, vecs
+
+    return eigh

tests/test_eig/test_eigh_partial.py

@@ -0,0 +1,38 @@
+"""Tests for eigenvalue functionality."""
+
+from matfree import decomp, eig, test_util
+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)
+
+    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)
+
+    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)

tests/test_eig/test_svd_partial.py

@@ -26,15 +26,12 @@ def test_equal_to_linalg_svd(A):
     nrows, ncols = np.shape(A)
     num_matvecs = min(nrows, ncols)
 
-    def Av(v):
-        return A @ v
-
     v0 = np.ones((ncols,))
     v0 /= linalg.vector_norm(v0)
 
     bidiag = decomp.bidiag(num_matvecs)
     svd = eig.svd_partial(bidiag)
-    U, S, Vt = svd(Av, v0)
+    U, S, Vt = svd(lambda v, p: p @ v, v0, A)
 
     U_, S_, Vt_ = linalg.svd(A, full_matrices=False)