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

@@ -94,6 +94,20 @@ Module containing default algorithm parameter values and arguments.
 * `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 dimension for 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
 
 * `reuse_env=$(Defaults.reuse_env)` : If `true`, the current optimization step is initialized on the previous environment, otherwise a random environment is used.
@@ -168,6 +182,18 @@ 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
 const optimizer_tol = 1.0e-4

src/PEPSKit.jl

@@ -143,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/peps_optimization.jl

@@ -17,37 +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}
+        ) where {B, G, P}
         _check_algorithm_combination(
             parent_alg(boundary_alg), parent_alg(gradient_alg), symmetrization
         )
-        return new{B, G}(boundary_alg, gradient_alg, optimizer_alg, reuse_env, 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
@@ -137,18 +143,34 @@ 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 and gradient algorithms each additionally accept the keyword arguments below,
-which wrap the corresponding algorithm in an `MPSKit.DynamicTols.DynamicTol` that
+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 tolerance is 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
+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)` and `$(Defaults.gradient_dynamic_tols)` for the boundary and gradient algorithm respectively.
+* `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.
@@ -227,10 +249,18 @@ function fixedpoint(
     # normalize the initial guess
     peps₀ = peps_normalize(peps₀)
 
+    # initialize the preconditioner
+    function precondition(x, g)
+        precondition_alg = updatetol(
+            alg.precondition_alg, tracked_finalizer.tol_state.iter, tracked_finalizer.tol_state.gradnorm
+        )
+        return peps_precondition(x, g, tracked_finalizer.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!) = tracked_finalizer,
     ) do (peps, env)
         start_time = time()

src/algorithms/optimization/preconditioning.jl

@@ -0,0 +1,125 @@
+abstract type PreconditionAlgorithm end
+
+const PRECONDITION_ALGORITHM_SYMBOLS = IdDict{Symbol, Type{<:PreconditionAlgorithm}}()
+
+"""
+    PreconditionAlgorithm(; kwargs...)
+
+Keyword argument parser returning the appropriate `PreconditionAlgorithm` algorithm struct.
+"""
+function PreconditionAlgorithm(;
+        alg = Defaults.precondition_alg,
+        tol = Defaults.precondition_tol,
+        maxiter = Defaults.precondition_maxiter,
+        verbosity = Defaults.precondition_verbosity,
+        krylovdim = Defaults.precondition_krylovdim,
+        regularization = Defaults.precondition_regularization,
+    )
+    # replace symbol with PreconditionAlgorithm alg type
+    haskey(PRECONDITION_ALGORITHM_SYMBOLS, alg) ||
+        throw(ArgumentError("unknown PreconditionAlgorithm algorithm: $alg"))
+    alg_type = PRECONDITION_ALGORITHM_SYMBOLS[alg]
+
+    return alg_type(GMRES(; tol, maxiter, verbosity, krylovdim), regularization)
+end
+
+"""
+$(TYPEDEF)
+
+Preconditioner for PEPS ground-state optimization based on the leading term of the PEPS
+metric tensor, i.e. the local norm matrix ``N`` obtained by contracting the full CTMRG
+environment around a given unit-cell site while leaving the ket and bra PEPS tensor at that
+site open.
+
+Instead of following the raw Euclidean gradient ``g``, the optimizer is fed the solution
+``g̃`` of the regularized local linear problem
+```math
+( N / ⟨ψ|ψ⟩ + δ ) g̃ = g
+```
+which is solved separately for every tensor in the unit cell using an iterative solver.
+The regularization ``δ`` prevents the (generally singular) metric from being inverted
+exactly, and decays during the optimization according to
+```math
+δ = \\texttt{regularization} × ‖∇f‖^2 / \\texttt{iter}^2
+```
+so that the preconditioner acts conservatively far from the minimum and becomes
+increasingly aggressive as the optimization converges.
+
+## Fields
+
+$(TYPEDFIELDS)
+
+## Constructors
+
+    LocalPreconditioner(; kwargs...)
+
+Construct a local preconditioner algorithm struct based on keyword arguments. The supported
+keywords are:
+
+* `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. Note that the default performs a single iteration, such that the metric is only inverted approximately.
+* `verbosity::Int=$(Defaults.precondition_verbosity)` : Preconditioner output verbosity, ≤0 by default to disable too verbose printing.
+* `krylovdim::Int=$(Defaults.precondition_krylovdim)` : Krylov dimension of the local linear problem.
+* `regularization::Real=$(Defaults.precondition_regularization)` : Prefactor setting the regularization strength ``δ``, see above.
+
+Reference: [Phys. Rev. B 113, 125111](@cite zhang_accelerating_2026)
+"""
+struct LocalPreconditioner{A} <: PreconditionAlgorithm
+    "solver algorithm used for the local linear problem"
+    solver_alg::A
+
+    "prefactor setting the regularization strength of the local linear problem"
+    regularization::Float64
+end
+LocalPreconditioner(; kwargs...) = PreconditionAlgorithm(; alg = :LocalPreconditioner, kwargs...)
+PRECONDITION_ALGORITHM_SYMBOLS[:LocalPreconditioner] = LocalPreconditioner
+
+# `LocalPreconditioner` 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::LocalPreconditioner, tol::Real) = @set alg.solver_alg.tol = tol
+
+"""
+$(SIGNATURES)
+
+Apply the regularized local metric ``N / ⟨ψ|ψ⟩ + δ`` to a single PEPS tensor `g_rc`, where
+``N`` is obtained by contracting the CTMRG environment `env` around the unit-cell site
+`(r, c)` and `norm_pref` is the corresponding local norm ``⟨ψ|ψ⟩``.
+"""
+function apply_local_preconditioner(g_rc::PEPSTensor, env::CTMRGEnv, δ, (r, c), norm_pref)
+    @autoopt @tensor g_rc_prec[d; D_N_below D_E_below D_S_below D_W_below] :=
+        g_rc[d; D_N_above D_E_above D_S_above D_W_above] *
+        corner(env, NORTHWEST, r - 1, c - 1)[χ_WNW; χ_NNW] *
+        edge(env, NORTH, r - 1, c)[χ_NNW D_N_above D_N_below; χ_NNE] *
+        corner(env, NORTHEAST, r - 1, c + 1)[χ_NNE; χ_ENE] *
+        edge(env, EAST, r, c + 1)[χ_ENE D_E_above D_E_below; χ_ESE] *
+        corner(env, SOUTHEAST, r + 1, c + 1)[χ_ESE; χ_SSE] *
+        edge(env, SOUTH, r + 1, c)[χ_SSE D_S_above D_S_below; χ_SSW] *
+        corner(env, SOUTHWEST, r + 1, c - 1)[χ_SSW; χ_WSW] *
+        edge(env, WEST, r, c - 1)[χ_WSW D_W_above D_W_below; χ_WNW]
+    g_rc_prec = twistdual(g_rc_prec, 2:5)
+    return g_rc_prec / norm_pref + δ * g_rc
+end
+
+"""
+    peps_precondition(x, g, tol_state::NamedTuple, alg)
+
+Precondition the PEPS gradient `g` at the point `x = (peps, env)` according to the
+preconditioner algorithm `alg`. Passing `alg = nothing` disables preconditioning and returns
+`g` unchanged. See [`LocalPreconditioner`](@ref) for details on the local metric
+preconditioner, and note that its regularization strength is set based on the current
+gradient norm and iteration count tracked in `tol_state`.
+"""
+peps_precondition(x, g, tol_state::NamedTuple, alg::Nothing) = g
+
+function peps_precondition(x, g, tol_state::NamedTuple, alg::LocalPreconditioner)
+    peps, env = x
+    δ = alg.regularization * tol_state.gradnorm^2 / max(tol_state.iter, 1)
+    g_prec_unitcell = map(eachcoordinate(g)) do (r, c)
+        nf = _contract_site((r, c), InfiniteSquareNetwork(peps), env)
+        g_rc_prec, = linsolve(g[r, c], g[r, c], alg.solver_alg) do g_in
+            return apply_local_preconditioner(g_in, env, δ, (r, c), nf)
+        end
+        return g_rc_prec
+    end
+    return InfinitePEPS(g_prec_unitcell)
+end

src/algorithms/select_algorithm.jl

@@ -1,6 +1,7 @@
 _alg_or_nt(::Type{T}, alg::NamedTuple) where {T} = T(; alg...)
 _alg_or_nt(::Type{T}, alg::A) where {T, A <: T} = alg
 _alg_or_nt(::Type{T}, alg::DynamicTol{<:T}) where {T} = alg
+_alg_or_nt(::Type, ::Nothing) = nothing
 _alg_or_nt(T, alg) = throw(ArgumentError("unkown $T: $alg"))
 
 """
