Add support for deflation in arnoldi_decomposition

Makefile

@@ -0,0 +1,9 @@
+.PHONY: lint tests
+
+lint:
+	ruff check src tests
+
+tests:
+	uv run pytest tests -s
+
+tests-all: lint tests

src/arnoldi/decomposition.py

@@ -34,7 +34,9 @@ def initialize(self, init_vector=None):
         self.V[:, 0] = init_vector
 
     def iterate(self, A):
-        _, _, m = arnoldi_decomposition(A, self.V, self.H, self._atol, self.m)
+        _, _, m = arnoldi_decomposition(A, self.V, self.H, self._atol)
+        if m != self.m:
+            raise ValueError("Lucky break not supported yet")
         return m
 
 
@@ -44,7 +46,7 @@ def _largest_eigvals(H, n_ev):
     return eigvals[ind], eigvecs[:, ind]
 
 
-def arnoldi_decomposition(A, V, H, invariant_tol, max_dim=None):
+def arnoldi_decomposition(A, V, H, invariant_tol=None, *, start_dim=0, max_dim=None):
     """Run the arnoldi decomposition for square matrix a of dimension n.
 
     Parameters
@@ -71,17 +73,23 @@ def arnoldi_decomposition(A, V, H, invariant_tol, max_dim=None):
         max_dim in case a "lucky break" is found, i.e. the Krylov basis
         invariant is lower dimension than max_dim
     """
+    # Logic of sqrt copied from Julia's ArnoldiMethod.jl package
+    if invariant_tol is None:
+        invariant_tol = np.sqrt(np.finfo(A.dtype).eps)
+
     n = A.shape[0]
     m = V.shape[1] - 1
 
     assert A.shape[1] == n, "A is expected to be square matrix"
     assert V.shape == (n, m+1), "V must have the same number of rows as A"
     assert H.shape == (m+1, m), f"H must be {m+1, m}, is {H.shape}"
 
-    max_dim = max_dim or m
+    if max_dim is None:
+        max_dim = m
+
     assert max_dim <= m, "max_dim > m violated"
 
-    for j in range(max_dim):
+    for j in range(start_dim, max_dim):
         v = V[:, j+1]
         v[:] = A @ V[:, j]
 
@@ -93,7 +101,8 @@ def arnoldi_decomposition(A, V, H, invariant_tol, max_dim=None):
         H[j + 1, j] = norm(v)
 
         if H[j + 1, j] < invariant_tol:
-            raise ValueError("Lucky break not supported yet")
+            max_dim = j + 1
+            return V[:, :max_dim+1], H[:max_dim+1, :max_dim], max_dim
         v /= H[j + 1, j]
 
     return V[:, :max_dim+1], H[:max_dim+1, :max_dim], max_dim

src/arnoldi/explicit_restarts.py

@@ -0,0 +1,31 @@
+import numpy as np
+
+
+from .decomposition import RitzDecomposition, arnoldi_decomposition
+from .utils import rand_normalized_vector
+
+
+def naive_explicit_restarts(A, m=None, *, stopping_criterion=None, max_restarts=10):
+    tol = stopping_criterion or np.sqrt(np.finfo(A.dtype).eps)
+    dtype = np.promote_types(A.dtype, np.complex64)
+
+    n = A.shape[0]
+    k = 1  # Naive arnoldi w/o restart only really works for 1 eigenvalue
+    m = m or min(max(2 * k + 1, 20), n)
+
+    V = np.zeros((n, m+1), dtype)
+    H = np.zeros((m+1, m), dtype)
+
+    v0 = rand_normalized_vector(n).astype(dtype)
+    for i in range(max_restarts):
+        V[:, 0] = v0
+        V, H, n_iter = arnoldi_decomposition(A, V, H)
+        ritz = RitzDecomposition.from_v_and_h(V, H, k)
+        if ritz.approximate_residuals[0] < tol:
+            residuals = ritz.compute_true_residuals(A)
+            if residuals[0] / max(ritz.values[0], tol) < tol:
+                return ritz, True, i
+        # FIXME: should take the ritz vector w/ the lowest residual
+        v0 = ritz.vectors[:, 0]
+
+    return ritz, False, max_restarts

