diff --git a/src/algorithms/optimization/peps_optimization.jl b/src/algorithms/optimization/peps_optimization.jl index 5bf49e8bb..4100611bb 100644 --- a/src/algorithms/optimization/peps_optimization.jl +++ b/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 ψ diff --git a/test/examples/heisenberg.jl b/test/examples/heisenberg.jl index 6cb797b2b..2968e09c1 100644 --- a/test/examples/heisenberg.jl +++ b/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