Add local preconditioner

docs/src/assets/pepskit.bib

@@ -157,3 +157,18 @@ @misc{zhang_accelerating_2025
       primaryClass={cond-mat.str-el},
       url={https://arxiv.org/abs/2505.00494}, 
 }
+
+@article{zhang_accelerating_2026,
+  title = {Accelerating two-dimensional tensor network optimization by preconditioning},
+  author = {Zhang, Xing-Yu and Yang, Qi and Corboz, Philippe and Haegeman, Jutho and Tang, Wei},
+  journal = {Phys. Rev. B},
+  volume = {113},
+  issue = {12},
+  pages = {125111},
+  numpages = {8},
+  year = {2026},
+  month = {Mar},
+  publisher = {American Physical Society},
+  doi = {10.1103/h396-yc28},
+  url = {https://link.aps.org/doi/10.1103/h396-yc28}
+}

src/Defaults.jl

@@ -12,6 +12,10 @@ Module containing default algorithm parameter values and arguments.
     - `:SimultaneousCTMRG` : Simultaneous expansion and renormalization of all sides.
     - `:SequentialCTMRG` : Sequential application of left moves and rotations.
 * `ctmrg_verbosity=$(Defaults.ctmrg_verbosity)` : CTMRG output information verbosity
+* `ctmrg_dynamic_tols=$(Defaults.ctmrg_dynamic_tols)` : If `true`, wrap the CTMRG algorithm used during variational optimization in an `MPSKit.DynamicTols.DynamicTol` that rescales its tolerance based on the current PEPS optimization gradient norm, see [`PEPSKit.PEPSOptimize`](@ref).
+* `ctmrg_tol_min=$(Defaults.ctmrg_tol_min)` : Minimal CTMRG tolerance used by `ctmrg_dynamic_tols`.
+* `ctmrg_tol_max=$(Defaults.ctmrg_tol_max)` : Maximal CTMRG tolerance used by `ctmrg_dynamic_tols`.
+* `ctmrg_tol_factor=$(Defaults.ctmrg_tol_factor)` : Tolerance scaling factor used by `ctmrg_dynamic_tols`.
 
 ## SVD forward & reverse
 
@@ -85,6 +89,24 @@ Module containing default algorithm parameter values and arguments.
     - `:GeomSum` : Geometric sum approximation of the Neumann series of the inverse Jacobian, see [`PEPSKit.GeomSum`](@ref) for details
     - `:ManualIter` : Manual fixed-point iteration, see [`PEPSKit.ManualIter`](@ref) for details
 * `gradient_fixedpoint_solver_eager=$(Defaults.gradient_fixedpoint_solver_eager)` : Enables `:Arnoldi` solver algorithm to finish before the full Krylov dimension is reached.
+* `gradient_dynamic_tols=$(Defaults.gradient_dynamic_tols)` : If `true`, wrap the gradient algorithm used during variational optimization in an `MPSKit.DynamicTols.DynamicTol` that rescales its tolerance based on the effective (possibly dynamically-scaled) tolerance of the boundary algorithm, see [`PEPSKit.PEPSOptimize`](@ref).
+* `gradient_tol_min=$(Defaults.gradient_tol_min)` : Minimal gradient algorithm tolerance used by `gradient_dynamic_tols`.
+* `gradient_tol_max=$(Defaults.gradient_tol_max)` : Maximal gradient algorithm tolerance used by `gradient_dynamic_tols`.
+* `gradient_tol_factor=$(Defaults.gradient_tol_factor)` : Tolerance scaling factor relative to the boundary algorithm's tolerance, used by `gradient_dynamic_tols` (e.g. `10` makes the gradient tolerance ~10x looser than the boundary tolerance).
+
+## Preconditioning
+
+* `precondition_alg=:$(Defaults.precondition_alg)` : Algorithm variant used for preconditioning the PEPS gradient.
+    - `:LocalPreconditioner` : Precondition using the leading (local) term of the PEPS metric, see [`PEPSKit.LocalPreconditioner`](@ref).
+* `precondition_tol=$(Defaults.precondition_tol)` : Convergence tolerance for the linear problem in the preconditioning step.
+* `precondition_maxiter=$(Defaults.precondition_maxiter)` : Maximal number of iterations for the linear problem in the preconditioning step.
+* `precondition_verbosity=$(Defaults.precondition_verbosity)` : Preconditioning output information verbosity.
+* `precondition_krylovdim=$(Defaults.precondition_krylovdim)` : Krylov dimensionfor the linear problem in the preconditioning step.
+* `precondition_regularization=$(Defaults.precondition_regularization)` : Prefactor setting the regularization strength of the local linear problem, see [`PEPSKit.LocalPreconditioner`](@ref).
+* `precondition_dynamic_tols=$(Defaults.precondition_dynamic_tols)` : If `true`, wrap the preconditioner algorithm in a `MPSKit.DynamicTol` that rescales its tolerance based on the current PEPS optimization gradient norm, see [`PEPSKit.PEPSOptimize`](@ref).
+* `precondition_tol_min=$(Defaults.precondition_tol_min)` : Minimal preconditioner tolerance used by `precondition_dynamic_tols`.
+* `precondition_tol_max=$(Defaults.precondition_tol_max)` : Maximal preconditioner tolerance used by `precondition_dynamic_tols`.
+* `precondition_tol_factor=$(Defaults.precondition_tol_factor)` : Tolerance scaling factor used by `precondition_dynamic_tols`.
 
 ## Optimization
 
@@ -117,6 +139,10 @@ const ctmrg_miniter = 4
 const ctmrg_alg = :SimultaneousCTMRG # ∈ {:SimultaneousCTMRG, :SequentialCTMRG}
 const ctmrg_verbosity = 2
 const sparse = false # TODO: implement sparse CTMRG
+const ctmrg_dynamic_tols = true
+const ctmrg_tol_min = 1.0e-12
+const ctmrg_tol_max = 1.0e-4
+const ctmrg_tol_factor = 1.0e-3
 
 # SVD forward & reverse
 const trunc = :FixedSpaceTruncation # ∈ {:FixedSpaceTruncation, :notrunc, :truncerror, :truncspace, :trunctol}
@@ -151,6 +177,22 @@ const gradient_verbosity = -1
 const gradient_alg = :FixedPointGradient
 const gradient_fixedpoint_solver_alg = :Arnoldi # ∈ {:GMRES, :BiCGStab, :Arnoldi, :GeomSum, :ManualIter}
 const gradient_fixedpoint_solver_eager = true
+const gradient_dynamic_tols = true
+const gradient_tol_min = 1.0e-10
+const gradient_tol_max = 1.0e-1
+const gradient_tol_factor = 1.0e1
+
+# Preconditioning
+const precondition_alg = :LocalPreconditioner
+const precondition_tol = 1.0e-6
+const precondition_maxiter = 1
+const precondition_verbosity = -1
+const precondition_krylovdim = 30
+const precondition_regularization = 100.0
+const precondition_dynamic_tols = true
+const precondition_tol_min = 1.0e-12
+const precondition_tol_max = 1.0e-4
+const precondition_tol_factor = 1.0e-2
 
 # Optimization
 const reuse_env = true

src/PEPSKit.jl

@@ -30,6 +30,8 @@ import TupleTools
 using MPSKit
 using MPSKit: MPSTensor, MPOTensor, GenericMPSTensor, MPSBondTensor, ProductTransferMatrix
 using MPSKit: InfiniteEnvironments
+using MPSKit: DynamicTol, updatetol
+import MPSKit.DynamicTols: _updatetol
 import MPSKit: tensorexpr, leading_boundary, loginit!, logiter!, logfinish!, logcancel!, physicalspace
 import MPSKit: infinite_temperature_density_matrix
 
@@ -141,6 +143,7 @@ include("algorithms/correlator_adapters.jl")
 include("algorithms/correlators.jl")
 
 include("algorithms/optimization/fixed_point_differentiation.jl")
+include("algorithms/optimization/preconditioning.jl")
 include("algorithms/optimization/peps_optimization.jl")
 
 include("algorithms/select_algorithm.jl")

src/algorithms/optimization/fixed_point_differentiation.jl

@@ -93,6 +93,10 @@ end
 FixedPointGradient(; kwargs...) = GradientAlgorithm(; alg = :FixedPointGradient, kwargs...)
 GRADIENT_ALGORITHM_SYMBOLS[:FixedPointGradient] = FixedPointGradient
 
+# `FixedPointGradient` has no top-level `tol` field (it lives on `solver_alg`), so the
+# default `MPSKit.DynamicTols._updatetol` (which sets `alg.tol`) doesn't apply
+_updatetol(alg::FixedPointGradient, tol::Real) = @set alg.solver_alg.tol = tol
+
 const FIXEDPOINT_SOLVER_SYMBOLS = IdDict{Symbol, Type{<:Any}}(
     :GMRES => GMRES, :BiCGStab => BiCGStab, :Arnoldi => Arnoldi,
 )

src/algorithms/optimization/peps_optimization.jl

@@ -17,35 +17,43 @@ For a full description, see [`fixedpoint`](@ref). The supported keywords are:
 * `boundary_alg::Union{NamedTuple,<:CTMRGAlgorithm,...}`
 * `gradient_alg::Union{NamedTuple,Nothing,<:GradientAlgorithm}`
 * `optimizer_alg::Union{NamedTuple,<:OptimKit.OptimizationAlgorithm}`
+* `precondition_alg::Union{NamedTuple,Nothing,<:PreconditionAlgorithm}`
 * `reuse_env::Bool=$(Defaults.reuse_env)`
 * `symmetrization::Union{Nothing,SymmetrizationStyle}=nothing`
 """
-struct PEPSOptimize{B, G}
+struct PEPSOptimize{B, G, P}
     boundary_alg::B
     gradient_alg::G
     optimizer_alg::OptimKit.OptimizationAlgorithm
+    precondition_alg::P
     reuse_env::Bool
     symmetrization::Union{Nothing, SymmetrizationStyle}
 
     function PEPSOptimize(  # Inner constructor to prohibit illegal setting combinations
-            boundary_alg::B, gradient_alg::G, optimizer_alg,
+            boundary_alg::B, gradient_alg::G, optimizer_alg, precondition_alg::P,
             reuse_env, symmetrization,
-        ) where {B, G}
-        _check_algorithm_combination(boundary_alg, gradient_alg, symmetrization)
-        return new{B, G}(boundary_alg, gradient_alg, optimizer_alg, reuse_env, symmetrization)
+        ) where {B, G, P}
+        _check_algorithm_combination(
+            parent_alg(boundary_alg), parent_alg(gradient_alg), symmetrization
+        )
+        return new{B, G, P}(
+            boundary_alg, gradient_alg, optimizer_alg, precondition_alg,
+            reuse_env, symmetrization,
+        )
     end
 end
 
 function PEPSOptimize(;
-        boundary_alg = (;), gradient_alg = (;), optimizer_alg = (;),
+        boundary_alg = (;), gradient_alg = (;), optimizer_alg = (;), precondition_alg = (;),
         reuse_env = Defaults.reuse_env, symmetrization = nothing,
     )
     boundary_algorithm = _alg_or_nt(CTMRGAlgorithm, boundary_alg)
     gradient_algorithm = _alg_or_nt(GradientAlgorithm, gradient_alg)
     optimizer_algorithm = _alg_or_nt(OptimKit.OptimizationAlgorithm, optimizer_alg)
+    precondition_algorithm = _alg_or_nt(PreconditionAlgorithm, precondition_alg)
 
     return PEPSOptimize(
-        boundary_algorithm, gradient_algorithm, optimizer_algorithm,
+        boundary_algorithm, gradient_algorithm, optimizer_algorithm, precondition_algorithm,
         reuse_env, symmetrization,
     )
 end
@@ -135,6 +143,38 @@ keyword arguments are:
     - `: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.
 
+### Preconditioner algorithm
+
+Supply preconditioner parameters via `precondition_alg::Union{NamedTuple,Nothing,<:PreconditionAlgorithm}`
+using either a `NamedTuple` of keyword arguments, `nothing`, or a `PreconditionAlgorithm`
+struct directly. By default, the gradient is preconditioned with the local PEPS metric, see
+[`LocalPreconditioner`](@ref); pass `nothing` to disable preconditioning and optimize using
+the raw Euclidean gradient. The supported `NamedTuple` keyword arguments are:
+
+* `alg::Symbol=:$(Defaults.precondition_alg)` : Preconditioner algorithm variant, can be one of the following:
+    - `:LocalPreconditioner` : Precondition using the leading (local) term of the PEPS metric, see [`LocalPreconditioner`](@ref)
+* `tol::Real=$(Defaults.precondition_tol)` : Convergence tolerance of the local linear problem.
+* `maxiter::Int=$(Defaults.precondition_maxiter)` : Maximal number of iterations of the local linear problem.
+* `verbosity::Int` : Preconditioner output verbosity, ≤0 by default to disable too verbose printing. Should only be >0 for debug purposes.
+* `krylovdim::Int=$(Defaults.precondition_krylovdim)` : Krylov dimension of the local linear problem.
+* `regularization::Real=$(Defaults.precondition_regularization)` : Prefactor setting the regularization strength of the local linear problem.
+
+### Dynamic tolerances
+
+The boundary, gradient and preconditioner algorithms each additionally accept the keyword
+arguments below, which wrap the corresponding algorithm in an `MPSKit.DynamicTols.DynamicTol` that
+rescales its tolerance over the course of the optimization. This allows the intermediate
+problems to be solved only as accurately as the current optimization step requires, which
+can significantly reduce the total runtime. The boundary and preconditioner tolerances are
+scaled relative to the current gradient norm, while the gradient tolerance is in turn scaled
+relative to the effective boundary tolerance. These settings are only available within a
+variational optimization, not for standalone [`leading_boundary`](@ref) calls.
+
+* `dynamic_tols::Bool` : Enable dynamic tolerance scaling for this algorithm. Defaults to `$(Defaults.ctmrg_dynamic_tols)`, `$(Defaults.gradient_dynamic_tols)` and `$(Defaults.precondition_dynamic_tols)` for the boundary, gradient and preconditioner algorithm respectively.
+* `tol_min::Real` : Lower clamp on the dynamically scaled tolerance.
+* `tol_max::Real` : Upper clamp on the dynamically scaled tolerance.
+* `tol_factor::Real` : Prefactor of the dynamically scaled tolerance.
+
 ### Optimizer settings
 
 Supply the optimizer algorithm via `optimizer_alg::Union{NamedTuple,<:OptimKit.OptimizationAlgorithm}`
@@ -189,36 +229,63 @@ function fixedpoint(
         )
     end
 
-    # initialize info collection vectors
     T = promote_type(real(scalartype(peps₀)), real(scalartype(env₀)))
+
+    # `tol_state` tracks (iter, gradnorm) of the last accepted optimization step
+    # (updated only in `finalize!`), used to adjust `alg.boundary_alg`/`alg.gradient_alg`
+    # via `MPSKit.updatetol` if they are wrapped in a `MPSKit.DynamicTol`. `latest_*`
+    # hold the values produced by the current `fg` call, and are only recorded into
+    # their respective history vectors once a step is accepted.
+    tol_state = Ref((iter = 0, gradnorm = one(T)))
+    latest_metrics = Ref{NamedTuple}()
+    latest_gradnorms = Ref{Matrix{T}}()
+    latest_time = Ref(0.0)
+
+    # initialize info collection vectors
     contraction_metrics = Vector{NamedTuple}()
     gradnorms_unitcell = Vector{Matrix{T}}()
     times = Vector{Float64}()
+    finalize! = track_state_and_finalize!(
+        tol_state, latest_metrics, latest_gradnorms, latest_time,
+        contraction_metrics, gradnorms_unitcell, times, finalize!,
+    )
 
     # normalize the initial guess
     peps₀ = peps_normalize(peps₀)
 
+    # initialize the preconditioner
+    function precondition(x, g)
+        precondition_alg = updatetol(
+            alg.precondition_alg, tol_state[].iter, tol_state[].gradnorm
+        )
+        return peps_precondition(x, g, tol_state, precondition_alg)
+    end
+
     # optimize operator cost function
     (peps_final, env_final), cost_final, ∂cost, numfg, convergence_history = optimize(
         (peps₀, env₀), alg.optimizer_alg;
-        retract, inner = real_inner, (transport!) = (peps_transport!),
+        retract, inner = real_inner, (transport!) = (peps_transport!), precondition,
         hasconverged, shouldstop, finalize!,
     ) do (peps, env)
         start_time = time_ns()
+        boundary_alg = updatetol(alg.boundary_alg, tol_state[].iter, tol_state[].gradnorm)
+        # gradient tolerance is scaled relative to the boundary algorithm's own
+        # (just-updated) effective tolerance, not directly to the gradient norm
+        gradient_alg = updatetol(alg.gradient_alg, tol_state[].iter, boundary_alg.tol)
         E, gs = withgradient(peps) do ψ
             env′, info = hook_pullback(
-                leading_boundary, env, ψ, alg.boundary_alg;
-                alg_rrule = alg.gradient_alg,
+                leading_boundary, env, ψ, boundary_alg;
+                alg_rrule = gradient_alg,
             )
             ignore_derivatives() do
                 alg.reuse_env && update!(env, env′)
-                push!(contraction_metrics, info.contraction_metrics)
+                latest_metrics[] = info.contraction_metrics
             end
             return cost_function(ψ, env′, operator)
         end
         g = only(gs)  # `withgradient` returns tuple of gradients `gs`
-        push!(gradnorms_unitcell, norm.(g.A))
-        push!(times, (time_ns() - start_time) * 1.0e-9)
+        latest_gradnorms[] = norm.(unitcell(g))
+        latest_time[] = (time_ns() - start_time) * 1.0e-9
         return E, g
     end
 
@@ -235,13 +302,14 @@ function fixedpoint(
 end
 
 """
-    check_input(::typeof(fixedpoint), peps₀, env₀, alg::PEPSOptimize{<:SimultaneousCTMRG})
+    check_input(::typeof(fixedpoint), peps₀, env₀, alg::PEPSOptimize)
 
 Check compatibility of an initial PEPS and environment with a specified PEPS optimization algorithm.
 """
-function check_input(::typeof(fixedpoint), peps₀, env₀, alg::PEPSOptimize) end
-function check_input(::typeof(fixedpoint), peps₀, env₀, alg::PEPSOptimize{<:SimultaneousCTMRG, <:FixedPointGradient})
-    if scalartype(env₀) <: Real # :fixed mode gauge fixing is incompatible with real environments
+function check_input(::typeof(fixedpoint), peps₀, env₀, alg::PEPSOptimize)
+    if parent_alg(alg.boundary_alg) isa SimultaneousCTMRG &&
+            parent_alg(alg.gradient_alg) isa FixedPointGradient &&
+            scalartype(env₀) <: Real # :fixed mode gauge fixing is incompatible with real environments
         msg = "the provided real environment is incompatible with :fixed mode \
         since :fixed mode generally produces complex gauges"
         throw(ArgumentError(msg))
@@ -338,3 +406,36 @@ function symmetrize_retract_and_finalize!(
     end
     return retract_then_symmetrize, symmetrize_then_finalize!
 end
+
+"""
+    track_state_and_finalize!(
+        tol_state::Base.RefValue, latest_metrics::Base.RefValue, latest_gradnorms::Base.RefValue,
+        latest_time::Base.RefValue, contraction_metrics::Vector, gradnorms_unitcell::Vector,
+        times::Vector, [finalize!],
+    )
+
+Return a `finalize!` function that, after calling `finalize!` (defaulting to
+`OptimKit._finalize!`):
+* updates `tol_state[]` to the `(iter, gradnorm)` of the now-accepted optimization step,
+  used to drive `MPSKit.updatetol` for any `alg.boundary_alg`/`alg.gradient_alg` wrapped
+  in a `MPSKit.DynamicTol`
+* records `latest_metrics[]`/`latest_gradnorms[]`/`latest_time[]`, i.e. the values
+  produced by the `fg` call corresponding to the accepted step, into
+  `contraction_metrics`/`gradnorms_unitcell`/`times`
+"""
+function track_state_and_finalize!(
+        tol_state::Base.RefValue, latest_metrics::Base.RefValue, latest_gradnorms::Base.RefValue,
+        latest_time::Base.RefValue, contraction_metrics::Vector, gradnorms_unitcell::Vector,
+        times::Vector, (finalize!) = OptimKit._finalize!,
+    )
+    function commit_state_and_finalize!((peps, env), E, grad, numiter)
+        (peps, env), E, grad = finalize!((peps, env), E, grad, numiter)
+        gradnorm = sqrt(real_inner((peps, env), grad, grad))
+        tol_state[] = (; iter = numiter, gradnorm)
+        push!(contraction_metrics, latest_metrics[])
+        push!(gradnorms_unitcell, latest_gradnorms[])
+        push!(times, latest_time[])
+        return (peps, env), E, grad
+    end
+    return commit_state_and_finalize!
+end