tests/common.py

@@ -0,0 +1,3 @@
+# Max retries for short tests
+MAX_RETRIES_SHORT = 3
+

tests/test_decomposition.py

@@ -7,11 +7,12 @@
 from arnoldi.matrices import mark, laplace
 from arnoldi.utils import rand_normalized_vector
 
+from .common import MAX_RETRIES_SHORT
+
+
 
 ATOL = 1e-8
 RTOL = 1e-4
-# Max retries for short tests
-MAX_RETRIES_SHORT = 3
 
 norm = np.linalg.norm
 
@@ -25,7 +26,7 @@ def inject_noise(A):
 
 
 def basis_vector(n, k, dtype=np.int64):
-    """Create the basis vector e_k in R^n, aka e_k is (n,), and with e_k[k] =
+    """Create the basis vector e_k in R^n, aka e_k is (n,), and with e_k[k-1] =
     1
     """
     ret = np.zeros(n, dtype=dtype)
@@ -54,7 +55,7 @@ def assert_invariants(A, V, H, m):
 
     # the arnoldi basis V is orthonormal
     np.testing.assert_allclose(
-        V.conj().T @ V, np.eye(m + 1), rtol=RTOL, atol=ATOL
+        V_m.conj().T @ V_m, np.eye(m), rtol=RTOL, atol=ATOL
     )
 
     # the arnoldi decomposition invariants are respected
@@ -171,6 +172,60 @@ def test_invariant_simple(self):
         ## Then
         assert_invariants(A, Va, Ha, n_iter)
 
+    def test_stop_restart_mid_point(self):
+        # Ensure stopping mid point and restarting keeps the decomposition
+        # invariants
+
+        ## Given
+        n = 10
+        m = 6
+        mid_point = 3
+        dtype = np.complex128
+
+        A = sp.random(n, n, density=5 / n, dtype=dtype)
+        A += sp.diags_array(np.ones(n))
+
+        V = np.zeros((n, m+1), dtype=dtype)
+        H = np.zeros((m+1, m), dtype=dtype)
+
+        V[:, 0] = rand_normalized_vector(n, dtype)
+
+        ## When
+        Va, Ha, n_iter = arnoldi_decomposition(A, V, H, ATOL, max_dim=mid_point)
+
+        ## Then
+        assert n_iter == mid_point
+        assert Va.shape == (n, mid_point+1)
+        assert Ha.shape == (mid_point+1, mid_point)
+        assert_invariants(A, Va, Ha, n_iter)
+
+        ## When
+        Va, Ha, n_iter = arnoldi_decomposition(A, V, H, ATOL, start_dim=mid_point)
+
+        ## Then
+        assert n_iter == m
+        assert Va.shape == (n, m+1)
+        assert Ha.shape == (m+1, m)
+        assert_invariants(A, Va, Ha, n_iter)
+
+        ## When
+        # Filling H/V with sentinel values
+        V[:, :mid_point] = 42
+        H[:, :mid_point] = 42
+
+        Va, Ha, n_iter = arnoldi_decomposition(A, V, H, ATOL, start_dim=mid_point)
+
+        ## Then
+        assert n_iter == m
+        assert Va.shape == (n, m+1)
+        assert Ha.shape == (m+1, m)
+        # Ensure initial columns are untouched
+        np.testing.assert_array_almost_equal(V[:, 0], 42)
+        np.testing.assert_array_almost_equal(H[:, 0], 42)
+        # Invariants should fail in this case
+        with pytest.raises(AssertionError):
+            assert_invariants(A, Va, Ha, n_iter)
+
     def test_max_dim_support(self):
         ## Given
         n = 10
@@ -187,13 +242,79 @@ def test_max_dim_support(self):
         V[:, 0] = rand_normalized_vector(n, dtype)
 
         ## When
-        Va, Ha, n_iter = arnoldi_decomposition(A, V, H, ATOL, max_dim)
+        Va, Ha, n_iter = arnoldi_decomposition(A, V, H, ATOL, max_dim=max_dim)
 
         ## Then
         assert Va.shape == (n, max_dim+1)
         assert Ha.shape == (max_dim+1, max_dim)
         assert_invariants(A, Va, Ha, n_iter)
 
+    def test_converge_first_iteration(self):
+        ## Given
+        n = 10
+        m = 6
+        dtype = np.complex128
+
+        A = sp.random(n, n, density=5 / n, dtype=dtype)
+        A += sp.diags_array(np.ones(n))
+
+        ## When
+        # the Arnoldi decomposition is initialized with an eigenvector
+        _, r_vecs = sp.linalg.eigs(A)
+
+        V = np.zeros((n, m+1), dtype=dtype)
+        H = np.zeros((m+1, m), dtype=dtype)
+        V[:, 0] = r_vecs[:, 0]
+
+        Vm, Hm, n_iter = arnoldi_decomposition(A, V, H, ATOL)
+
+        ## Then
+        # Only one iteration is expected
+        assert n_iter == 1
+        assert Vm.shape == (n, n_iter+1)
+        assert Hm.shape == (n_iter+1, 1)
+        assert_invariants(A, Vm, Hm, n_iter)
+
+    def test_start_dim_invariant(self):
+        # Ensure starting from k-1 eigenpairs works with start_dim to find
+        # approximation of k-th eigenvalue.
+
+        ## Given
+        n = 10
+        m = n-1
+        k = 2
+        dtype = np.complex128
+
+        A = sp.random(n, n, density=5 / n, dtype=dtype)
+        A += sp.diags_array(np.ones(n))
+        r_vals, r_vecs = sp.linalg.eigs(A, k)
+
+        ## When
+        # the Arnoldi decomposition is initialized with k-1 eigenvector
+        V = np.zeros((n, m+1), dtype=dtype)
+        H = np.zeros((m+1, m), dtype=dtype)
+
+        V[:, :k-1] = r_vecs[:, :k-1]
+        for i in range(k-1):
+            H[i, i] = r_vals[i]
+        V[:, k-1] = rand_normalized_vector(n, dtype)
+
+        # Orthonormalize
+        V[:, :k], _ = np.linalg.qr(V[:, :k], mode="reduced")
+        V_init = V[:, :k].copy()
+        H_init = H[:, :k-1].copy()
+
+        Va, Ha, n_iter = arnoldi_decomposition(A, V, H, ATOL, start_dim=k-1)
+
+        ## Then
+        assert Va.shape == (n, n_iter+1)
+        assert Ha.shape == (n_iter+1, n_iter)
+        assert_invariants(A, Va, Ha, n_iter)
+
+        # Ensure initialized arrays are not modified
+        np.testing.assert_equal(Va[:, :k], V_init)
+        np.testing.assert_equal(Ha[:, :k-1], H_init)
+
 
 class TestEigenValues:
     @pytest.mark.parametrize(

tests/test_explicit_restarts.py

@@ -0,0 +1,42 @@
+import pytest
+
+from arnoldi.explicit_restarts import naive_explicit_restarts
+from arnoldi.matrices import mark
+
+from .common import MAX_RETRIES_SHORT
+
+
+class TestNaiveExplicitRestarts:
+    @pytest.mark.parametrize(
+        "restarts, digits", [(1, 0), (2, 1), (3, 3), (4, 5), (5, 6)]
+    )
+    @pytest.mark.flaky(reruns=MAX_RETRIES_SHORT)
+    def test_mark10(self, restarts, digits):
+        # For the numerical value, see table 6.2 of Numerical Methods for Large
+        # Eigenvalue Problems, 2nd edition.
+
+        ## Given
+        A = mark(10)
+        m = 10
+
+        ## When
+        ritz, *_ = naive_explicit_restarts(A, m, max_restarts=restarts)
+
+        ## Then
+        assert ritz.compute_true_residuals(A) <= 2 * 10**(-digits)
+
+    @pytest.mark.flaky(reruns=MAX_RETRIES_SHORT)
+    def test_convergence(self):
+        ## Given
+        A = mark(10)
+        m = 20
+        atol = 1e-6
+
+        ## When
+        ritz, has_converged, *_ = naive_explicit_restarts(A, m,
+                                                          max_restarts=200,
+                                                          stopping_criterion=atol)
+
+        ## Then
+        assert ritz.compute_true_residuals(A) <= atol
+        assert has_converged