Simplified and reorganized gradient descent and Arnoldi

src/seemps/optimization/arnoldi.py

@@ -1,11 +1,12 @@
-from typing import Callable, Optional, Union
+from __future__ import annotations
+from typing import Callable, Optional, Union, Any
 
 import numpy as np
 import scipy.linalg  # type: ignore
 from numpy.typing import NDArray
 
 from ..expectation import scprod
-from ..mpo import MPO
+from ..operators import MPO, MPOSum, MPOList
 from ..state import MPS, CanonicalMPS, MPSSum, random_mps
 from ..truncate.simplify import simplify
 from .descent import DESCENT_STRATEGY, OptimizeResults, Strategy
@@ -18,13 +19,17 @@ def _ill_conditioned_norm_matrix(N, tol=np.finfo(float).eps * 10):
 
 class MPSArnoldiRepresentation:
     empty: NDArray = np.zeros((0, 0))
-    operator: MPO
+    operator: Union[MPO, MPOSum]
     H: NDArray
     N: NDArray
     V: list[MPS]
     strategy: Strategy
 
-    def __init__(self, operator: MPO, strategy: Strategy = DESCENT_STRATEGY):
+    def __init__(
+        self,
+        operator: Union[MPO, MPOSum],
+        strategy: Strategy = DESCENT_STRATEGY,
+    ):
         self.operator = operator
         self.H = self.empty
         self.N = self.empty
@@ -71,13 +76,13 @@ def variance_estimate(self) -> float:
         # This means
         # H[0,0] = <v|H|v>
         # H[0,1] = <v|H|Hv> / sqrt(<Hv|Hv>) = sqrt(<Hv|Hv>)
-        return np.abs(self.H[0, 1]) ** 2 - np.abs(self.H[0, 0]) ** 2
+        return abs(np.abs(self.H[0, 1]) ** 2 - np.abs(self.H[0, 0]) ** 2)
 
     def exponential(self, factor: Union[complex, float]) -> MPS:
         w = np.zeros(len(self.V))
         w[0] = 1.0
         w = scipy.sparse.linalg.expm_multiply(factor * self.H, w)
