diff --git a/docs/src/assets/pepskit.bib b/docs/src/assets/pepskit.bib index 15949e7c3..00c0c5341 100644 --- a/docs/src/assets/pepskit.bib +++ b/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} +} diff --git a/src/Defaults.jl b/src/Defaults.jl index 9bcbb6244..9cf00233c 100644 --- a/src/Defaults.jl +++ b/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 diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 0139d60d0..4a6956069 100644 --- a/src/PEPSKit.jl +++ b/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") diff --git a/src/algorithms/optimization/fixed_point_differentiation.jl b/src/algorithms/optimization/fixed_point_differentiation.jl index 2c2a54991..bf4f2403d 100644 --- a/src/algorithms/optimization/fixed_point_differentiation.jl +++ b/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, ) diff --git a/src/algorithms/optimization/peps_optimization.jl b/src/algorithms/optimization/peps_optimization.jl index 4100611bb..48733766e 100644 --- a/src/algorithms/optimization/peps_optimization.jl +++ b/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 diff --git a/src/algorithms/optimization/preconditioning.jl b/src/algorithms/optimization/preconditioning.jl new file mode 100644 index 000000000..e3b843d76 --- /dev/null +++ b/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::Base.RefValue, 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::Base.RefValue, alg::Nothing) = g + +function peps_precondition(x, g, tol_state::Base.RefValue, 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 diff --git a/src/algorithms/select_algorithm.jl b/src/algorithms/select_algorithm.jl index 93445a6b5..c121fb3fd 100644 --- a/src/algorithms/select_algorithm.jl +++ b/src/algorithms/select_algorithm.jl @@ -1,7 +1,48 @@ _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")) +""" + parent_alg(alg) + +Unwrap an algorithm from a `MPSKit.DynamicTols.DynamicTol` wrapper, if +present, returning the algorithm it wraps. Falls back to returning `alg` unchanged for any +other input, including `nothing`. +""" +parent_alg(alg) = alg +parent_alg(alg::DynamicTol) = parent_alg(alg.alg) + +""" + _dynamic_tol_or_alg(alg; dynamic_tols::Bool, tol_min::Real, tol_max::Real, tol_factor::Real) + +Wrap `alg` in a `MPSKit.DynamicTols.DynamicTol` with the given +tolerance-scaling settings if `dynamic_tols` is `true`, otherwise return `alg` unchanged. +""" +function _dynamic_tol_or_alg(alg; dynamic_tols::Bool, tol_min::Real, tol_max::Real, tol_factor::Real) + return dynamic_tols ? DynamicTol(alg, tol_min, tol_max, tol_factor) : alg +end + +const DYNAMIC_TOL_KWARGS = (; dynamic_tols = nothing, tol_min = nothing, tol_max = nothing, tol_factor = nothing) + +""" + _pop_dynamic_tol_kwargs(kwargs::NamedTuple) -> dynamic_tol_kwargs, remaining_kwargs + +Split off the `dynamic_tols`/`tol_min`/`tol_max`/`tol_factor` entries from `kwargs`, +returning them separately (to be passed on to [`_dynamic_tol_or_alg`](@ref)) from the +remaining keyword arguments (to be passed on to the algorithm constructor). +""" +function _pop_dynamic_tol_kwargs(kwargs::NamedTuple) + dynamic_tol_kwargs = (; + dynamic_tols = kwargs.dynamic_tols, + tol_min = kwargs.tol_min, + tol_max = kwargs.tol_max, + tol_factor = kwargs.tol_factor, + ) + return dynamic_tol_kwargs, Base.structdiff(kwargs, DYNAMIC_TOL_KWARGS) +end + """ select_algorithm(func_or_alg, args...; kwargs...) -> Algorithm @@ -22,29 +63,61 @@ 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 if boundary_alg isa NamedTuple - defaults = (; verbosity = verbosity ≤ 3 ? -1 : 3, tol = 1.0e-4tol) + defaults = (; + verbosity = verbosity ≤ 3 ? -1 : 3, tol = 1.0e-4tol, + dynamic_tols = Defaults.ctmrg_dynamic_tols, + tol_min = Defaults.ctmrg_tol_min, + tol_max = Defaults.ctmrg_tol_max, + tol_factor = Defaults.ctmrg_tol_factor, + ) boundary_kwargs = merge(defaults, boundary_alg) + dynamic_tol_kwargs, boundary_kwargs = _pop_dynamic_tol_kwargs(boundary_kwargs) boundary_alg = select_algorithm(leading_boundary, env₀; boundary_kwargs...) + boundary_alg = _dynamic_tol_or_alg(boundary_alg; dynamic_tol_kwargs...) end # C4vCTMRG-specific defaults - if boundary_alg isa C4vCTMRG + if parent_alg(boundary_alg) isa C4vCTMRG # symmetrize state and gradient if isnothing(symmetrization) symmetrization = RotateReflect() end end - # adjust gradient verbosity + # adjust gradient verbosity and construct the gradient algorithm if gradient_alg isa NamedTuple # TODO: check this: - defaults = (; verbosity = verbosity ≤ 3 ? -1 : 3, tol = 1.0e-2tol) - gradient_alg = merge(defaults, gradient_alg) + defaults = (; + verbosity = verbosity ≤ 3 ? -1 : 3, tol = 1.0e-2tol, + dynamic_tols = Defaults.gradient_dynamic_tols, + tol_min = Defaults.gradient_tol_min, + tol_max = Defaults.gradient_tol_max, + tol_factor = Defaults.gradient_tol_factor, + ) + gradient_kwargs = merge(defaults, gradient_alg) + dynamic_tol_kwargs, gradient_kwargs = _pop_dynamic_tol_kwargs(gradient_kwargs) + gradient_alg = GradientAlgorithm(; gradient_kwargs...) + 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 @@ -53,7 +126,10 @@ function select_algorithm( 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( diff --git a/src/states/infinitepeps.jl b/src/states/infinitepeps.jl index d25b19e7a..a1d31c8ff 100644 --- a/src/states/infinitepeps.jl +++ b/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)) diff --git a/test/examples/heisenberg.jl b/test/examples/heisenberg.jl index 5a06daed9..8441998f6 100644 --- a/test/examples/heisenberg.jl +++ b/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)),