Add keyword arguments to fixedpoint for hasconverged and shouldstop stopping criteria

src/algorithms/optimization/peps_optimization.jl

@@ -109,7 +109,9 @@ The optimization parameters can be supplied via the keyword arguments or directl
     4. Debug info including AD outputs
 * `reuse_env::Bool=$(Defaults.reuse_env)` : If `true`, the current optimization step is initialized on the previous environment, otherwise a random environment is used.
 * `symmetrization::Union{Nothing,SymmetrizationStyle}=nothing` : Accepts `nothing` or a `SymmetrizationStyle`, in which case the PEPS and PEPS gradient are symmetrized after each optimization iteration.
-* `(finalize!)=OptimKit._finalize!` : Inserts a `finalize!` function call after each optimization step by utilizing the `finalize!` kwarg of `OptimKit.optimize`. The function maps `(peps, env), f, g = finalize!((peps, env), f, g, numiter)`.
+* `hasconverged=OptimKit.DefaultHasConverged(optimizer_alg.gradtol)` : Function specifying the convergence criterion with signature `bool = hasconverged(state, cost, grad, gradnorm)`, see [OptimKit.optimize](https://github.com/Jutho/OptimKit.jl/blob/master/src/OptimKit.jl) for specifics. Note that this overrides the default convergence criterion.
+* `shouldstop=OptimKit.DefaultShouldStop(optimizer_alg.maxiter)` : Function specifying the stopping criterion with signature `bool = shouldstop(state, cost, grad, numfg, iter, timespent)`, see [OptimKit.optimize](https://github.com/Jutho/OptimKit.jl/blob/master/src/OptimKit.jl) for specifics. Note that this overrides the default stopping criterion.
+* `(finalize!)=OptimKit._finalize!` : Inserts a `finalize!` function call after each optimization step by utilizing the `finalize!` kwarg of `OptimKit.optimize`. The function maps `(state, env), f, g = finalize!((state, env), cost, grad, numiter)`.
 
 ### Boundary algorithm
 
@@ -131,7 +133,7 @@ keyword arguments are:
 * `verbosity::Int` : Gradient output verbosity, ≤0 by default to disable too verbose printing. Should only be >0 for debug purposes.
 * `alg::Symbol=:$(Defaults.gradient_alg)` : Implicit gradient algorithm variant, can be one of the following:
     - `:FixedPointGradient` : Compute the gradient via fixed-point differentiation, see [`FixedPointGradient`](@ref)
-* `solver_alg::`Union{Algorithm,NamedTuple}`: Solver algorithm for computing the implicit gradient; see [`FixedPointGradient`](@ref) for supported algorithms.
+* `solver_alg::Union{Algorithm,NamedTuple}`: Solver algorithm for computing the implicit gradient; see [`FixedPointGradient`](@ref) for supported algorithms.
 
 ### Optimizer settings
 
@@ -164,14 +166,16 @@ information `NamedTuple` which contains the following entries:
 * `gradnorms_unitcell` : History of gradient norms for each respective unit cell entry.
 * `times` : History of optimization step execution times.
 """
-function fixedpoint(
-        operator, peps₀::InfinitePEPS, env₀; (finalize!) = OptimKit._finalize!, kwargs...,
-    )
-    alg = select_algorithm(fixedpoint, env₀; kwargs...)
-    return fixedpoint(operator, peps₀, env₀, alg; finalize!)
+function fixedpoint(operator, peps₀::InfinitePEPS, env₀; kwargs...)
+    extra_kwarg_keys = (:finalize!, :hasconverged, :shouldstop) # these will not be passed to `select_algorithm`, only to the 2nd `fixedpoint` call
+    alg = select_algorithm(fixedpoint, env₀; filter(kw -> !(first(kw) in extra_kwarg_keys), kwargs)...)
+    return fixedpoint(operator, peps₀, env₀, alg; filter(kw -> first(kw) in extra_kwarg_keys, kwargs)...)
 end
 function fixedpoint(
-        operator, peps₀::InfinitePEPS, env₀, alg::PEPSOptimize; (finalize!) = OptimKit._finalize!,
+        operator, peps₀::InfinitePEPS, env₀, alg::PEPSOptimize;
+        (finalize!) = OptimKit._finalize!,
+        hasconverged = OptimKit.DefaultHasConverged(alg.optimizer_alg.gradtol),
+        shouldstop = OptimKit.DefaultShouldStop(alg.optimizer_alg.maxiter),
     )
     # validate inputs
     check_input(fixedpoint, peps₀, env₀, alg)
@@ -197,7 +201,8 @@ function fixedpoint(
     # optimize operator cost function
     (peps_final, env_final), cost_final, ∂cost, numfg, convergence_history = optimize(
         (peps₀, env₀), alg.optimizer_alg;
-        retract, inner = real_inner, finalize!, (transport!) = (peps_transport!),
+        retract, inner = real_inner, (transport!) = (peps_transport!),
+        hasconverged, shouldstop, finalize!,
     ) do (peps, env)
         start_time = time_ns()
         E, gs = withgradient(peps) do ψ

test/examples/heisenberg.jl

@@ -150,3 +150,52 @@ end
     @test e_site2 ≈ E_ref atol = 1.0e-2
     @test all(@. ξ_h > 0 && ξ_v > 0)
 end
+
+@kwdef mutable struct ΔEnergyShouldStop
+    E_last::Float64 = 0.0
+    tol::Float64 = 1.0e-12
+    maxiter::Int = 100
+end
+function (es::ΔEnergyShouldStop)(x, f, g, numfg, numiter, t)
+    Δenergy = f - es.E_last
+    es.E_last = f
+    return (abs(Δenergy) <= es.tol) || numiter >= es.maxiter
+end
+
+@kwdef mutable struct ΔEnergyHasConverged
+    E_last::Float64 = 0.0
+    tol::Float64 = 1.0e-12
+end
+function (es::ΔEnergyHasConverged)(x, f, g, normgrad)
+    Δenergy = f - es.E_last
+    es.E_last = f
+    return abs(Δenergy) <= es.tol
+end
+
+@testset "Early stopping with hasconverged and shouldstop" begin
+    maxiter = 100
+    Random.seed!(123)
+    H = heisenberg_XYZ(InfiniteSquare())
+    peps₀ = InfinitePEPS(ComplexSpace(2), ComplexSpace(Dbond))
+
+    # should converge when energy difference becomes small enough
+    env₀, = leading_boundary(CTMRGEnv(peps₀, ComplexSpace(χenv)), peps₀)
+    ΔEconverged = ΔEnergyHasConverged(; tol = 1.0e-4)
+    peps, env, E, info = fixedpoint(
+        H, peps₀, env₀;
+        optimizer_alg = (; maxiter, tol = 0),
+        hasconverged = ΔEconverged,
+    )
+    @test length(info.costs) < maxiter
+
+    # should stop when linesearching fails and returns α = 0 (i.e. ΔE = 0)
+    χenvsmall = 2
+    env₀, = leading_boundary(CTMRGEnv(peps₀, ComplexSpace(χenvsmall)), peps₀)
+    ΔEstop = ΔEnergyShouldStop(; tol = 1.0e-8)
+    peps, env, E, info = fixedpoint(
+        H, peps₀, env₀;
+        optimizer_alg = (; maxiter, ls_maxfg = 2, ls_maxiter = 2),
+        shouldstop = ΔEstop,
+    )
+    @test length(info.costs) < maxiter
+end