-        return simplify(MPSSum(w, self.V), strategy=self.strategy)
+        return simplify(MPSSum(w.tolist(), self.V), strategy=self.strategy)
 
     def build_Krylov_basis(self, v: MPS, order: int) -> bool:
         """Build a Krylov basis up to given order. Returns False
@@ -94,67 +99,88 @@ def build_Krylov_basis(self, v: MPS, order: int) -> bool:
 
 
 def arnoldi_eigh(
-    operator: MPO,
+    operator: Union[MPO, MPOSum],
     v0: Optional[MPS] = None,
-    maxiter: int = 100,
     nvectors: int = 10,
+    maxiter: int = 100,
     tol: float = 1e-13,
     strategy: Strategy = DESCENT_STRATEGY,
     miniter: int = 1,
-    callback: Optional[Callable] = None,
+    callback: Optional[Callable[[MPS, float, OptimizeResults], Any]] = None,
 ) -> OptimizeResults:
+    """Ground state search of Hamiltonian `H` by Arnoldi diagonalization.
+
+    Parameters
+    ----------
+    H : Union[MPO, MPOSum]
+        Hamiltonian in MPO form.
+    state : MPS
+        Initial guess of the ground state.
+    nvectors : int
+        Number of vectors in the Arnoldi algorithm (defaults to 10).
+        It must be an integer between 2 and 100, both included.
+    maxiter : int
+        Maximum number of iterations (defaults to 1000).
+    tol : float
+        Energy variation that indicates termination (defaults to 1e-13).
+    strategy : Optional[Strategy]
+        Linear combination of MPS truncation strategy. Defaults to
+        DESCENT_STRATEGY.
+    callback : Optional[Callable[[MPS, float, OptimizeResults], Any]]
+        A callable called after each iteration with the current state,
+        an estimate of the energy, and the accumulated results object.
+        Defaults to None.
+
+    Results
+    -------
+    OptimizeResults
+        Results from the optimization. See :class:`OptimizeResults`.
+    """
     if v0 is None:
         v0 = random_mps(operator.dimensions(), D=2)
+    if nvectors < 2 or nvectors > 100 or not isinstance(nvectors, int):
+        raise ValueError("nvectors must be an integer between 2 and 100")
     arnoldi = MPSArnoldiRepresentation(operator, strategy)
     arnoldi.add_vector(v0)
     v: MPS = operator @ v0  # type: ignore
-    best_energy = arnoldi.H[0, 0].real
+    energy = arnoldi.H[0, 0].real
+    variance = abs(scprod(v, v) - energy * energy)
+    results = OptimizeResults(
+        state=v0,
+        energy=energy,
+        trajectory=[energy],
+        variances=[variance],
+        message=f"Exceeded maximum number of steps {maxiter}",
+        converged=False,
+    )
     if callback is not None:
-        callback(arnoldi.V[0])
-    variance = abs(scprod(v, v)) - best_energy * best_energy
-    best_vector = v0
-    energies: list[float] = [best_energy]
-    variances: list[float] = [variance]
-    last_eigenvalue = variance = np.Inf
-    message = f"Exceeded maximum number of steps {maxiter}"
-    converged = True
+        callback(arnoldi.V[0], energy, results)
+    last_eigenvalue = np.Inf
     for i in range(maxiter):
         v, success = arnoldi.add_vector(v)
         if not success and nvectors == 2:
-            message = "Unable to construct Arnoldi matrix"
-            converged = False
+            results.message = "Unable to construct Arnoldi matrix"
+            results.converged = False
             break
         L = len(arnoldi.V)
-        if L == 2 and i > 0:
-            if L > 1:
-                variance = arnoldi.variance_estimate()
-            variances.append(variance)
-        elif L != 2 and i > 0:
-            variances.append(variances[-1])
+        if L > 1:
+            variance = arnoldi.variance_estimate()
+        results.variances.append(variance)
         if L == nvectors or not success:
             v, eigenvalue = arnoldi.restart_with_ground_state()
             eigenvalue_change, last_eigenvalue = (
                 eigenvalue - last_eigenvalue,
                 eigenvalue,
             )
-            if (
-                eigenvalue_change >= abs(tol) or eigenvalue_change >= -abs(tol)
-            ) and i > miniter:
-                message = f"Eigenvalue converged within tolerance {tol}"
-                converged = True
+            if (eigenvalue_change >= -abs(tol)) and i > miniter:
+                results.message = f"Eigenvalue converged within tolerance {tol}"
+                results.converged = True
                 break
-        v = operator @ v  # type: ignore
         energy = arnoldi.H[0, 0].real
+        results.trajectory.append(energy)
+        if energy < results.energy:
+            results.energy, results.state = energy, arnoldi.V[0]
         if callback is not None:
-            callback(arnoldi.V[0])
-        energies.append(energy)
-        if energy < best_energy:
-            best_energy, best_vector = energy, arnoldi.V[0]
-    return OptimizeResults(
-        state=best_vector,
-        energy=best_energy,
-        converged=converged,
-        message=message,
-        trajectory=energies,
-        variances=variances,
-    )
+            callback(arnoldi.V[0], energy, results)
+        v = operator @ v  # type: ignore
+    return results

src/seemps/optimization/descent.py

@@ -1,5 +1,6 @@
+from __future__ import annotations
 import dataclasses
-from typing import Callable, Union
+from typing import Callable, Union, Any
 
 import numpy as np
 
@@ -57,7 +58,7 @@ def gradient_descent(
     tol: float = 1e-13,
     tol_variance: float = 1e-14,
     strategy: Strategy = DESCENT_STRATEGY,
-    callback: Optional[Callable] = None,
+    callback: Optional[Callable[[MPS, float, OptimizeResults], Any]] = None,
 ) -> OptimizeResults:
     """Ground state search of Hamiltonian `H` by gradient descent.
 
