From fbd467d12dc3a0d8ecc0d5ebdb89e00f68fbb63d Mon Sep 17 00:00:00 2001 From: Juan Jose Garcia-Ripoll Date: Sun, 21 Jan 2024 18:20:42 +0100 Subject: [PATCH] Simplified and reorganized gradient descent and Arnoldi --- src/seemps/optimization/arnoldi.py | 114 +++++++++++------- src/seemps/optimization/descent.py | 11 +- tests/test_optimization/test_arnoldi.py | 35 +++++- .../test_gradient_descent.py | 5 +- 4 files changed, 111 insertions(+), 54 deletions(-) diff --git a/src/seemps/optimization/arnoldi.py b/src/seemps/optimization/arnoldi.py index d8e2d07a..f1129704 100644 --- a/src/seemps/optimization/arnoldi.py +++ b/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] = # H[0,1] = / sqrt() = sqrt() - 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 diff --git a/src/seemps/optimization/descent.py b/src/seemps/optimization/descent.py index 708206ff..26b847ff 100644 --- a/src/seemps/optimization/descent.py +++ b/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 ------- diff --git a/tests/test_optimization/test_arnoldi.py b/tests/test_optimization/test_arnoldi.py index a49498d5..59358a8f 100644 --- a/tests/test_optimization/test_arnoldi.py +++ b/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)) diff --git a/tests/test_optimization/test_gradient_descent.py b/tests/test_optimization/test_gradient_descent.py index fe9ac4c2..de9cb66a 100644 --- a/tests/test_optimization/test_gradient_descent.py +++ b/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