Remove pinv and implement matfree.eig (#198)

matfree/bounds.py

@@ -1,4 +1,4 @@
-"""Bounds on functions of matrices."""
+"""Matrix-free bounds on functions of matrices."""
 
 from matfree.backend import linalg, np
 

matfree/decomp.py

@@ -12,45 +12,6 @@
 from matfree.backend import containers, control_flow, func, linalg, np, tree_util
 from matfree.backend.typing import Array, Callable, Tuple
 
-# todo: rename svd_approx to svd_partial() because the algorithm is called
-#  "Partial SVD", not "Approximate SVD".
-# todo: move svd_approx to a dedicated eigenvalue-module?
-
-
-def svd_approx(
-    v0: Array, depth: int, Av: Callable, vA: Callable, matrix_shape: Tuple[int, ...]
-):
-    """Approximate singular value decomposition.
-
-    Uses GKL with full reorthogonalisation to bi-diagonalise the target matrix
-    and computes the full SVD of the (small) bidiagonal matrix.
-
-    Parameters
-    ----------
-    v0:
-        Initial vector for Golub-Kahan-Lanczos bidiagonalisation.
-    depth:
-        Depth of the Krylov space constructed by Golub-Kahan-Lanczos bidiagonalisation.
-        Choosing `depth = min(nrows, ncols) - 1` would yield behaviour similar to
-        e.g. `np.linalg.svd`.
-    Av:
-        Matrix-vector product function.
-    vA:
-        Vector-matrix product function.
-    matrix_shape:
-        Shape of the matrix involved in matrix-vector and vector-matrix products.
-    """
-    # Factorise the matrix
-    algorithm = bidiag(Av, vA, depth, matrix_shape=matrix_shape)
-    u, (d, e), vt, _ = algorithm(v0)
-
-    # Compute SVD of factorisation
-    B = _bidiagonal_dense(d, e)
-    U, S, Vt = linalg.svd(B, full_matrices=False)
-
-    # Combine orthogonal transformations
-    return u @ U, S, Vt @ vt
-
 
 class _LanczosAlg(containers.NamedTuple):
     """Lanczos decomposition algorithm."""
@@ -670,21 +631,3 @@ def body_fun(_, s):
 
     result = control_flow.fori_loop(lower, upper, body_fun=body_fun, init_val=init_val)
     return extract(result)
-
-
-def _bidiagonal_dense(d, e):
-    diag = linalg.diagonal_matrix(d)
-    offdiag = linalg.diagonal_matrix(e, 1)
-    return diag + offdiag
-
-
-def _eigh_tridiag_sym(diag, off_diag):
-    # todo: once jax supports eigh_tridiagonal(eigvals_only=False),
-    #  use it here. Until then: an eigen-decomposition of size (order + 1)
-    #  does not hurt too much...
-    diag = linalg.diagonal_matrix(diag)
-    offdiag1 = linalg.diagonal_matrix(off_diag, -1)
-    offdiag2 = linalg.diagonal_matrix(off_diag, 1)
-    dense_matrix = diag + offdiag1 + offdiag2
-    eigvals, eigvecs = linalg.eigh(dense_matrix)
-    return eigvals, eigvecs

matfree/eig.py

@@ -0,0 +1,46 @@
+"""Matrix-free eigenvalue and singular-value analysis."""
+
+from matfree import decomp
+from matfree.backend import linalg
+from matfree.backend.typing import Array, Callable, Tuple
+
+
+def svd_partial(
+    v0: Array, depth: int, Av: Callable, vA: Callable, matrix_shape: Tuple[int, ...]
+):
+    """Partial singular value decomposition.
+
+    Combines bidiagonalisation with full reorthogonalisation
+    and computes the full SVD of the (small) bidiagonal matrix.
+
+    Parameters
+    ----------
+    v0:
+        Initial vector for Golub-Kahan-Lanczos bidiagonalisation.
+    depth:
+        Depth of the Krylov space constructed by Golub-Kahan-Lanczos bidiagonalisation.
+        Choosing `depth = min(nrows, ncols) - 1` would yield behaviour similar to
+        e.g. `np.linalg.svd`.
+    Av:
+        Matrix-vector product function.
+    vA:
+        Vector-matrix product function.
+    matrix_shape:
+        Shape of the matrix involved in matrix-vector and vector-matrix products.
+    """
+    # Factorise the matrix
+    algorithm = decomp.bidiag(Av, vA, depth, matrix_shape=matrix_shape)
+    u, (d, e), vt, _ = algorithm(v0)
+
+    # Compute SVD of factorisation
+    B = _bidiagonal_dense(d, e)
+    U, S, Vt = linalg.svd(B, full_matrices=False)
+
+    # Combine orthogonal transformations
+    return u @ U, S, Vt @ vt
+
+
+def _bidiagonal_dense(d, e):
+    diag = linalg.diagonal_matrix(d)
+    offdiag = linalg.diagonal_matrix(e, 1)
+    return diag + offdiag

matfree/funm.py

@@ -1,5 +1,4 @@
-"""Functions of matrices implemented as matrix-function-vector products.
-
+"""Matrix-free implementations of functions of matrices.
 
 Examples
 --------

tests/test_decomp/test_svd_approx.py → tests/test_eig/test_svd_partial.py

@@ -1,6 +1,6 @@
 """Tests for SVD functionality."""
 
-from matfree import decomp, test_util
+from matfree import eig, test_util
 from matfree.backend import linalg, np, testing
 
 
@@ -34,7 +34,7 @@ def vA(v):
 
     v0 = np.ones((ncols,))
     v0 /= linalg.vector_norm(v0)
-    U, S, Vt = decomp.svd_approx(v0, depth, Av, vA, matrix_shape=np.shape(A))
+    U, S, Vt = eig.svd_partial(v0, depth, Av, vA, matrix_shape=np.shape(A))
     U_, S_, Vt_ = linalg.svd(A, full_matrices=False)
 
     tols_decomp = {"atol": 1e-5, "rtol": 1e-5}