@@ -76,8 +77,10 @@ def gradient_descent(
     strategy : Optional[Strategy]
         Linear combination of MPS truncation strategy. Defaults to
         DESCENT_STRATEGY.
-    callback : Optional[callable]
-        A callable called after each iteration (defaults to None).
+    callback : Optional[Callable[[MPS, float, OptimizeResult], Any]]
+        A callable called after each iteration with the current state,
+        an estimate of the energy, and the accumulated results object.
+        Defaults to None.
 
     Results
     -------

tests/test_optimization/test_arnoldi.py

@@ -20,11 +20,40 @@ def make_local_Sz_mpo(self, size: int) -> MPO:
         tensors[-1] = tensors[-1][:, :, :, [1]]
         return MPO(tensors)
 
-    def test_arnoldi_eigh_with_local_field(self):
+    def call_arnoldi(self, *args, **kwdargs):
         N = 4
         H = self.make_local_Sz_mpo(N)
         guess = product_state(np.asarray([1, 1]) / np.sqrt(2.0), N)
         exact = product_state([0, 1], N)
-        result = arnoldi_eigh(H, guess, maxiter=20, tol=1e-15)
-        self.assertAlmostEqual(result.energy, H.expectation(exact))
+        result = arnoldi_eigh(H, guess, *args, **kwdargs)
+        exact_energy = H.expectation(exact)
+        return result, exact, exact_energy
+
+    def test_arnoldi_solves_small_problem(self):
+        result, exact, exact_energy = self.call_arnoldi(
+            nvectors=10, maxiter=20, tol=1e-15
+        )
+        self.assertAlmostEqual(result.energy, exact_energy)
         self.assertSimilarStates(result.state, exact, atol=1e-7)
+
+    def test_arnoldi_stops_if_tolerance_reached(self):
+        result, exact, exact_energy = self.call_arnoldi(
+            nvectors=10, maxiter=300, tol=1e-15
+        )
+        self.assertTrue(len(result.trajectory) < 300)
+
+    def test_arnoldi_invokes_callback(self):
+        callback_calls = 0
+
+        def callback(state, E, results):
+            nonlocal callback_calls
+            callback_calls += 1
+
+        result, *_ = self.call_arnoldi(maxiter=20, tol=1e-15, callback=callback)
+        self.assertTrue(callback_calls > 0)
+        self.assertEqual(callback_calls, len(result.trajectory))
+
+    def test_arnoldi_signals_non_convergence(self):
+        result, _, exact_energy = self.call_arnoldi(nvectors=2, maxiter=3, tol=1e-15)
+        self.assertFalse(result.converged)
+        self.assertFalse(np.isclose(result.energy, exact_energy))

tests/test_optimization/test_gradient_descent.py

@@ -2,7 +2,7 @@
 
 from seemps import MPO, product_state
 from seemps.hamiltonians import HeisenbergHamiltonian
-from seemps.optimization.descent import gradient_descent
+from seemps.optimization.descent import gradient_descent, OptimizeResults
 
 from ..tools import *
 
@@ -38,14 +38,13 @@ def test_gradient_descent_acknowledges_tolerance(self):
             random_uniform_mps(2, N, rng=self.rng), center=0, normalize=True
         )
         result = gradient_descent(H, guess, tol=tol, maxiter=1000)
-        print(result)
         self.assertTrue(result.converged)
         self.assertTrue(abs(result.trajectory[-1] - result.trajectory[-2]) < tol)
 
     def callback(self):
         norms = []
 
-        def callback_func(state: MPS):
+        def callback_func(state: MPS, energy: float, results: OptimizeResults):
             norms.append(np.sqrt(state.norm_squared()))
             return None