@@ -62,7 +63,7 @@ function select_algorithm(
         env₀;
         tol = Defaults.optimizer_tol, # top-level tolerance
         verbosity = 3, # top-level verbosity
-        boundary_alg = (;), gradient_alg = (;), optimizer_alg = (;),
+        boundary_alg = (;), gradient_alg = (;), optimizer_alg = (;), precondition_alg = (;),
         symmetrization = nothing, kwargs...,
     )
     # adjust CTMRG tols and verbosity
@@ -104,13 +105,31 @@ function select_algorithm(
         gradient_alg = _dynamic_tol_or_alg(gradient_alg; dynamic_tol_kwargs...)
     end
 
+    # adjust preconditioner verbosity and construct the preconditioner algorithm
+    if precondition_alg isa NamedTuple
+        defaults = (;
+            verbosity = verbosity ≤ 3 ? -1 : 3,
+            dynamic_tols = Defaults.precondition_dynamic_tols,
+            tol_min = Defaults.precondition_tol_min,
+            tol_max = Defaults.precondition_tol_max,
+            tol_factor = Defaults.precondition_tol_factor,
+        )
+        precondition_kwargs = merge(defaults, precondition_alg)
+        dynamic_tol_kwargs, precondition_kwargs = _pop_dynamic_tol_kwargs(precondition_kwargs)
+        precondition_alg = PreconditionAlgorithm(; precondition_kwargs...)
+        precondition_alg = _dynamic_tol_or_alg(precondition_alg; dynamic_tol_kwargs...)
+    end
+
     # adjust optimizer tol and verbosity
     if optimizer_alg isa NamedTuple
         defaults = (; tol, verbosity)
         optimizer_alg = merge(defaults, optimizer_alg)
     end
 
-    return PEPSOptimize(; boundary_alg, gradient_alg, optimizer_alg, symmetrization, kwargs...)
+    return PEPSOptimize(;
+        boundary_alg, gradient_alg, optimizer_alg, precondition_alg,
+        symmetrization, kwargs...,
+    )
 end
 
 function select_algorithm(

src/states/infinitepeps.jl

@@ -142,6 +142,9 @@ function eachcoordinate(A::InfinitePEPS, dirs)
     return collect(Iterators.product(dirs, axes(A, 1), axes(A, 2)))
 end
 
+Base.real(A::InfinitePEPS) = InfinitePEPS(real.(unitcell(A)))
+Base.complex(A::InfinitePEPS) = InfinitePEPS(complex.(unitcell(A)))
+
 ## Spaces
 
 TensorKit.spacetype(::Type{T}) where {T <: InfinitePEPS} = spacetype(eltype(T))

test/examples/heisenberg.jl

@@ -139,7 +139,7 @@ end
 
     # continue with auto differentiation
     peps_final, env_final, E_final, = fixedpoint(
-        ham, peps, complex(env); # make environment complex explicitly
+        ham, complex(peps), complex(env); # make environment complex explicitly
         optimizer_alg = (; tol = gradtol, maxiter = 25),
         boundary_alg = (; maxiter = ctmrg_maxiter),
         gradient_alg = (; solver_alg = (; alg = :GMRES)),