From 5767ba3e1d4caa8936f150049ec2e5afe6217afb Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Fri, 13 Dec 2024 16:23:54 +0100 Subject: [PATCH 01/28] Add rough draft --- Project.toml | 4 + src/PEPSKit.jl | 73 +++++++++---- src/algorithms/peps_opt.jl | 208 +++++++++++++++++++++++++------------ src/algorithms/toolbox.jl | 9 -- 4 files changed, 198 insertions(+), 96 deletions(-) diff --git a/Project.toml b/Project.toml index ba49bc982..3d0a77b60 100644 --- a/Project.toml +++ b/Project.toml @@ -13,6 +13,8 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36" MPSKit = "bb1c41ca-d63c-52ed-829e-0820dda26502" MPSKitModels = "ca635005-6f8c-4cd1-b51d-8491250ef2ab" +Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e" +Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5" OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5" OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" @@ -33,6 +35,8 @@ LinearAlgebra = "1" LoggingExtras = "1" MPSKit = "0.11" MPSKitModels = "0.3.5" +Manifolds = "0.10.8" +Manopt = "0.5.4" OhMyThreads = "0.7" OptimKit = "0.3" Printf = "1" diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 74d68e03b..17e85914f 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -5,7 +5,7 @@ using Base: @kwdef using Compat using Accessors: @set using VectorInterface -using TensorKit, KrylovKit, MPSKit, OptimKit, TensorOperations +using TensorKit, KrylovKit, MPSKit, TensorOperations using ChainRulesCore, Zygote using LoggingExtras using MPSKit: loginit!, logiter!, logfinish!, logcancel! @@ -59,27 +59,44 @@ include("utility/symmetrization.jl") """ module Defaults + # CTMRG + const ctmrg_tol = 1e-8 const ctmrg_maxiter = 100 const ctmrg_miniter = 4 - const ctmrg_tol = 1e-8 - const fpgrad_maxiter = 30 - const fpgrad_tol = 1e-6 - const reuse_env = true + const sparse = false const trscheme = FixedSpaceTruncation() const fwd_alg = TensorKit.SDD() - const rrule_alg = Arnoldi(; tol=1e-2fpgrad_tol, krylovdim=48, verbosity=-1) + const rrule_alg = Arnoldi(; tol=ctmrg_tol, krylovdim=48, verbosity=-1) const svd_alg = SVDAdjoint(; fwd_alg, rrule_alg) const projector_alg_type = HalfInfiniteProjector const projector_alg = projector_alg_type(svd_alg, trscheme, 2) const ctmrg_alg = SimultaneousCTMRG( ctmrg_tol, ctmrg_maxiter, ctmrg_miniter, 2, projector_alg ) - const optimizer = LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=2) - const gradient_linsolver = KrylovKit.BiCGStab(; - maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol + + # Optimization + const optim_alg = quasi_Newton + const record_group = [ + RecordCost(), + RecordGradientNorm(), + RecordConditionNumber(), + RecordCostUnitCell(), + RecordTime(), + ] + const optim_kwargs = (; + memory_size=32, + stopping_criterion=StopAfterIteration(100) | StopWhenGradientNormLess(1e-4), + record=record_group, + return_state=true, ) + const fpgrad_maxiter = 30 + const fpgrad_tol = 1e-6 + const gradient_linsolver = KrylovKit.BiCGStab(; maxiter=fpgrad_maxiter, tol=fpgrad_tol) const iterscheme = :fixed + const reuse_env = true const gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme) + + # OhMyThreads scheduler defaults const scheduler = Ref{Scheduler}(Threads.nthreads() == 1 ? SerialScheduler() : DynamicScheduler()) end @@ -88,9 +105,6 @@ Module containing default values that represent typical algorithm parameters. - `ctmrg_maxiter`: Maximal number of CTMRG iterations per run - `ctmrg_miniter`: Minimal number of CTMRG carried out - `ctmrg_tol`: Tolerance checking singular value and norm convergence -- `fpgrad_maxiter`: Maximal number of iterations for computing the CTMRG fixed-point gradient -- `fpgrad_tol`: Convergence tolerance for the fixed-point gradient iteration -- `reuse_env`: If `true`, the current optimization step is initialized on the previous environment - `trscheme`: Truncation scheme for SVDs and other decompositions - `fwd_alg`: SVD algorithm that is used in the forward pass - `rrule_alg`: Reverse-rule for differentiating that SVD @@ -98,41 +112,60 @@ Module containing default values that represent typical algorithm parameters. - `projector_alg_type`: Default type of projector algorithm - `projector_alg`: Algorithm to compute CTMRG projectors - `ctmrg_alg`: Algorithm for performing CTMRG runs -- `optimizer`: Optimization algorithm for PEPS ground-state optimization +- `fpgrad_maxiter`: Maximal number of iterations for computing the CTMRG fixed-point gradient +- `fpgrad_tol`: Convergence tolerance for the fixed-point gradient iteration +- `reuse_env`: If `true`, the current optimization step is initialized on the previous environment - `gradient_linsolver`: Default linear solver for the `LinSolver` gradient algorithm - `iterscheme`: Scheme for differentiating one CTMRG iteration - `gradient_alg`: Algorithm to compute the gradient fixed-point +# TODO - `scheduler`: Multi-threading scheduler which can be accessed via `set_scheduler!` """ module Defaults using TensorKit, KrylovKit, OptimKit, OhMyThreads + using Manopt using PEPSKit: LinSolver, FixedSpaceTruncation, SVDAdjoint, HalfInfiniteProjector, SimultaneousCTMRG + + # CTMRG const ctmrg_tol = 1e-8 const ctmrg_maxiter = 100 const ctmrg_miniter = 4 - const fpgrad_maxiter = 30 - const fpgrad_tol = 1e-6 const sparse = false - const reuse_env = true const trscheme = FixedSpaceTruncation() const fwd_alg = TensorKit.SDD() - const rrule_alg = Arnoldi(; tol=1e-2fpgrad_tol, krylovdim=48, verbosity=-1) + const rrule_alg = Arnoldi(; tol=ctmrg_tol, krylovdim=48, verbosity=-1) const svd_alg = SVDAdjoint(; fwd_alg, rrule_alg) const projector_alg_type = HalfInfiniteProjector const projector_alg = projector_alg_type(svd_alg, trscheme, 2) const ctmrg_alg = SimultaneousCTMRG( ctmrg_tol, ctmrg_maxiter, ctmrg_miniter, 2, projector_alg ) - const optimizer = LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=2) - const gradient_linsolver = KrylovKit.BiCGStab(; - maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol + + # Optimization + const optim_alg = quasi_Newton + const record_group = [ + RecordCost(), + RecordGradientNorm(), + RecordConditionNumber(), # TODO: implement PEPS record actions + RecordCostUnitCell(), + RecordTime(), + ] + const optim_kwargs = (; + memory_size=32, + stopping_criterion=StopAfterIteration(100) | StopWhenGradientNormLess(1e-4), + record=record_group, + return_state=true, ) + const fpgrad_maxiter = 30 + const fpgrad_tol = 1e-6 + const gradient_linsolver = KrylovKit.BiCGStab(; maxiter=fpgrad_maxiter, tol=fpgrad_tol) const iterscheme = :fixed + const reuse_env = true const gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme) # OhMyThreads scheduler defaults diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl index 7be93c92d..3a52d0636 100644 --- a/src/algorithms/peps_opt.jl +++ b/src/algorithms/peps_opt.jl @@ -1,3 +1,5 @@ +using Manifolds, Manopt + abstract type GradMode{F} end iterscheme(::GradMode{F}) where {F} = F @@ -90,13 +92,13 @@ step. The CTMRG gradient itself is computed using the `gradient_alg` algorithm. """ struct PEPSOptimize{G} boundary_alg::CTMRGAlgorithm - optimizer::OptimKit.OptimizationAlgorithm + optim_alg reuse_env::Bool gradient_alg::G function PEPSOptimize( # Inner constructor to prohibit illegal setting combinations boundary_alg::CTMRGAlgorithm, - optimizer, + optim_alg, reuse_env, gradient_alg::G, ) where {G} @@ -105,57 +107,70 @@ struct PEPSOptimize{G} throw(ArgumentError(":sequential and :fixed are not compatible")) end end - return new{G}(boundary_alg, optimizer, reuse_env, gradient_alg) + return new{G}(boundary_alg, optim_alg, reuse_env, gradient_alg) end end function PEPSOptimize(; boundary_alg=Defaults.ctmrg_alg, - optimizer=Defaults.optimizer, + optim_alg=Defaults.optim_alg, reuse_env=Defaults.reuse_env, gradient_alg=Defaults.gradient_alg, ) - return PEPSOptimize(boundary_alg, optimizer, reuse_env, gradient_alg) + return PEPSOptimize(boundary_alg, optim_alg, reuse_env, gradient_alg) +end + +struct PEPSEnergyCost + env::CTMRGEnv + hamiltonian::LocalOperator + alg::PEPSOptimize + from_vec::Function +end + +function (pec::PEPSEnergyCost)(peps_vec::Vector) + peps = pec.from_vec(peps_vec) + E, gs = withgradient(peps) do ψ + env′ = hook_pullback( + leading_boundary, + pec.env, + ψ, + pec.alg.boundary_alg; + alg_rrule=pec.alg.gradient_alg, + ) + ignore_derivatives() do + pec.alg.reuse_env && update!(pec.env, env′) + end + + E = expectation_value(peps, pec.hamiltonian, env′) + ignore_derivatives() do + isapprox(imag(E), 0; atol=sqrt(eps(real(E)))) || + @warn "Expectation value is not real: $E." + end + return real(E) + end + g = only(gs) # `withgradient` returns tuple of gradients `gs` + return E, to_vec(g)[1] end """ - fixedpoint(ψ₀::InfinitePEPS{T}, H, alg::PEPSOptimize, [env₀::CTMRGEnv]; - finalize!=OptimKit._finalize!, symmetrization=nothing) where {T} - -Optimize `ψ₀` with respect to the Hamiltonian `H` according to the parameters supplied -in `alg`. The initial environment `env₀` serves as an initial guess for the first CTMRG run. -By default, a random initial environment is used. - -The `finalize!` kwarg can be used to insert a function call after each optimization step -by utilizing the `finalize!` kwarg of `OptimKit.optimize`. -The function maps `(peps, envs), f, g = finalize!((peps, envs), f, g, numiter)`. -The `symmetrization` kwarg accepts `nothing` or a `SymmetrizationStyle`, in which case the -PEPS and PEPS gradient are symmetrized after each optimization iteration. Note that this -requires a symmmetric `ψ₀` and `env₀` to converge properly. - -The function returns a `NamedTuple` which contains the following entries: -- `peps`: final `InfinitePEPS` -- `env`: `CTMRGEnv` corresponding to the final PEPS -- `E`: final energy -- `E_history`: convergence history of the energy function -- `grad`: final energy gradient -- `gradnorm_history`: convergence history of the energy gradient norms -- `numfg`: total number of calls to the energy function +TODO """ function fixedpoint( - ψ₀::InfinitePEPS{F}, + peps₀::InfinitePEPS{T}, H, alg::PEPSOptimize, - env₀::CTMRGEnv=CTMRGEnv(ψ₀, field(F)^20); - (finalize!)=OptimKit._finalize!, - symmetrization=nothing, -) where {F} - if isnothing(symmetrization) - retract = peps_retract - else - retract, symm_finalize! = symmetrize_retract_and_finalize!(symmetrization) - fin! = finalize! # Previous finalize! - finalize! = (x, f, g, numiter) -> fin!(symm_finalize!(x, f, g, numiter)..., numiter) - end + env₀::CTMRGEnv=CTMRGEnv(peps₀, field(T)^20); + # (finalize!)=OptimKit._finalize!, + # symmetrization=nothing, + optim_kwargs=Defaults.optim_kwargs +) where {T} + # TODO: reimplement symmetrization using Manopt + # if isnothing(symmetrization) + # retract = peps_retract + # else + # retract, symm_finalize! = symmetrize_retract_and_finalize!(symmetrization) + # fin! = finalize! # Previous finalize! + # finalize! = (x, f, g, numiter) -> fin!(symm_finalize!(x, f, g, numiter)..., numiter) + # end if scalartype(env₀) <: Real env₀ = complex(env₀) @@ -164,37 +179,96 @@ function fixedpoint( with purely real environments" end - (peps, env), E, ∂E, numfg, convhistory = optimize( - (ψ₀, env₀), alg.optimizer; retract, inner=real_inner, finalize! - ) do (peps, envs) - E, gs = withgradient(peps) do ψ - envs´ = hook_pullback( - leading_boundary, - envs, - ψ, - alg.boundary_alg; - alg_rrule=alg.gradient_alg, - ) - ignore_derivatives() do - alg.reuse_env && update!(envs, envs´) - end - return costfun(ψ, envs´, H) - end - g = only(gs) # `withgradient` returns tuple of gradients `gs` - return E, g - end + peps₀_vec, from_vec = to_vec(peps₀) + pec = PEPSEnergyCost(env₀, H, alg, from_vec) - return (; - peps, - env, - E, - E_history=convhistory[:, 1], - grad=∂E, - gradnorm_history=convhistory[:, 2], - numfg, - ) + fld = scalartype(peps₀) <: Real ? Manifolds.ℝ : Manifolds.ℂ + cost_and_grad = ManifoldCostGradientObjective(pec) + result = alg.optimizer(Euclidean(; field=fld), cost_and_grad, peps₀_vec; optim_kwargs...) + + # TODO: extract quantities from result + + return peps, env, E, info end +# """ +# fixedpoint(ψ₀::InfinitePEPS{T}, H, alg::PEPSOptimize, [env₀::CTMRGEnv]; +# finalize!=OptimKit._finalize!, symmetrization=nothing) where {T} + +# Optimize `ψ₀` with respect to the Hamiltonian `H` according to the parameters supplied +# in `alg`. The initial environment `env₀` serves as an initial guess for the first CTMRG run. +# By default, a random initial environment is used. + +# The `finalize!` kwarg can be used to insert a function call after each optimization step +# by utilizing the `finalize!` kwarg of `OptimKit.optimize`. +# The function maps `(peps, envs), f, g = finalize!((peps, envs), f, g, numiter)`. +# The `symmetrization` kwarg accepts `nothing` or a `SymmetrizationStyle`, in which case the +# PEPS and PEPS gradient are symmetrized after each optimization iteration. Note that this +# requires a symmmetric `ψ₀` and `env₀` to converge properly. + +# The function returns a `NamedTuple` which contains the following entries: +# - `peps`: final `InfinitePEPS` +# - `env`: `CTMRGEnv` corresponding to the final PEPS +# - `E`: final energy +# - `E_history`: convergence history of the energy function +# - `grad`: final energy gradient +# - `gradnorm_history`: convergence history of the energy gradient norms +# - `numfg`: total number of calls to the energy function +# """ +# function fixedpoint( +# ψ₀::InfinitePEPS{F}, +# H, +# alg::PEPSOptimize, +# env₀::CTMRGEnv=CTMRGEnv(ψ₀, field(F)^20); +# (finalize!)=OptimKit._finalize!, +# symmetrization=nothing, +# ) where {F} +# if isnothing(symmetrization) +# retract = peps_retract +# else +# retract, symm_finalize! = symmetrize_retract_and_finalize!(symmetrization) +# fin! = finalize! # Previous finalize! +# finalize! = (x, f, g, numiter) -> fin!(symm_finalize!(x, f, g, numiter)..., numiter) +# end + +# if scalartype(env₀) <: Real +# env₀ = complex(env₀) +# @warn "the provided real environment was converted to a complex environment since \ +# :fixed mode generally produces complex gauges; use :diffgauge mode instead to work \ +# with purely real environments" +# end + +# (peps, env), E, ∂E, numfg, convhistory = optimize( +# (ψ₀, env₀), alg.optimizer; retract, inner=real_inner, finalize! +# ) do (peps, envs) +# E, gs = withgradient(peps) do ψ +# envs´ = hook_pullback( +# leading_boundary, +# envs, +# ψ, +# alg.boundary_alg; +# alg_rrule=alg.gradient_alg, +# ) +# ignore_derivatives() do +# alg.reuse_env && update!(envs, envs´) +# end +# return costfun(ψ, envs´, H) +# end +# g = only(gs) # `withgradient` returns tuple of gradients `gs` +# return E, g +# end + +# return (; +# peps, +# env, +# E, +# E_history=convhistory[:, 1], +# grad=∂E, +# gradnorm_history=convhistory[:, 2], +# numfg, +# ) +# end + # Update PEPS unit cell in non-mutating way # Note: Both x and η are InfinitePEPS during optimization function peps_retract(x, η, α) diff --git a/src/algorithms/toolbox.jl b/src/algorithms/toolbox.jl index b89a072a8..44a0fe3c9 100644 --- a/src/algorithms/toolbox.jl +++ b/src/algorithms/toolbox.jl @@ -7,15 +7,6 @@ function MPSKit.expectation_value(peps::InfinitePEPS, O::LocalOperator, envs::CT return sum(term_vals) end -function costfun(peps::InfinitePEPS, envs::CTMRGEnv, O::LocalOperator) - E = MPSKit.expectation_value(peps, O, envs) - ignore_derivatives() do - isapprox(imag(E), 0; atol=sqrt(eps(real(E)))) || - @warn "Expectation value is not real: $E." - end - return real(E) -end - function LinearAlgebra.norm(peps::InfinitePEPS, env::CTMRGEnv) total = one(scalartype(peps)) From db9821ab9f7d80615ce7072080641bc59673e27b Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Mon, 16 Dec 2024 18:34:42 +0100 Subject: [PATCH 02/28] Add new RecordActions and return info in leading_boundary --- src/PEPSKit.jl | 21 ++-- src/algorithms/ctmrg/ctmrg.jl | 8 +- src/algorithms/peps_opt.jl | 158 ++++++++++--------------- src/environments/ctmrg_environments.jl | 19 +++ 4 files changed, 102 insertions(+), 104 deletions(-) diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 17e85914f..ea9a50feb 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -6,6 +6,7 @@ using Compat using Accessors: @set using VectorInterface using TensorKit, KrylovKit, MPSKit, TensorOperations +using TensorKit: ℂ, ℝ # To avoid conflict with Manifolds using ChainRulesCore, Zygote using LoggingExtras using MPSKit: loginit!, logiter!, logfinish!, logcancel! @@ -53,10 +54,10 @@ include("algorithms/time_evolution/simpleupdate.jl") include("algorithms/toolbox.jl") -include("algorithms/peps_opt.jl") - include("utility/symmetrization.jl") +include("algorithms/peps_opt.jl") + """ module Defaults # CTMRG @@ -83,6 +84,7 @@ include("utility/symmetrization.jl") RecordCostUnitCell(), RecordTime(), ] + const optim_kwargs = (; memory_size=32, stopping_criterion=StopAfterIteration(100) | StopWhenGradientNormLess(1e-4), @@ -129,7 +131,9 @@ module Defaults FixedSpaceTruncation, SVDAdjoint, HalfInfiniteProjector, - SimultaneousCTMRG + SimultaneousCTMRG, + RecordConditionNumber, + RecordUnitCellGradientNorm # CTMRG const ctmrg_tol = 1e-8 @@ -149,12 +153,13 @@ module Defaults # Optimization const optim_alg = quasi_Newton const record_group = [ - RecordCost(), - RecordGradientNorm(), - RecordConditionNumber(), # TODO: implement PEPS record actions - RecordCostUnitCell(), - RecordTime(), + RecordCost() => :cost, + RecordGradientNorm() => :gradient_norm, + RecordConditionNumber() => :condition, # TODO: implement PEPS record actions + RecordUnitCellGradientNorm() => :unitcell_gradient_norm, + RecordTime() => :time, ] + const stopping_criterion = StopAfterIteration(100) | StopWhenGradientNormLess(1e-4) const optim_kwargs = (; memory_size=32, stopping_criterion=StopAfterIteration(100) | StopWhenGradientNormLess(1e-4), diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index ac5444f21..f1799bfba 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -40,10 +40,14 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRGAlgorithm) env = deepcopy(envinit) log = ignore_derivatives(() -> MPSKit.IterLog("CTMRG")) + truncation_error = 0.0 + condition_number = 1.0 return LoggingExtras.withlevel(; alg.verbosity) do ctmrg_loginit!(log, η, N) for iter in 1:(alg.maxiter) - env, = ctmrg_iteration(state, env, alg) # Grow and renormalize in all 4 directions + env, info = ctmrg_iteration(state, env, alg) # Grow and renormalize in all 4 directions + truncation_error = info.truncation_error + condition_number = info.condition_number η, CS, TS = calc_convergence(env, CS, TS) N = norm(state, env) @@ -57,7 +61,7 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRGAlgorithm) ctmrg_logiter!(log, iter, η, N) end end - return env + return env, (; truncation_error, condition_number) end end diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl index 3a52d0636..880e445e6 100644 --- a/src/algorithms/peps_opt.jl +++ b/src/algorithms/peps_opt.jl @@ -78,6 +78,40 @@ function LinSolver(; return LinSolver{iterscheme}(solver) end +mutable struct RecordTruncationError + recorded_values::Vector{Float64} + RecordTruncationError() = new(Vector{Float64}()) +end +function (r::RecordTruncationError)( + p::AbstractManoptProblem, ::AbstractManoptSolverState, i::Int +) + pec = Manopt.get_cost_function(get_objective(p)) + return Manopt.record_or_reset!(r, pec.env_info.truncation_error, i) +end + +mutable struct RecordConditionNumber + recorded_values::Vector{Float64} + RecordConditionNumber() = new(Vector{Float64}()) +end +function (r::RecordConditionNumber)( + p::AbstractManoptProblem, ::AbstractManoptSolverState, i::Int +) + pec = Manopt.get_cost_function(get_objective(p)) + return Manopt.record_or_reset!(r, pec.env_info.condition_number, i) +end + +mutable struct RecordUnitCellGradientNorm + recorded_values::Matrix{Float64} + RecordUnitCellGradientNorm() = new(Vector{Float64}()) +end +function (r::RecordUnitCellGradientNorm)( + p::AbstractManoptProblem, s::AbstractManoptSolverState, i::Int +) + pec = Manopt.get_cost_function(get_objective(p)) + grad = pec.from_vec(get_gradient(s)) + return Manopt.record_or_reset!(r, norm.(grad.A), i) +end + """ PEPSOptimize{G}(; boundary_alg=Defaults.ctmrg_alg, optimizer::OptimKit.OptimizationAlgorithm=Defaults.optimizer reuse_env::Bool=true, gradient_alg::G=Defaults.gradient_alg) @@ -93,8 +127,11 @@ step. The CTMRG gradient itself is computed using the `gradient_alg` algorithm. struct PEPSOptimize{G} boundary_alg::CTMRGAlgorithm optim_alg - reuse_env::Bool + optim_kwargs gradient_alg::G + reuse_env::Bool + # reuse_env_tol::Float64 # TODO: add option for reuse tolerance + symmetrization::SymmetrizationStyle function PEPSOptimize( # Inner constructor to prohibit illegal setting combinations boundary_alg::CTMRGAlgorithm, @@ -119,29 +156,29 @@ function PEPSOptimize(; return PEPSOptimize(boundary_alg, optim_alg, reuse_env, gradient_alg) end -struct PEPSEnergyCost +mutable struct PEPSEnergyCost env::CTMRGEnv hamiltonian::LocalOperator alg::PEPSOptimize from_vec::Function + env_info::NamedTuple end function (pec::PEPSEnergyCost)(peps_vec::Vector) peps = pec.from_vec(peps_vec) + env₀ = reuse_env ? pec.env : CTMRGEnv(randn, scalartype(pec.env), env) E, gs = withgradient(peps) do ψ - env′ = hook_pullback( + env′, info = hook_pullback( leading_boundary, - pec.env, + env₀, ψ, pec.alg.boundary_alg; alg_rrule=pec.alg.gradient_alg, ) - ignore_derivatives() do - pec.alg.reuse_env && update!(pec.env, env′) - end - + pec.env_info = info E = expectation_value(peps, pec.hamiltonian, env′) ignore_derivatives() do + update!(pec.env, env′) # Update environment in-place isapprox(imag(E), 0; atol=sqrt(eps(real(E)))) || @warn "Expectation value is not real: $E." end @@ -159,9 +196,8 @@ function fixedpoint( H, alg::PEPSOptimize, env₀::CTMRGEnv=CTMRGEnv(peps₀, field(T)^20); - # (finalize!)=OptimKit._finalize!, - # symmetrization=nothing, - optim_kwargs=Defaults.optim_kwargs + stopping_criterion::StoppingCriterion=Defaults.stopping_criterion, + record=Defaults.record_group, ) where {T} # TODO: reimplement symmetrization using Manopt # if isnothing(symmetrization) @@ -179,96 +215,30 @@ function fixedpoint( with purely real environments" end + # construct cost function struct peps₀_vec, from_vec = to_vec(peps₀) - pec = PEPSEnergyCost(env₀, H, alg, from_vec) - + pec = PEPSEnergyCost(env₀, H, alg, from_vec, (;)) fld = scalartype(peps₀) <: Real ? Manifolds.ℝ : Manifolds.ℂ cost_and_grad = ManifoldCostGradientObjective(pec) - result = alg.optimizer(Euclidean(; field=fld), cost_and_grad, peps₀_vec; optim_kwargs...) - # TODO: extract quantities from result + # optimize + result = alg.optimizer( + Euclidean(; field=fld), + cost_and_grad, + peps₀_vec; + stopping_criterion, + record, + return_state=true, + alg.optim_kwargs..., + ) - return peps, env, E, info + # extract final result + peps_final = from_vec(get_solver_result(result)) + env_final = leading_boundary(pec.env, peps_final, alg.boundary_alg) + E_final = expectation_value(peps_final, H, env_final) + return peps_final, env_final, E_final, result end -# """ -# fixedpoint(ψ₀::InfinitePEPS{T}, H, alg::PEPSOptimize, [env₀::CTMRGEnv]; -# finalize!=OptimKit._finalize!, symmetrization=nothing) where {T} - -# Optimize `ψ₀` with respect to the Hamiltonian `H` according to the parameters supplied -# in `alg`. The initial environment `env₀` serves as an initial guess for the first CTMRG run. -# By default, a random initial environment is used. - -# The `finalize!` kwarg can be used to insert a function call after each optimization step -# by utilizing the `finalize!` kwarg of `OptimKit.optimize`. -# The function maps `(peps, envs), f, g = finalize!((peps, envs), f, g, numiter)`. -# The `symmetrization` kwarg accepts `nothing` or a `SymmetrizationStyle`, in which case the -# PEPS and PEPS gradient are symmetrized after each optimization iteration. Note that this -# requires a symmmetric `ψ₀` and `env₀` to converge properly. - -# The function returns a `NamedTuple` which contains the following entries: -# - `peps`: final `InfinitePEPS` -# - `env`: `CTMRGEnv` corresponding to the final PEPS -# - `E`: final energy -# - `E_history`: convergence history of the energy function -# - `grad`: final energy gradient -# - `gradnorm_history`: convergence history of the energy gradient norms -# - `numfg`: total number of calls to the energy function -# """ -# function fixedpoint( -# ψ₀::InfinitePEPS{F}, -# H, -# alg::PEPSOptimize, -# env₀::CTMRGEnv=CTMRGEnv(ψ₀, field(F)^20); -# (finalize!)=OptimKit._finalize!, -# symmetrization=nothing, -# ) where {F} -# if isnothing(symmetrization) -# retract = peps_retract -# else -# retract, symm_finalize! = symmetrize_retract_and_finalize!(symmetrization) -# fin! = finalize! # Previous finalize! -# finalize! = (x, f, g, numiter) -> fin!(symm_finalize!(x, f, g, numiter)..., numiter) -# end - -# if scalartype(env₀) <: Real -# env₀ = complex(env₀) -# @warn "the provided real environment was converted to a complex environment since \ -# :fixed mode generally produces complex gauges; use :diffgauge mode instead to work \ -# with purely real environments" -# end - -# (peps, env), E, ∂E, numfg, convhistory = optimize( -# (ψ₀, env₀), alg.optimizer; retract, inner=real_inner, finalize! -# ) do (peps, envs) -# E, gs = withgradient(peps) do ψ -# envs´ = hook_pullback( -# leading_boundary, -# envs, -# ψ, -# alg.boundary_alg; -# alg_rrule=alg.gradient_alg, -# ) -# ignore_derivatives() do -# alg.reuse_env && update!(envs, envs´) -# end -# return costfun(ψ, envs´, H) -# end -# g = only(gs) # `withgradient` returns tuple of gradients `gs` -# return E, g -# end - -# return (; -# peps, -# env, -# E, -# E_history=convhistory[:, 1], -# grad=∂E, -# gradnorm_history=convhistory[:, 2], -# numfg, -# ) -# end - # Update PEPS unit cell in non-mutating way # Note: Both x and η are InfinitePEPS during optimization function peps_retract(x, η, α) diff --git a/src/environments/ctmrg_environments.jl b/src/environments/ctmrg_environments.jl index 97d890de9..3fae1c367 100644 --- a/src/environments/ctmrg_environments.jl +++ b/src/environments/ctmrg_environments.jl @@ -330,6 +330,25 @@ function CTMRGEnv( end @non_differentiable CTMRGEnv(peps::InfinitePEPS, args...) +""" + CTMRGEnv([f=randn, T=scalartype(env)], env::CTMRGEnv) + +Copy constructor for easily initializing a similar environment, optionally providing a +function `f` for element generation with type `T`. +""" +function CTMRGEnv(env::CTMRGEnv) + return CTMRTEnv(randn, scalartype(env), env) +end +function CTMRGEnv(f, T, env::CTMRGEnv) + corners = map(env.corners) do corner + return TensorMap(f, T, space(corner)) + end + edges = map(env.edges) do edge + return TensorMap(f, T, space(edge)) + end + return CTMRGEnv(corners, edges) +end + # Custom adjoint for CTMRGEnv constructor, needed for fixed-point differentiation function ChainRulesCore.rrule(::Type{CTMRGEnv}, corners, edges) ctmrgenv_pullback(ē) = NoTangent(), ē.corners, ē.edges From 23b9e28a9fbeeeb2a0156adcb1f4ee97c6607c20 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Tue, 17 Dec 2024 17:23:34 +0100 Subject: [PATCH 03/28] Add symmetrization --- src/algorithms/ctmrg/sequential.jl | 22 +++++++---- src/algorithms/ctmrg/simultaneous.jl | 16 ++++++-- src/algorithms/peps_opt.jl | 57 ++++++++++++++++------------ src/utility/symmetrization.jl | 20 ---------- 4 files changed, 61 insertions(+), 54 deletions(-) diff --git a/src/algorithms/ctmrg/sequential.jl b/src/algorithms/ctmrg/sequential.jl index 94f4846e5..6ae2a8ee9 100644 --- a/src/algorithms/ctmrg/sequential.jl +++ b/src/algorithms/ctmrg/sequential.jl @@ -31,18 +31,20 @@ function SequentialCTMRG(; end function ctmrg_iteration(state, envs::CTMRGEnv, alg::SequentialCTMRG) - ϵ = zero(real(scalartype(state))) + truncation_error = zero(real(scalartype(state))) + condition_number = zero(real(scalartype(state))) for _ in 1:4 # rotate for col in 1:size(state, 2) # left move column-wise projectors, info = sequential_projectors(col, state, envs, alg.projector_alg) envs = renormalize_sequentially(col, projectors, state, envs) - ϵ = max(ϵ, info.err) + truncation_error = max(truncation_error, info.truncation_error) + condition_number = max(condition_number, info.condition_number) end state = rotate_north(state, EAST) envs = rotate_north(envs, EAST) end - return envs, (; err=ϵ) + return envs, (; truncation_error, condition_number) end """ @@ -53,10 +55,13 @@ Compute CTMRG projectors in the `:sequential` scheme either for an entire column for a specific `coordinate` (where `dir=WEST` is already implied in the `:sequential` scheme). """ function sequential_projectors( - col::Int, state::InfinitePEPS, envs::CTMRGEnv, alg::ProjectorAlgorithm -) + col::Int, state::InfinitePEPS, envs::CTMRGEnv{C,T}, alg::ProjectorAlgorithm +) where {C,T} # SVD half-infinite environment column-wise - ϵ = Zygote.Buffer(zeros(real(scalartype(envs)), size(envs, 2))) + ϵ = Zygote.Buffer(zeros(real(scalartype(T)), size(envs, 2))) + S = Zygote.Buffer( + zeros(size(envs, 2)), tensormaptype(spacetype(T), 1, 1, real(scalartype(T))) + ) coordinates = eachcoordinate(envs)[:, col] projectors = dtmap(coordinates) do (r, c) trscheme = truncation_scheme(alg, envs.edges[WEST, _prev(r, size(envs, 2)), c]) @@ -64,10 +69,13 @@ function sequential_projectors( (WEST, r, c), state, envs, @set(alg.trscheme = trscheme) ) ϵ[r] = info.err / norm(info.S) + S[r] = info.S return proj end - return (map(first, projectors), map(last, projectors)), (; err=maximum(copy(ϵ))) + condition_number = maximum(_condition_number, copy(S)) + info = (; truncation_error=maximum(copy(ϵ)), condition_number) + return (map(first, projectors), map(last, projectors)), info end function sequential_projectors( coordinate::NTuple{3,Int}, diff --git a/src/algorithms/ctmrg/simultaneous.jl b/src/algorithms/ctmrg/simultaneous.jl index fcae797bb..870922be4 100644 --- a/src/algorithms/ctmrg/simultaneous.jl +++ b/src/algorithms/ctmrg/simultaneous.jl @@ -44,7 +44,7 @@ function _prealloc_svd(edges::Array{E,N}, ::HalfInfiniteProjector) where {E,N} Sc = scalartype(E) U = Zygote.Buffer(map(e -> TensorMap(zeros, Sc, space(e)), edges)) V = Zygote.Buffer(map(e -> TensorMap(zeros, Sc, domain(e), codomain(e)), edges)) - S = Zygote.Buffer(U.data, tensormaptype(spacetype(E), 1, 1, real(Sc))) # Corner type but with real numbers + S = Zygote.Buffer(edges, tensormaptype(spacetype(E), 1, 1, real(Sc))) # Corner type but with real numbers return U, S, V end function _prealloc_svd(edges::Array{E,N}, ::FullInfiniteProjector) where {E,N} @@ -52,10 +52,17 @@ function _prealloc_svd(edges::Array{E,N}, ::FullInfiniteProjector) where {E,N} Rspace(x) = fuse(codomain(x)) U = Zygote.Buffer(map(e -> TensorMap(zeros, Sc, Rspace(e), domain(e)), edges)) V = Zygote.Buffer(map(e -> TensorMap(zeros, Sc, domain(e), Rspace(e)), edges)) - S = Zygote.Buffer(U.data, tensormaptype(spacetype(E), 1, 1, real(Sc))) # Corner type but with real numbers + S = Zygote.Buffer(edges, tensormaptype(spacetype(E), 1, 1, real(Sc))) # Corner type but with real numbers return U, S, V end +# Compute condition number σ_max / σ_min for diagonal singular value TensorMap +function _condition_number(S::AbstractTensorMap) + S_diag = diag(S.data) + return maximum(S_diag) / minimum(S_diag) +end +@non_differentiable _condition_number(S::AbstractTensorMap) + """ simultaneous_projectors(enlarged_corners::Array{E,3}, envs::CTMRGEnv, alg::ProjectorAlgorithm) simultaneous_projectors(coordinate, enlarged_corners::Array{E,3}, alg::ProjectorAlgorithm) @@ -84,7 +91,10 @@ function simultaneous_projectors( P_left = map(first, projectors) P_right = map(last, projectors) - return (P_left, P_right), (; err=maximum(copy(ϵ)), U=copy(U), S=copy(S), V=copy(V)) + S = copy(S) + condition_number = maximum(_condition_number, S) + info = (; truncation_error=maximum(copy(ϵ)), condition_number, U=copy(U), S, V=copy(V)) + return (P_left, P_right), info end function simultaneous_projectors( coordinate, enlarged_corners::Array{E,3}, alg::HalfInfiniteProjector diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl index 880e445e6..0a3854979 100644 --- a/src/algorithms/peps_opt.jl +++ b/src/algorithms/peps_opt.jl @@ -101,8 +101,8 @@ function (r::RecordConditionNumber)( end mutable struct RecordUnitCellGradientNorm - recorded_values::Matrix{Float64} - RecordUnitCellGradientNorm() = new(Vector{Float64}()) + recorded_values::Vector{Matrix{Float64}} + RecordUnitCellGradientNorm() = new(Vector{Matrix{Float64}}()) end function (r::RecordUnitCellGradientNorm)( p::AbstractManoptProblem, s::AbstractManoptSolverState, i::Int @@ -167,6 +167,8 @@ end function (pec::PEPSEnergyCost)(peps_vec::Vector) peps = pec.from_vec(peps_vec) env₀ = reuse_env ? pec.env : CTMRGEnv(randn, scalartype(pec.env), env) + + # compute cost and gradient E, gs = withgradient(peps) do ψ env′, info = hook_pullback( leading_boundary, @@ -185,9 +187,30 @@ function (pec::PEPSEnergyCost)(peps_vec::Vector) return real(E) end g = only(gs) # `withgradient` returns tuple of gradients `gs` + + # symmetrize gradient + if !isnothing(pec.alg.symmetrization) + g = symmetrize!(g, pec.alg.symmetrization) + end + return E, to_vec(g)[1] end +# First retract and then symmetrize the resulting PEPS +# (ExponentialRetraction is the default retraction for Euclidean manifolds) +struct SymmetrizeExponentialRetraction <: AbstractRetractionMethod + symmetrization::SymmetrizationStyle + from_vec::Function +end + +function Manifolds.retract( + M::AbstractManifold, p, X, t::Number, sr::SymmetrizeExponentialRetraction +) + q = retract(M, p, X, t, ExponentialRetraction()) + q_symm_peps = symmetrize!(sr.from_vec(q), sr.symmetrization) + return to_vec(q_symm_peps) +end + """ TODO """ @@ -199,15 +222,6 @@ function fixedpoint( stopping_criterion::StoppingCriterion=Defaults.stopping_criterion, record=Defaults.record_group, ) where {T} - # TODO: reimplement symmetrization using Manopt - # if isnothing(symmetrization) - # retract = peps_retract - # else - # retract, symm_finalize! = symmetrize_retract_and_finalize!(symmetrization) - # fin! = finalize! # Previous finalize! - # finalize! = (x, f, g, numiter) -> fin!(symm_finalize!(x, f, g, numiter)..., numiter) - # end - if scalartype(env₀) <: Real env₀ = complex(env₀) @warn "the provided real environment was converted to a complex environment since \ @@ -222,13 +236,20 @@ function fixedpoint( cost_and_grad = ManifoldCostGradientObjective(pec) # optimize + M = Euclidean(; field=fld) + retraction_method = if isnothing(alg.symmetrization) + default_retraction_method(M) + else + SymmetrizeExponentialRetraction(alg.symmetrization, from_vec) + end result = alg.optimizer( - Euclidean(; field=fld), + M, cost_and_grad, peps₀_vec; stopping_criterion, record, return_state=true, + retraction_method, alg.optim_kwargs..., ) @@ -239,18 +260,6 @@ function fixedpoint( return peps_final, env_final, E_final, result end -# Update PEPS unit cell in non-mutating way -# Note: Both x and η are InfinitePEPS during optimization -function peps_retract(x, η, α) - peps = deepcopy(x[1]) - peps.A .+= η.A .* α - env = deepcopy(x[2]) - return (peps, env), η -end - -# Take real valued part of dot product -real_inner(_, η₁, η₂) = real(dot(η₁, η₂)) - #= Evaluating the gradient of the cost function for CTMRG: - The gradient of the cost function for CTMRG can be computed using automatic differentiation (AD) or explicit evaluation of the geometric sum. diff --git a/src/utility/symmetrization.jl b/src/utility/symmetrization.jl index 5bec6a961..abfdb4744 100644 --- a/src/utility/symmetrization.jl +++ b/src/utility/symmetrization.jl @@ -216,23 +216,3 @@ function symmetrize!(peps::InfinitePEPS, symm::RotateReflect) end return peps end - -""" - symmetrize_retract_and_finalize!(symm::SymmetrizationStyle) - -Return the `retract` and `finalize!` function for symmetrizing the `peps` and `grad` tensors. -""" -function symmetrize_retract_and_finalize!(symm::SymmetrizationStyle) - finf = function symmetrize_finalize!((peps, envs), E, grad, _) - grad_symm = symmetrize!(grad, symm) - return (peps, envs), E, grad_symm - end - retf = function symmetrize_retract((peps, envs), η, α) - peps_symm = deepcopy(peps) - peps_symm.A .+= η.A .* α - e = deepcopy(envs) - symmetrize!(peps_symm, symm) - return (peps_symm, e), η - end - return retf, finf -end From ce68d55223ba6bfa318588438df6b4a42c811e30 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Tue, 17 Dec 2024 18:26:18 +0100 Subject: [PATCH 04/28] Remove OptimKit, fix PEPSOptimize constructors --- Project.toml | 2 -- examples/heisenberg.jl | 4 ++-- src/PEPSKit.jl | 7 +++++-- src/algorithms/peps_opt.jl | 43 +++++++++++++++++++++++--------------- src/states/infinitepeps.jl | 2 -- test/ctmrg/gradients.jl | 36 +++++++++++++++---------------- test/j1j2_model.jl | 3 +-- test/pwave.jl | 3 +-- test/tf_ising.jl | 3 +-- 9 files changed, 54 insertions(+), 49 deletions(-) diff --git a/Project.toml b/Project.toml index 3d0a77b60..0deecf910 100644 --- a/Project.toml +++ b/Project.toml @@ -16,7 +16,6 @@ MPSKitModels = "ca635005-6f8c-4cd1-b51d-8491250ef2ab" Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e" Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5" OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5" -OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" @@ -38,7 +37,6 @@ MPSKitModels = "0.3.5" Manifolds = "0.10.8" Manopt = "0.5.4" OhMyThreads = "0.7" -OptimKit = "0.3" Printf = "1" Random = "1" Statistics = "1" diff --git a/examples/heisenberg.jl b/examples/heisenberg.jl index 34baed189..98739780b 100644 --- a/examples/heisenberg.jl +++ b/examples/heisenberg.jl @@ -1,6 +1,6 @@ using LinearAlgebra -using TensorKit, OptimKit -using PEPSKit, KrylovKit +using TensorKit, KrylovKit, PEPSKit +using Manopt # Square lattice Heisenberg Hamiltonian # We use the parameters (J₁, J₂, J₃) = (-1, 1, -1) by default to capture diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index ea9a50feb..c82dd7313 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -124,7 +124,7 @@ Module containing default values that represent typical algorithm parameters. - `scheduler`: Multi-threading scheduler which can be accessed via `set_scheduler!` """ module Defaults - using TensorKit, KrylovKit, OptimKit, OhMyThreads + using TensorKit, KrylovKit, OhMyThreads using Manopt using PEPSKit: LinSolver, @@ -160,9 +160,12 @@ module Defaults RecordTime() => :time, ] const stopping_criterion = StopAfterIteration(100) | StopWhenGradientNormLess(1e-4) + const optim_maxiter = 100 + const optim_tol = 1e-4 const optim_kwargs = (; memory_size=32, - stopping_criterion=StopAfterIteration(100) | StopWhenGradientNormLess(1e-4), + stopping_criterion=StopAfterIteration(optim_maxiter) | + StopWhenGradientNormLess(optim_tol), record=record_group, return_state=true, ) diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl index 0a3854979..72cf40fed 100644 --- a/src/algorithms/peps_opt.jl +++ b/src/algorithms/peps_opt.jl @@ -113,8 +113,7 @@ function (r::RecordUnitCellGradientNorm)( end """ - PEPSOptimize{G}(; boundary_alg=Defaults.ctmrg_alg, optimizer::OptimKit.OptimizationAlgorithm=Defaults.optimizer - reuse_env::Bool=true, gradient_alg::G=Defaults.gradient_alg) +TODO Algorithm struct that represent PEPS ground-state optimization using AD. Set the algorithm to contract the infinite PEPS in `boundary_alg`; @@ -126,34 +125,45 @@ step. The CTMRG gradient itself is computed using the `gradient_alg` algorithm. """ struct PEPSOptimize{G} boundary_alg::CTMRGAlgorithm - optim_alg - optim_kwargs + optim_alg::Function + optim_kwargs::NamedTuple gradient_alg::G reuse_env::Bool # reuse_env_tol::Float64 # TODO: add option for reuse tolerance - symmetrization::SymmetrizationStyle + symmetrization::Union{Nothing,SymmetrizationStyle} function PEPSOptimize( # Inner constructor to prohibit illegal setting combinations boundary_alg::CTMRGAlgorithm, optim_alg, - reuse_env, + optim_kwargs, gradient_alg::G, + reuse_env, + symmetrization, ) where {G} if gradient_alg isa GradMode if boundary_alg isa SequentialCTMRG && iterscheme(gradient_alg) === :fixed throw(ArgumentError(":sequential and :fixed are not compatible")) end end - return new{G}(boundary_alg, optim_alg, reuse_env, gradient_alg) + return new{G}( + boundary_alg, optim_alg, optim_kwargs, gradient_alg, reuse_env, symmetrization + ) end end function PEPSOptimize(; boundary_alg=Defaults.ctmrg_alg, optim_alg=Defaults.optim_alg, - reuse_env=Defaults.reuse_env, + maxiter=Defaults.optim_maxiter, + tol=Defaults.optim_tol, gradient_alg=Defaults.gradient_alg, + reuse_env=Defaults.reuse_env, + symmetrization=nothing, ) - return PEPSOptimize(boundary_alg, optim_alg, reuse_env, gradient_alg) + stopping_criterion = StopAfterIteration(maxiter) | StopWhenGradientNormLess(tol) + optim_kwargs = (; stopping_criterion, record=Defaults.record_group) + return PEPSOptimize( + boundary_alg, optim_alg, optim_kwargs, gradient_alg, reuse_env, symmetrization + ) end mutable struct PEPSEnergyCost @@ -164,6 +174,7 @@ mutable struct PEPSEnergyCost env_info::NamedTuple end +# TODO: split this up into f and grad_f function (pec::PEPSEnergyCost)(peps_vec::Vector) peps = pec.from_vec(peps_vec) env₀ = reuse_env ? pec.env : CTMRGEnv(randn, scalartype(pec.env), env) @@ -219,8 +230,6 @@ function fixedpoint( H, alg::PEPSOptimize, env₀::CTMRGEnv=CTMRGEnv(peps₀, field(T)^20); - stopping_criterion::StoppingCriterion=Defaults.stopping_criterion, - record=Defaults.record_group, ) where {T} if scalartype(env₀) <: Real env₀ = complex(env₀) @@ -233,7 +242,8 @@ function fixedpoint( peps₀_vec, from_vec = to_vec(peps₀) pec = PEPSEnergyCost(env₀, H, alg, from_vec, (;)) fld = scalartype(peps₀) <: Real ? Manifolds.ℝ : Manifolds.ℂ - cost_and_grad = ManifoldCostGradientObjective(pec) + cost = # TODO #ManifoldCostGradientObjective(pec) + grad = # TODO # optimize M = Euclidean(; field=fld) @@ -242,15 +252,14 @@ function fixedpoint( else SymmetrizeExponentialRetraction(alg.symmetrization, from_vec) end - result = alg.optimizer( + result = alg.optim_alg( M, - cost_and_grad, + cost, + grad, peps₀_vec; - stopping_criterion, - record, + alg.optim_kwargs..., return_state=true, retraction_method, - alg.optim_kwargs..., ) # extract final result diff --git a/src/states/infinitepeps.jl b/src/states/infinitepeps.jl index fea4ab1fd..846ad707c 100644 --- a/src/states/infinitepeps.jl +++ b/src/states/infinitepeps.jl @@ -153,13 +153,11 @@ function Base.isapprox(ψ₁::InfinitePEPS, ψ₂::InfinitePEPS; kwargs...) end end -# Used in _scale during OptimKit.optimize function LinearAlgebra.rmul!(ψ::InfinitePEPS, α::Number) rmul!(ψ.A, α) return ψ end -# Used in _add during OptimKit.optimize function LinearAlgebra.axpy!(α::Number, ψ₁::InfinitePEPS, ψ₂::InfinitePEPS) ψ₂.A .+= α * ψ₁.A return ψ₂ diff --git a/test/ctmrg/gradients.jl b/test/ctmrg/gradients.jl index 648d48bdb..5659aa6d3 100644 --- a/test/ctmrg/gradients.jl +++ b/test/ctmrg/gradients.jl @@ -1,10 +1,9 @@ using Test using Random -using PEPSKit using TensorKit -using Zygote -using OptimKit using KrylovKit +using PEPSKit +using Zygote ## Test models, gradmodes and CTMRG algorithm # ------------------------------------------- @@ -57,25 +56,26 @@ steps = -0.01:0.005:0.01 psi_init = InfinitePEPS(Pspace, Vspace, Vspace) @testset "$ctmrg_alg and $alg_rrule" for (ctmrg_alg, alg_rrule) in Iterators.product(calgs, gms) - @info "optimtest of $ctmrg_alg and $alg_rrule on $(names[i])" + # @info "optimtest of $ctmrg_alg and $alg_rrule on $(names[i])" Random.seed!(42039482030) dir = InfinitePEPS(Pspace, Vspace, Vspace) psi = InfinitePEPS(Pspace, Vspace, Vspace) env = leading_boundary(CTMRGEnv(psi, Espace), psi, ctmrg_alg) - alphas, fs, dfs1, dfs2 = OptimKit.optimtest( - (psi, env), - dir; - alpha=steps, - retract=PEPSKit.peps_retract, - inner=PEPSKit.real_inner, - ) do (peps, envs) - E, g = Zygote.withgradient(peps) do psi - envs2 = PEPSKit.hook_pullback(leading_boundary, envs, psi, ctmrg_alg; alg_rrule) - return costfun(psi, envs2, models[i]) - end + # TODO: redo this test using Manopt + # alphas, fs, dfs1, dfs2 = OptimKit.optimtest( + # (psi, env), + # dir; + # alpha=steps, + # retract=PEPSKit.peps_retract, + # inner=PEPSKit.real_inner, + # ) do (peps, envs) + # E, g = Zygote.withgradient(peps) do psi + # envs2 = PEPSKit.hook_pullback(leading_boundary, envs, psi, ctmrg_alg; alg_rrule) + # return costfun(psi, envs2, models[i]) + # end - return E, only(g) - end - @test dfs1 ≈ dfs2 atol = 1e-2 + # return E, only(g) + # end + # @test dfs1 ≈ dfs2 atol = 1e-2 end end diff --git a/test/j1j2_model.jl b/test/j1j2_model.jl index 8dd883972..3df78f881 100644 --- a/test/j1j2_model.jl +++ b/test/j1j2_model.jl @@ -1,9 +1,8 @@ using Test using Random -using PEPSKit using TensorKit using KrylovKit -using OptimKit +using PEPSKit # initialize parameters χbond = 2 diff --git a/test/pwave.jl b/test/pwave.jl index dcdecf7ea..1007b05ac 100644 --- a/test/pwave.jl +++ b/test/pwave.jl @@ -1,9 +1,8 @@ using Test using Random -using PEPSKit using TensorKit using KrylovKit -using OptimKit +using PEPSKit # Initialize parameters unitcell = (2, 2) diff --git a/test/tf_ising.jl b/test/tf_ising.jl index 8b84696f2..90f21b1b4 100644 --- a/test/tf_ising.jl +++ b/test/tf_ising.jl @@ -1,9 +1,8 @@ using Test using Random -using PEPSKit using TensorKit using KrylovKit -using OptimKit +using PEPSKit # References # ---------- From 182210483a45c92d0992ac77cbd1d30762bc7ac6 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Tue, 17 Dec 2024 18:29:41 +0100 Subject: [PATCH 05/28] Fix OhMyThreads precompilation warning --- src/PEPSKit.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index c82dd7313..f1a7d5064 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -12,7 +12,7 @@ using LoggingExtras using MPSKit: loginit!, logiter!, logfinish!, logcancel! using MPSKitModels using FiniteDifferences -using OhMyThreads +using OhMyThreads: tmap include("utility/util.jl") include("utility/diffable_threads.jl") From 9e2b4d6191164168153fbaae434329638ec2db8a Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Wed, 18 Dec 2024 16:04:34 +0100 Subject: [PATCH 06/28] Update cost function cache --- src/algorithms/peps_opt.jl | 108 ++++++++++++++++++++++++------------- test/heisenberg.jl | 8 ++- 2 files changed, 73 insertions(+), 43 deletions(-) diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl index 72cf40fed..65b3e6a35 100644 --- a/src/algorithms/peps_opt.jl +++ b/src/algorithms/peps_opt.jl @@ -166,45 +166,85 @@ function PEPSOptimize(; ) end -mutable struct PEPSEnergyCost - env::CTMRGEnv - hamiltonian::LocalOperator +mutable struct PEPSCostFunctionCache{T} + operator::LocalOperator alg::PEPSOptimize - from_vec::Function + env::CTMRGEnv + from_vec + peps_vec::Vector{T} + grad_vec::Vector{T} + cost::Float64 env_info::NamedTuple end -# TODO: split this up into f and grad_f -function (pec::PEPSEnergyCost)(peps_vec::Vector) - peps = pec.from_vec(peps_vec) - env₀ = reuse_env ? pec.env : CTMRGEnv(randn, scalartype(pec.env), env) +function PEPSCostFunctionCache( + operator::LocalOperator, alg::PEPSOptimize, peps_vec::Vector, from_vec, env::CTMRGEnv +) + return PEPSCostFunctionCache( + operator, + alg, + env, + from_vec, + peps_vec, + similar(peps_vec), + 0.0, + (; truncation_error=0.0, condition_number=1.0), + ) +end + +function cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) where {T} + cache.peps_vec .= peps_vec # update point in manifold + peps = cache.from_vec(peps_vec) # convert back to InfinitePEPS + env₀ = + cache.alg.reuse_env ? cache.env : CTMRGEnv(randn, scalartype(cache.env), cache.env) # compute cost and gradient - E, gs = withgradient(peps) do ψ - env′, info = hook_pullback( + cost, grads = withgradient(peps) do ψ + env, info = hook_pullback( leading_boundary, env₀, ψ, - pec.alg.boundary_alg; - alg_rrule=pec.alg.gradient_alg, + cache.alg.boundary_alg; + alg_rrule=cache.alg.gradient_alg, ) - pec.env_info = info - E = expectation_value(peps, pec.hamiltonian, env′) + cache.env_info = info + cost = expectation_value(peps, cache.operator, env) ignore_derivatives() do - update!(pec.env, env′) # Update environment in-place - isapprox(imag(E), 0; atol=sqrt(eps(real(E)))) || - @warn "Expectation value is not real: $E." + update!(cache.env, env) # update environment in-place + isapprox(imag(cost), 0; atol=sqrt(eps(real(cost)))) || + @warn "Expectation value is not real: $cost." end - return real(E) + return real(cost) end - g = only(gs) # `withgradient` returns tuple of gradients `gs` + grad = only(grads) # `withgradient` returns tuple of gradients `grads` # symmetrize gradient - if !isnothing(pec.alg.symmetrization) - g = symmetrize!(g, pec.alg.symmetrization) + if !isnothing(cache.alg.symmetrization) + grad = symmetrize!(grad, cache.alg.symmetrization) end - return E, to_vec(g)[1] + # update cache + cache.cost = cost + cache.grad_vec .= to_vec(grad)[1] + return cache.cost, cache.grad_vec +end + +function cost_function(cache::PEPSCostFunctionCache{T}) + return function cost_function(::Euclidean, peps_vec::Vector{T}) + if !(peps_vec == cache.peps_vec) # update cache + cost_and_grad!(cache, peps_vec) + end + return cache.cost + end +end + +function grad_function(cache::PEPSCostFunctionCache) + return function grad_function(::Euclidean, peps_vec::Vector{T}) + if !(peps_vec == cache.peps_vec) # update cache + cost_and_grad!(cache, peps_vec) + end + return cache.grad_vec + end end # First retract and then symmetrize the resulting PEPS @@ -227,7 +267,7 @@ TODO """ function fixedpoint( peps₀::InfinitePEPS{T}, - H, + operator::LocalOperator, alg::PEPSOptimize, env₀::CTMRGEnv=CTMRGEnv(peps₀, field(T)^20); ) where {T} @@ -238,14 +278,14 @@ function fixedpoint( with purely real environments" end - # construct cost function struct + # construct cost and grad functions peps₀_vec, from_vec = to_vec(peps₀) - pec = PEPSEnergyCost(env₀, H, alg, from_vec, (;)) - fld = scalartype(peps₀) <: Real ? Manifolds.ℝ : Manifolds.ℂ - cost = # TODO #ManifoldCostGradientObjective(pec) - grad = # TODO + cache = PEPSCostFunctionCache(operator, alg, peps_vec, from_vec, env₀) + cost = cost_function(cache) + grad = grad_function(cache) # optimize + fld = scalartype(peps₀) <: Real ? Manifolds.ℝ : Manifolds.ℂ M = Euclidean(; field=fld) retraction_method = if isnothing(alg.symmetrization) default_retraction_method(M) @@ -253,20 +293,12 @@ function fixedpoint( SymmetrizeExponentialRetraction(alg.symmetrization, from_vec) end result = alg.optim_alg( - M, - cost, - grad, - peps₀_vec; - alg.optim_kwargs..., - return_state=true, - retraction_method, + M, cost, grad, peps₀_vec; alg.optim_kwargs..., return_state=true, retraction_method ) # extract final result peps_final = from_vec(get_solver_result(result)) - env_final = leading_boundary(pec.env, peps_final, alg.boundary_alg) - E_final = expectation_value(peps_final, H, env_final) - return peps_final, env_final, E_final, result + return peps_final, cost.env, cache.cost, result end #= diff --git a/test/heisenberg.jl b/test/heisenberg.jl index 5865ba468..a78e98af0 100644 --- a/test/heisenberg.jl +++ b/test/heisenberg.jl @@ -1,18 +1,16 @@ using Test using Random using Accessors -using PEPSKit using TensorKit using KrylovKit -using OptimKit +using PEPSKit +using Manopt # initialize parameters Dbond = 2 χenv = 16 ctm_alg = SimultaneousCTMRG() -opt_alg = PEPSOptimize(; - boundary_alg=ctm_alg, optimizer=LBFGS(4; gradtol=1e-3, verbosity=2) -) +opt_alg = PEPSOptimize(; boundary_alg=ctm_alg, tol=1e-3) # compare against Juraj Hasik's data: # https://github.com/jurajHasik/j1j2_ipeps_states/blob/main/single-site_pg-C4v-A1/j20.0/state_1s_A1_j20.0_D2_chi_opt48.dat E_ref = -0.6602310934799577 From ca54e0bea940d7b06c3bf2af5ac8d3a3c09e9082 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Wed, 18 Dec 2024 17:23:51 +0100 Subject: [PATCH 07/28] Make fixedpoint runnable --- src/algorithms/peps_opt.jl | 55 ++++++++++++++++++-------------------- src/states/infinitepeps.jl | 4 +++ test/heisenberg.jl | 7 ++++- 3 files changed, 36 insertions(+), 30 deletions(-) diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl index 65b3e6a35..1c67d287b 100644 --- a/src/algorithms/peps_opt.jl +++ b/src/algorithms/peps_opt.jl @@ -78,37 +78,37 @@ function LinSolver(; return LinSolver{iterscheme}(solver) end -mutable struct RecordTruncationError +mutable struct RecordTruncationError <: RecordAction recorded_values::Vector{Float64} RecordTruncationError() = new(Vector{Float64}()) end function (r::RecordTruncationError)( p::AbstractManoptProblem, ::AbstractManoptSolverState, i::Int ) - pec = Manopt.get_cost_function(get_objective(p)) - return Manopt.record_or_reset!(r, pec.env_info.truncation_error, i) + cache = Manopt.get_cost_function(get_objective(p)) + return Manopt.record_or_reset!(r, cache.env_info.truncation_error, i) end -mutable struct RecordConditionNumber +mutable struct RecordConditionNumber <: RecordAction recorded_values::Vector{Float64} RecordConditionNumber() = new(Vector{Float64}()) end function (r::RecordConditionNumber)( p::AbstractManoptProblem, ::AbstractManoptSolverState, i::Int ) - pec = Manopt.get_cost_function(get_objective(p)) - return Manopt.record_or_reset!(r, pec.env_info.condition_number, i) + cache = Manopt.get_cost_function(get_objective(p)) + return Manopt.record_or_reset!(r, cache.env_info.condition_number, i) end -mutable struct RecordUnitCellGradientNorm +mutable struct RecordUnitCellGradientNorm <: RecordAction recorded_values::Vector{Matrix{Float64}} RecordUnitCellGradientNorm() = new(Vector{Matrix{Float64}}()) end function (r::RecordUnitCellGradientNorm)( p::AbstractManoptProblem, s::AbstractManoptSolverState, i::Int ) - pec = Manopt.get_cost_function(get_objective(p)) - grad = pec.from_vec(get_gradient(s)) + cache = Manopt.get_cost_function(get_objective(p)) + grad = cache.from_vec(get_gradient(s)) return Manopt.record_or_reset!(r, norm.(grad.A), i) end @@ -176,7 +176,6 @@ mutable struct PEPSCostFunctionCache{T} cost::Float64 env_info::NamedTuple end - function PEPSCostFunctionCache( operator::LocalOperator, alg::PEPSOptimize, peps_vec::Vector, from_vec, env::CTMRGEnv ) @@ -207,7 +206,7 @@ function cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) wh cache.alg.boundary_alg; alg_rrule=cache.alg.gradient_alg, ) - cache.env_info = info + cache.env_info = info # update environment information (truncation error, ...) cost = expectation_value(peps, cache.operator, env) ignore_derivatives() do update!(cache.env, env) # update environment in-place @@ -223,24 +222,22 @@ function cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) wh grad = symmetrize!(grad, cache.alg.symmetrization) end - # update cache - cache.cost = cost - cache.grad_vec .= to_vec(grad)[1] + cache.cost = cost # update cost function value + cache.grad_vec .= to_vec(grad)[1] # update vectorized gradient return cache.cost, cache.grad_vec end -function cost_function(cache::PEPSCostFunctionCache{T}) - return function cost_function(::Euclidean, peps_vec::Vector{T}) - if !(peps_vec == cache.peps_vec) # update cache - cost_and_grad!(cache, peps_vec) - end - return cache.cost +# Functor returns only the cost such that we can access the cache +# using get_objective(::AbstractManoptProblem) in RecordActions +function (cache::PEPSCostFunctionCache{T})(::Euclidean, peps_vec::Vector{T}) where {T} + if !(peps_vec == cache.peps_vec) # update cache if at new point + cost_and_grad!(cache, peps_vec) end + return cache.cost end - -function grad_function(cache::PEPSCostFunctionCache) - return function grad_function(::Euclidean, peps_vec::Vector{T}) - if !(peps_vec == cache.peps_vec) # update cache +function gradient_function(cache::PEPSCostFunctionCache{T}) where {T} + return function gradient_function(::Euclidean, peps_vec::Vector{T}) + if !(peps_vec == cache.peps_vec) # update cache if at new point cost_and_grad!(cache, peps_vec) end return cache.grad_vec @@ -280,13 +277,13 @@ function fixedpoint( # construct cost and grad functions peps₀_vec, from_vec = to_vec(peps₀) - cache = PEPSCostFunctionCache(operator, alg, peps_vec, from_vec, env₀) - cost = cost_function(cache) - grad = grad_function(cache) + cache = PEPSCostFunctionCache(operator, alg, peps₀_vec, from_vec, env₀) + cost = cache + grad = gradient_function(cache) # optimize - fld = scalartype(peps₀) <: Real ? Manifolds.ℝ : Manifolds.ℂ - M = Euclidean(; field=fld) + # fld = scalartype(peps₀) <: Real ? Manifolds.ℝ : Manifolds.ℂ # Manopt can't optimize over ℂ? + M = Euclidean(length(peps₀_vec)) retraction_method = if isnothing(alg.symmetrization) default_retraction_method(M) else diff --git a/src/states/infinitepeps.jl b/src/states/infinitepeps.jl index 846ad707c..f08acaf8e 100644 --- a/src/states/infinitepeps.jl +++ b/src/states/infinitepeps.jl @@ -122,6 +122,7 @@ Base.copy(T::InfinitePEPS) = InfinitePEPS(copy(T.A)) Base.similar(T::InfinitePEPS, args...) = InfinitePEPS(similar(T.A, args...)) Base.repeat(T::InfinitePEPS, counts...) = InfinitePEPS(repeat(T.A, counts...)) +## Indexing Base.getindex(T::InfinitePEPS, args...) = Base.getindex(T.A, args...) Base.setindex!(T::InfinitePEPS, args...) = (Base.setindex!(T.A, args...); T) Base.axes(T::InfinitePEPS, args...) = axes(T.A, args...) @@ -131,7 +132,10 @@ end function eachcoordinate(x::InfinitePEPS, dirs) return collect(Iterators.product(dirs, axes(x, 1), axes(x, 2))) end + +## Tensor properties TensorKit.space(t::InfinitePEPS, i, j) = space(t[i, j], 1) +TensorKit.dim(t::InfinitePEPS) = sum(dim.(t.A)) ## Math Base.:+(ψ₁::InfinitePEPS, ψ₂::InfinitePEPS) = InfinitePEPS(ψ₁.A + ψ₂.A) diff --git a/test/heisenberg.jl b/test/heisenberg.jl index a78e98af0..2171d0687 100644 --- a/test/heisenberg.jl +++ b/test/heisenberg.jl @@ -15,12 +15,17 @@ opt_alg = PEPSOptimize(; boundary_alg=ctm_alg, tol=1e-3) # https://github.com/jurajHasik/j1j2_ipeps_states/blob/main/single-site_pg-C4v-A1/j20.0/state_1s_A1_j20.0_D2_chi_opt48.dat E_ref = -0.6602310934799577 +Random.seed!(123) +H = heisenberg_XYZ(InfiniteSquare()) +psi_init = InfinitePEPS(2, Dbond) +env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg) + @testset "(1, 1) unit cell AD optimization" begin # initialize states Random.seed!(123) H = heisenberg_XYZ(InfiniteSquare()) psi_init = InfinitePEPS(2, Dbond) - env_init = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg) + env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg) # optimize energy and compute correlation lengths result = fixedpoint(psi_init, H, opt_alg, env_init) From d36b2ff923c03a801bb889e08b8c0cb3af7e4bd8 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Wed, 18 Dec 2024 19:37:24 +0100 Subject: [PATCH 08/28] Make fixedpoint optimize (kind of) --- src/PEPSKit.jl | 4 ++-- src/algorithms/peps_opt.jl | 28 ++++++++++++++++++---------- test/heisenberg.jl | 5 ----- 3 files changed, 20 insertions(+), 17 deletions(-) diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index f1a7d5064..0b750c800 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -159,15 +159,15 @@ module Defaults RecordUnitCellGradientNorm() => :unitcell_gradient_norm, RecordTime() => :time, ] + const debug_group = [:Iteration, :Cost, :GradientNorm, :Time, :Stop] const stopping_criterion = StopAfterIteration(100) | StopWhenGradientNormLess(1e-4) const optim_maxiter = 100 const optim_tol = 1e-4 const optim_kwargs = (; - memory_size=32, stopping_criterion=StopAfterIteration(optim_maxiter) | StopWhenGradientNormLess(optim_tol), record=record_group, - return_state=true, + debug=debug_group, ) const fpgrad_maxiter = 30 const fpgrad_tol = 1e-6 diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl index 1c67d287b..7bd1576a0 100644 --- a/src/algorithms/peps_opt.jl +++ b/src/algorithms/peps_opt.jl @@ -1,4 +1,11 @@ -using Manifolds, Manopt +using Manifolds: Manifolds +using Manifolds: + AbstractManifold, + AbstractRetractionMethod, + Euclidean, + default_retraction_method, + retract +using Manopt abstract type GradMode{F} end @@ -158,9 +165,10 @@ function PEPSOptimize(; gradient_alg=Defaults.gradient_alg, reuse_env=Defaults.reuse_env, symmetrization=nothing, + kwargs..., ) stopping_criterion = StopAfterIteration(maxiter) | StopWhenGradientNormLess(tol) - optim_kwargs = (; stopping_criterion, record=Defaults.record_group) + optim_kwargs = merge(Defaults.optim_kwargs, (; stopping_criterion, kwargs...)) return PEPSOptimize( boundary_alg, optim_alg, optim_kwargs, gradient_alg, reuse_env, symmetrization ) @@ -184,7 +192,7 @@ function PEPSCostFunctionCache( alg, env, from_vec, - peps_vec, + similar(peps_vec), similar(peps_vec), 0.0, (; truncation_error=0.0, condition_number=1.0), @@ -206,10 +214,10 @@ function cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) wh cache.alg.boundary_alg; alg_rrule=cache.alg.gradient_alg, ) - cache.env_info = info # update environment information (truncation error, ...) cost = expectation_value(peps, cache.operator, env) ignore_derivatives() do update!(cache.env, env) # update environment in-place + cache.env_info = info # update environment information (truncation error, ...) isapprox(imag(cost), 0; atol=sqrt(eps(real(cost)))) || @warn "Expectation value is not real: $cost." end @@ -315,7 +323,7 @@ function _rrule( ) envs = leading_boundary(envinit, state, alg) - function leading_boundary_diffgauge_pullback(Δenvs′) + function leading_boundary_diffgauge_pullback((Δenvs′, Δinfo)) Δenvs = unthunk(Δenvs′) # find partial gradients of gauge_fixed single CTMRG iteration @@ -343,9 +351,9 @@ function _rrule( alg::SimultaneousCTMRG, ) @assert !isnothing(alg.projector_alg.svd_alg.rrule_alg) - envs = leading_boundary(envinit, state, alg) - envsconv, info = ctmrg_iteration(state, envs, alg) - envs_fixed, signs = gauge_fix(envs, envsconv) + envs, = leading_boundary(envinit, state, alg) + envs_conv, info = ctmrg_iteration(state, envs, alg) + envs_fixed, signs = gauge_fix(envs, envs_conv) # Fix SVD Ufix, Vfix = fix_relative_phases(info.U, info.V, signs) @@ -355,7 +363,7 @@ function _rrule( alg_fixed = @set alg.projector_alg.svd_alg = svd_alg_fixed alg_fixed = @set alg_fixed.projector_alg.trscheme = notrunc() - function leading_boundary_fixed_pullback(Δenvs′) + function leading_boundary_fixed_pullback((Δenvs′, Δinfo)) Δenvs = unthunk(Δenvs′) f(A, x) = fix_global_phases(x, ctmrg_iteration(A, x, alg_fixed)[1]) @@ -369,7 +377,7 @@ function _rrule( return NoTangent(), ZeroTangent(), ∂F∂envs, NoTangent() end - return envs_fixed, leading_boundary_fixed_pullback + return (envs_fixed, info), leading_boundary_fixed_pullback end @doc """ diff --git a/test/heisenberg.jl b/test/heisenberg.jl index 2171d0687..03a748c51 100644 --- a/test/heisenberg.jl +++ b/test/heisenberg.jl @@ -15,11 +15,6 @@ opt_alg = PEPSOptimize(; boundary_alg=ctm_alg, tol=1e-3) # https://github.com/jurajHasik/j1j2_ipeps_states/blob/main/single-site_pg-C4v-A1/j20.0/state_1s_A1_j20.0_D2_chi_opt48.dat E_ref = -0.6602310934799577 -Random.seed!(123) -H = heisenberg_XYZ(InfiniteSquare()) -psi_init = InfinitePEPS(2, Dbond) -env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg) - @testset "(1, 1) unit cell AD optimization" begin # initialize states Random.seed!(123) From bb6cdf12c57bde4ac6ad841baf35f55c333c3bbd Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Wed, 18 Dec 2024 20:17:11 +0100 Subject: [PATCH 09/28] Improve debug printing --- src/PEPSKit.jl | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 0b750c800..7d0ad4cc1 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -159,7 +159,16 @@ module Defaults RecordUnitCellGradientNorm() => :unitcell_gradient_norm, RecordTime() => :time, ] - const debug_group = [:Iteration, :Cost, :GradientNorm, :Time, :Stop] + const debug_group = [ + (:Iteration, "Optim %-5d"), + (:Cost, "f(x) = %.8f"), + (:GradientNorm, " ‖∂f‖ = %.8f "), + (:Stepsize, " step size = %.8f "), + DebugTime(; prefix="time =", mode=:iterative), + DebugWarnIfCostIncreases(:Always; tol=1e-12), + :Stop, + "\n", + ] const stopping_criterion = StopAfterIteration(100) | StopWhenGradientNormLess(1e-4) const optim_maxiter = 100 const optim_tol = 1e-4 From 443e58b00a30daf02f325a7351d7f85493d30f52 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Thu, 19 Dec 2024 19:33:31 +0100 Subject: [PATCH 10/28] Add LineSearches, add docstrings --- Project.toml | 2 + src/PEPSKit.jl | 77 +++++++++++------ src/algorithms/ctmrg/sequential.jl | 3 +- src/algorithms/ctmrg/simultaneous.jl | 3 +- src/algorithms/peps_opt.jl | 124 ++++++++++++++++++++++++--- 5 files changed, 169 insertions(+), 40 deletions(-) diff --git a/Project.toml b/Project.toml index 0deecf910..1284b5584 100644 --- a/Project.toml +++ b/Project.toml @@ -9,6 +9,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77" +LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36" MPSKit = "bb1c41ca-d63c-52ed-829e-0820dda26502" @@ -30,6 +31,7 @@ ChainRulesCore = "1.0" Compat = "3.46, 4.2" FiniteDifferences = "0.12" KrylovKit = "0.8" +LineSearches = "7.3.0" LinearAlgebra = "1" LoggingExtras = "1" MPSKit = "0.11" diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 7d0ad4cc1..2d50093dc 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -77,19 +77,33 @@ include("algorithms/peps_opt.jl") # Optimization const optim_alg = quasi_Newton + const optim_maxiter = 100 + const optim_tol = 1e-4 + const stopping_criterion = StopAfterIteration(optim_maxiter) | StopWhenGradientNormLess(optim_tol) const record_group = [ - RecordCost(), - RecordGradientNorm(), - RecordConditionNumber(), - RecordCostUnitCell(), - RecordTime(), + RecordCost() => :Cost, + RecordGradientNorm() => :GradientNorm, + RecordTruncationError() => :TruncationError, + RecordConditionNumber() => :ConditionNumber, + RecordUnitCellGradientNorm() => :UnitcellGradientNorm, + RecordTime(; mode=:iterative) => :Time, + ] + const debug_group = [ + (:Iteration, "Optim %-4d "), + (:Cost, "f(x) = %.8f "), + (:GradientNorm, "‖∂f‖ = %.8f "), + (:Stepsize, "step size = %.4f "), + DebugTime(; prefix="time =", mode=:iterative), + "\n", + DebugWarnIfCostIncreases(:Always; tol=1e-10), + "\n", + :Stop, ] - const optim_kwargs = (; - memory_size=32, - stopping_criterion=StopAfterIteration(100) | StopWhenGradientNormLess(1e-4), + stopping_criterion=StopAfterIteration(optim_maxiter) | + StopWhenGradientNormLess(optim_tol), record=record_group, - return_state=true, + debug=debug_group, ) const fpgrad_maxiter = 30 const fpgrad_tol = 1e-6 @@ -104,6 +118,7 @@ include("algorithms/peps_opt.jl") Module containing default values that represent typical algorithm parameters. +# CTMRG - `ctmrg_maxiter`: Maximal number of CTMRG iterations per run - `ctmrg_miniter`: Minimal number of CTMRG carried out - `ctmrg_tol`: Tolerance checking singular value and norm convergence @@ -114,13 +129,23 @@ Module containing default values that represent typical algorithm parameters. - `projector_alg_type`: Default type of projector algorithm - `projector_alg`: Algorithm to compute CTMRG projectors - `ctmrg_alg`: Algorithm for performing CTMRG runs + +# Optimization +- `optim_alg`: Manopt optimizer function +- `optim_maxiter`: Maximal number of optimization iterations +- `optim_tol`: Gradient norm convergence tolerance +- `stopping_criterion`: Manopt stopping criterion +- `record_group`: Values that recorded during Manopt optimization +- `debug_group`: Values that are printed during Manopt optimization +- `optim_kwargs`: All keyword arguments that are passed onto a Manopt optimization call - `fpgrad_maxiter`: Maximal number of iterations for computing the CTMRG fixed-point gradient - `fpgrad_tol`: Convergence tolerance for the fixed-point gradient iteration - `reuse_env`: If `true`, the current optimization step is initialized on the previous environment - `gradient_linsolver`: Default linear solver for the `LinSolver` gradient algorithm - `iterscheme`: Scheme for differentiating one CTMRG iteration - `gradient_alg`: Algorithm to compute the gradient fixed-point -# TODO + +# OhMyThreads scheduler - `scheduler`: Multi-threading scheduler which can be accessed via `set_scheduler!` """ module Defaults @@ -152,26 +177,29 @@ module Defaults # Optimization const optim_alg = quasi_Newton + const optim_maxiter = 100 + const optim_tol = 1e-4 + const stopping_criterion = + StopAfterIteration(optim_maxiter) | StopWhenGradientNormLess(optim_tol) const record_group = [ - RecordCost() => :cost, - RecordGradientNorm() => :gradient_norm, - RecordConditionNumber() => :condition, # TODO: implement PEPS record actions - RecordUnitCellGradientNorm() => :unitcell_gradient_norm, - RecordTime() => :time, + RecordCost() => :Cost, + RecordGradientNorm() => :GradientNorm, + RecordTruncationError() => :TruncationError, + RecordConditionNumber() => :ConditionNumber, + RecordUnitCellGradientNorm() => :UnitcellGradientNorm, + RecordTime(; mode=:iterative) => :Time, ] const debug_group = [ - (:Iteration, "Optim %-5d"), - (:Cost, "f(x) = %.8f"), - (:GradientNorm, " ‖∂f‖ = %.8f "), - (:Stepsize, " step size = %.8f "), + (:Iteration, "Optim %-4d "), + (:Cost, "f(x) = %.8f "), + (:GradientNorm, "‖∂f‖ = %.8f "), + (:Stepsize, "step size = %.4f "), DebugTime(; prefix="time =", mode=:iterative), - DebugWarnIfCostIncreases(:Always; tol=1e-12), - :Stop, "\n", + DebugWarnIfCostIncreases(:Always; tol=1e-10), + "\n", + :Stop, ] - const stopping_criterion = StopAfterIteration(100) | StopWhenGradientNormLess(1e-4) - const optim_maxiter = 100 - const optim_tol = 1e-4 const optim_kwargs = (; stopping_criterion=StopAfterIteration(optim_maxiter) | StopWhenGradientNormLess(optim_tol), @@ -239,6 +267,7 @@ export LocalOperator export expectation_value, costfun, product_peps, correlation_length export leading_boundary export PEPSOptimize, GeomSum, ManualIter, LinSolver +export RecordTruncationError, RecordConditionNumber, RecordUnitCellGradientNorm export fixedpoint export absorb_weight diff --git a/src/algorithms/ctmrg/sequential.jl b/src/algorithms/ctmrg/sequential.jl index 6ae2a8ee9..964e0c098 100644 --- a/src/algorithms/ctmrg/sequential.jl +++ b/src/algorithms/ctmrg/sequential.jl @@ -73,8 +73,9 @@ function sequential_projectors( return proj end + truncation_error = maximum(copy(ϵ)) condition_number = maximum(_condition_number, copy(S)) - info = (; truncation_error=maximum(copy(ϵ)), condition_number) + info = (; truncation_error, condition_number) return (map(first, projectors), map(last, projectors)), info end function sequential_projectors( diff --git a/src/algorithms/ctmrg/simultaneous.jl b/src/algorithms/ctmrg/simultaneous.jl index 870922be4..e23ce83ac 100644 --- a/src/algorithms/ctmrg/simultaneous.jl +++ b/src/algorithms/ctmrg/simultaneous.jl @@ -92,8 +92,9 @@ function simultaneous_projectors( P_left = map(first, projectors) P_right = map(last, projectors) S = copy(S) + truncation_error=maximum(copy(ϵ)) condition_number = maximum(_condition_number, S) - info = (; truncation_error=maximum(copy(ϵ)), condition_number, U=copy(U), S, V=copy(V)) + info = (; truncation_error, condition_number, U=copy(U), S, V=copy(V)) return (P_left, P_right), info end function simultaneous_projectors( diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl index 7bd1576a0..b08a45db9 100644 --- a/src/algorithms/peps_opt.jl +++ b/src/algorithms/peps_opt.jl @@ -85,6 +85,12 @@ function LinSolver(; return LinSolver{iterscheme}(solver) end +""" + mutable struct RecordTruncationError <: RecordAction + +Record the maximal truncation error of all `boundary_alg` runs of the corresponding +optimization iteration. +""" mutable struct RecordTruncationError <: RecordAction recorded_values::Vector{Float64} RecordTruncationError() = new(Vector{Float64}()) @@ -96,6 +102,12 @@ function (r::RecordTruncationError)( return Manopt.record_or_reset!(r, cache.env_info.truncation_error, i) end +""" + mutable struct RecordConditionNumber <: RecordAction + +Record the maximal condition number of all `boundary_alg` runs of the corresponding +optimization iteration. +""" mutable struct RecordConditionNumber <: RecordAction recorded_values::Vector{Float64} RecordConditionNumber() = new(Vector{Float64}()) @@ -107,6 +119,12 @@ function (r::RecordConditionNumber)( return Manopt.record_or_reset!(r, cache.env_info.condition_number, i) end +""" + mutable struct RecordUnitCellGradientNorm <: RecordAction + +Record the PEPS gradient norms unit cell entry-wise, i.e. an array +of norms `norm.(peps.A)`. +""" mutable struct RecordUnitCellGradientNorm <: RecordAction recorded_values::Vector{Matrix{Float64}} RecordUnitCellGradientNorm() = new(Vector{Matrix{Float64}}()) @@ -120,15 +138,19 @@ function (r::RecordUnitCellGradientNorm)( end """ -TODO - -Algorithm struct that represent PEPS ground-state optimization using AD. -Set the algorithm to contract the infinite PEPS in `boundary_alg`; -currently only `CTMRGAlgorithm`s are supported. The `optimizer` computes the gradient directions -based on the CTMRG gradient and updates the PEPS parameters. In this optimization, -the CTMRG runs can be started on the converged environments of the previous optimizer -step by setting `reuse_env` to true. Otherwise a random environment is used at each -step. The CTMRG gradient itself is computed using the `gradient_alg` algorithm. + struct PEPSOptimize{G} + +Algorithm struct for PEPS optimization using automatic differentiation. + +# Fields +- `boundary_alg::CTMRGAlgorithm`: algorithm for determining the PEPS environment +- `optim_alg::Function`: Manopt optimization algorithm +- `optim_kwargs::NamedTuple`: Keyword arguments provided to the Manopt `optim_alg` call +- `gradient_alg::G`: Algorithm computing the cost function gradient in reverse-mode +- `reuse_env::Bool`: If `true` the previous environment is used to initialize the next + `leading_boundary` call +- `symmetrization::Union{Nothing,SymmetrizationStyle}`: Symmetrize the PEPS and PEPS + gradient after each optimization iteration (does nothing if `nothing` is provided) """ struct PEPSOptimize{G} boundary_alg::CTMRGAlgorithm @@ -157,6 +179,25 @@ struct PEPSOptimize{G} ) end end + +""" + PEPSOptimize(; + boundary_alg=Defaults.ctmrg_alg, + optim_alg=Defaults.optim_alg, + maxiter=Defaults.optim_maxiter, + tol=Defaults.optim_tol, + gradient_alg=Defaults.gradient_alg, + reuse_env=Defaults.reuse_env, + symmetrization=nothing, + kwargs..., + ) + +Convenience keyword argument constructor for `PEPSOptimize` algorithms. +Here, `maxiter` and `tol` are passed onto `StopAfterIteration` and `StopWhenGradientNormLess` +stopping criteria, respectively. Additionally, any keyword arguments can be provided which +are then stored inside `optim_kwargs` and passed to the Manopt optimization call, such that +that all arguments of the respective `optim_alg` can be used. +""" function PEPSOptimize(; boundary_alg=Defaults.ctmrg_alg, optim_alg=Defaults.optim_alg, @@ -174,6 +215,22 @@ function PEPSOptimize(; ) end +""" + mutable struct PEPSCostFunctionCache{T} + +Stores objects used for computing PEPS cost functions during optimization that are +needed apart from the PEPS that is being optimized. + +# Fields +- `operator::LocalOperator`: cost function operator +- `alg::PEPSOptimize`: optimization parameters +- `env::CTMRGEnv`: environment of the current PEPS +- `from_vec`: map which returns vectorized PEPS as an `InfinitePEPS` +- `peps_vec::Vector{T}`: current vectorized PEPS +- `grad_vec::Vector{T}`: current vectorized gradient +- `cost::Float64`: current cost function value +- `env_info::NamedTuple`: return info of `leading_boundary` used by `RecordAction`s +""" mutable struct PEPSCostFunctionCache{T} operator::LocalOperator alg::PEPSOptimize @@ -184,6 +241,15 @@ mutable struct PEPSCostFunctionCache{T} cost::Float64 env_info::NamedTuple end + +""" + PEPSCostFunctionCache( + operator::LocalOperator, alg::PEPSOptimize, peps_vec::Vector, from_vec, env::CTMRGEnv + ) + +Initialize a `PEPSCostFunctionCache` using `peps_vec` from which the vector dimension +and scalartype are derived. +""" function PEPSCostFunctionCache( operator::LocalOperator, alg::PEPSOptimize, peps_vec::Vector, from_vec, env::CTMRGEnv ) @@ -199,6 +265,12 @@ function PEPSCostFunctionCache( ) end +""" + cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) where {T} + +Update the cost and gradient of the `PEPSCostFunctionCache` with respect to the new point +`peps_vec`. +""" function cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) where {T} cache.peps_vec .= peps_vec # update point in manifold peps = cache.from_vec(peps_vec) # convert back to InfinitePEPS @@ -235,14 +307,27 @@ function cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) wh return cache.cost, cache.grad_vec end -# Functor returns only the cost such that we can access the cache -# using get_objective(::AbstractManoptProblem) in RecordActions +""" + (cache::PEPSCostFunctionCache{T})(::Euclidean, peps_vec::Vector{T}) where {T} + +Return the cost of `cache` and recompute if `peps_vec` is a new point. +""" function (cache::PEPSCostFunctionCache{T})(::Euclidean, peps_vec::Vector{T}) where {T} + # Note that it is necessary to implement the cost function as a functor so that the + # `PEPSCostFunctionCache` is available through `get_objective(::AbstractManoptProblem)` + # to the `RecordAction`s during optimization if !(peps_vec == cache.peps_vec) # update cache if at new point cost_and_grad!(cache, peps_vec) end return cache.cost end + +""" + gradient_function(cache::PEPSCostFunctionCache{T}) where {T} + +Get the gradient function of `cache` which returns the gradient vector +and recomputes it if the provided point is new. +""" function gradient_function(cache::PEPSCostFunctionCache{T}) where {T} return function gradient_function(::Euclidean, peps_vec::Vector{T}) if !(peps_vec == cache.peps_vec) # update cache if at new point @@ -252,8 +337,11 @@ function gradient_function(cache::PEPSCostFunctionCache{T}) where {T} end end -# First retract and then symmetrize the resulting PEPS -# (ExponentialRetraction is the default retraction for Euclidean manifolds) +""" + SymmetrizeExponentialRetraction <: AbstractRetractionMethod + +Exponential retraction followed by a symmetrization step. +""" struct SymmetrizeExponentialRetraction <: AbstractRetractionMethod symmetrization::SymmetrizationStyle from_vec::Function @@ -268,7 +356,15 @@ function Manifolds.retract( end """ -TODO + fixedpoint( + peps₀::InfinitePEPS{T}, + operator::LocalOperator, + alg::PEPSOptimize, + env₀::CTMRGEnv=CTMRGEnv(peps₀, field(T)^20); + ) where {T} + +Optimize a PEPS starting from `peps₀` and `env₀` by minimizing the cost function given +by expectation value of `operator`. All optimization parameters are provided through `alg`. """ function fixedpoint( peps₀::InfinitePEPS{T}, From d0d42f657e1387fd3f6ab13b77055b2b2bf8c559 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Thu, 19 Dec 2024 21:46:58 +0100 Subject: [PATCH 11/28] Update Defaults and Defaults docstring --- src/PEPSKit.jl | 151 +++++++++++++++++++++++-------------------------- 1 file changed, 70 insertions(+), 81 deletions(-) diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 2d50093dc..a32731dc5 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -60,93 +60,87 @@ include("algorithms/peps_opt.jl") """ module Defaults - # CTMRG - const ctmrg_tol = 1e-8 - const ctmrg_maxiter = 100 - const ctmrg_miniter = 4 - const sparse = false - const trscheme = FixedSpaceTruncation() - const fwd_alg = TensorKit.SDD() - const rrule_alg = Arnoldi(; tol=ctmrg_tol, krylovdim=48, verbosity=-1) - const svd_alg = SVDAdjoint(; fwd_alg, rrule_alg) - const projector_alg_type = HalfInfiniteProjector - const projector_alg = projector_alg_type(svd_alg, trscheme, 2) - const ctmrg_alg = SimultaneousCTMRG( - ctmrg_tol, ctmrg_maxiter, ctmrg_miniter, 2, projector_alg - ) - - # Optimization - const optim_alg = quasi_Newton - const optim_maxiter = 100 - const optim_tol = 1e-4 - const stopping_criterion = StopAfterIteration(optim_maxiter) | StopWhenGradientNormLess(optim_tol) - const record_group = [ - RecordCost() => :Cost, - RecordGradientNorm() => :GradientNorm, - RecordTruncationError() => :TruncationError, - RecordConditionNumber() => :ConditionNumber, - RecordUnitCellGradientNorm() => :UnitcellGradientNorm, - RecordTime(; mode=:iterative) => :Time, - ] - const debug_group = [ - (:Iteration, "Optim %-4d "), - (:Cost, "f(x) = %.8f "), - (:GradientNorm, "‖∂f‖ = %.8f "), - (:Stepsize, "step size = %.4f "), - DebugTime(; prefix="time =", mode=:iterative), - "\n", - DebugWarnIfCostIncreases(:Always; tol=1e-10), - "\n", - :Stop, - ] - const optim_kwargs = (; - stopping_criterion=StopAfterIteration(optim_maxiter) | - StopWhenGradientNormLess(optim_tol), - record=record_group, - debug=debug_group, - ) - const fpgrad_maxiter = 30 - const fpgrad_tol = 1e-6 - const gradient_linsolver = KrylovKit.BiCGStab(; maxiter=fpgrad_maxiter, tol=fpgrad_tol) - const iterscheme = :fixed - const reuse_env = true - const gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme) - - # OhMyThreads scheduler defaults - const scheduler = Ref{Scheduler}(Threads.nthreads() == 1 ? SerialScheduler() : DynamicScheduler()) - end -Module containing default values that represent typical algorithm parameters. +Module containing default algorithm parameter values and arguments. # CTMRG -- `ctmrg_maxiter`: Maximal number of CTMRG iterations per run -- `ctmrg_miniter`: Minimal number of CTMRG carried out -- `ctmrg_tol`: Tolerance checking singular value and norm convergence -- `trscheme`: Truncation scheme for SVDs and other decompositions -- `fwd_alg`: SVD algorithm that is used in the forward pass +- `ctmrg_tol=1e-8`: Tolerance checking singular value and norm convergence +- `ctmrg_maxiter=100`: Maximal number of CTMRG iterations per run +- `ctmrg_miniter=4`: Minimal number of CTMRG carried out +- `trscheme=FixedSpaceTruncation()`: Truncation scheme for SVDs and other decompositions +- `fwd_alg=TensorKit.SDD()`: SVD algorithm that is used in the forward pass - `rrule_alg`: Reverse-rule for differentiating that SVD -- `svd_alg`: Combination of `fwd_alg` and `rrule_alg` -- `projector_alg_type`: Default type of projector algorithm + + ``` + rrule_alg = Arnoldi(; tol=ctmrg_tol, krylovdim=48, verbosity=-1) + ``` + +- `svd_alg=SVDAdjoint(; fwd_alg, rrule_alg)`: Combination of `fwd_alg` and `rrule_alg` +- `projector_alg_type=HalfInfiniteProjector`: Default type of projector algorithm - `projector_alg`: Algorithm to compute CTMRG projectors + + ``` + projector_alg = projector_alg_type(; svd_alg, trscheme, verbosity=0) + ``` + - `ctmrg_alg`: Algorithm for performing CTMRG runs + ``` + ctmrg_alg = SimultaneousCTMRG( + ctmrg_tol, ctmrg_maxiter, ctmrg_miniter, 2, projector_alg + ) + ``` + # Optimization -- `optim_alg`: Manopt optimizer function -- `optim_maxiter`: Maximal number of optimization iterations -- `optim_tol`: Gradient norm convergence tolerance +- `optim_alg=quasi_Newton`: Manopt optimizer function +- `optim_maxiter=100`: Maximal number of optimization iterations +- `optim_tol=1e-4`: Gradient norm convergence tolerance - `stopping_criterion`: Manopt stopping criterion -- `record_group`: Values that recorded during Manopt optimization -- `debug_group`: Values that are printed during Manopt optimization + + ``` + stopping_criterion = StopAfterIteration(optim_maxiter) | StopWhenGradientNormLess(optim_tol) + ``` + +- `record_group`: Values (`RecordAction`s) that are recorded during Manopt optimization + + ``` + record_group = [ + RecordCost() => :Cost, + RecordGradientNorm() => :GradientNorm, + RecordTruncationError() => :TruncationError, + RecordConditionNumber() => :ConditionNumber, + RecordUnitCellGradientNorm() => :UnitcellGradientNorm, + RecordTime(; mode=:iterative) => :Time, + ] + ``` + +- `debug_group=[...]`: Values (`DebugAction`s) that are printed during Manopt optimization + - `optim_kwargs`: All keyword arguments that are passed onto a Manopt optimization call -- `fpgrad_maxiter`: Maximal number of iterations for computing the CTMRG fixed-point gradient -- `fpgrad_tol`: Convergence tolerance for the fixed-point gradient iteration -- `reuse_env`: If `true`, the current optimization step is initialized on the previous environment + + ``` + optim_kwargs = (; stopping_criterion, record=record_group, debug=debug_group) + ``` + +- `fpgrad_maxiter=30`: Maximal number of iterations for computing the CTMRG fixed-point gradient +- `fpgrad_tol=1e-6`: Convergence tolerance for the fixed-point gradient iteration +- `iterscheme=:fixed`: Scheme for differentiating one CTMRG iteration - `gradient_linsolver`: Default linear solver for the `LinSolver` gradient algorithm -- `iterscheme`: Scheme for differentiating one CTMRG iteration + + ``` + gradient_linsolver=KrylovKit.BiCGStab(; maxiter=fpgrad_maxiter, tol=fpgrad_tol) + ``` + - `gradient_alg`: Algorithm to compute the gradient fixed-point + ``` + gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme) + ``` + +- `reuse_env=true`: If `true`, the current optimization step is initialized on the previous environment + # OhMyThreads scheduler -- `scheduler`: Multi-threading scheduler which can be accessed via `set_scheduler!` +- `scheduler=Ref{Scheduler}(...)`: Multi-threading scheduler which can be accessed via `set_scheduler!` """ module Defaults using TensorKit, KrylovKit, OhMyThreads @@ -170,7 +164,7 @@ module Defaults const rrule_alg = Arnoldi(; tol=ctmrg_tol, krylovdim=48, verbosity=-1) const svd_alg = SVDAdjoint(; fwd_alg, rrule_alg) const projector_alg_type = HalfInfiniteProjector - const projector_alg = projector_alg_type(svd_alg, trscheme, 2) + const projector_alg = projector_alg_type(; svd_alg, trscheme, verbosity=0) const ctmrg_alg = SimultaneousCTMRG( ctmrg_tol, ctmrg_maxiter, ctmrg_miniter, 2, projector_alg ) @@ -200,18 +194,13 @@ module Defaults "\n", :Stop, ] - const optim_kwargs = (; - stopping_criterion=StopAfterIteration(optim_maxiter) | - StopWhenGradientNormLess(optim_tol), - record=record_group, - debug=debug_group, - ) + const optim_kwargs = (; stopping_criterion, record=record_group, debug=debug_group) const fpgrad_maxiter = 30 const fpgrad_tol = 1e-6 const gradient_linsolver = KrylovKit.BiCGStab(; maxiter=fpgrad_maxiter, tol=fpgrad_tol) const iterscheme = :fixed - const reuse_env = true const gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme) + const reuse_env = true # OhMyThreads scheduler defaults const scheduler = Ref{Scheduler}() From e38dbce8d990257788e68901e336ce2c64fcd982 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Thu, 19 Dec 2024 22:08:47 +0100 Subject: [PATCH 12/28] Adjust scripts to new `leading_boundary` return values --- README.md | 2 +- docs/src/index.md | 2 +- examples/boundary_mps.jl | 4 ++-- examples/heisenberg.jl | 2 +- examples/hubbard_su.jl | 2 +- src/algorithms/peps_opt.jl | 4 ++-- test/boundarymps/vumps.jl | 4 ++-- test/ctmrg/fixed_iterscheme.jl | 4 ++-- test/ctmrg/flavors.jl | 6 +++--- test/ctmrg/gaugefix.jl | 4 ++-- test/ctmrg/gradients.jl | 2 +- test/heisenberg.jl | 4 ++-- test/j1j2_model.jl | 2 +- test/pwave.jl | 2 +- test/tf_ising.jl | 2 +- 15 files changed, 23 insertions(+), 23 deletions(-) diff --git a/README.md b/README.md index 25f9a9248..8dd2d139f 100644 --- a/README.md +++ b/README.md @@ -54,7 +54,7 @@ opt_alg = PEPSOptimize(; # ground state search state = InfinitePEPS(2, D) -ctm = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg) +ctm, = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg) result = fixedpoint(state, H, opt_alg, ctm) @show result.E # -0.6625... diff --git a/docs/src/index.md b/docs/src/index.md index 709f42222..617753100 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -36,7 +36,7 @@ opt_alg = PEPSOptimize(; # ground state search state = InfinitePEPS(2, D) -ctm = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg) +ctm, = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg) result = fixedpoint(state, H, opt_alg, ctm) @show result.E # -0.6625... diff --git a/examples/boundary_mps.jl b/examples/boundary_mps.jl index 51af6deae..6e82091d1 100644 --- a/examples/boundary_mps.jl +++ b/examples/boundary_mps.jl @@ -32,7 +32,7 @@ mps, envs, ϵ = leading_boundary(mps, T, VUMPS()) N = abs(prod(expectation_value(mps, T))) # This can be compared to the result obtained using the CTMRG algorithm -ctm = leading_boundary( +ctm, = leading_boundary( peps, SimultaneousCTMRG(; verbosity=1), CTMRGEnv(peps, ComplexSpace(20)) ) N´ = abs(norm(peps, ctm)) @@ -55,7 +55,7 @@ mps2 = PEPSKit.initializeMPS(T2, fill(ComplexSpace(20), 2, 2)) mps2, envs2, ϵ = leading_boundary(mps2, T2, VUMPS()) N2 = abs(prod(expectation_value(mps2, T2))) -ctm2 = leading_boundary( +ctm2, = leading_boundary( peps2, SimultaneousCTMRG(; verbosity=1), CTMRGEnv(peps2, ComplexSpace(20)) ) N2´ = abs(norm(peps2, ctm2)) diff --git a/examples/heisenberg.jl b/examples/heisenberg.jl index 98739780b..651f46aa1 100644 --- a/examples/heisenberg.jl +++ b/examples/heisenberg.jl @@ -26,6 +26,6 @@ opt_alg = PEPSOptimize(; # E/N = −0.6694421, which is a QMC estimate from https://arxiv.org/abs/1101.3281. # Of course there is a noticable bias for small χbond and χenv. ψ₀ = InfinitePEPS(2, χbond) -env₀ = leading_boundary(CTMRGEnv(ψ₀, ℂ^χenv), ψ₀, ctm_alg) +env₀, = leading_boundary(CTMRGEnv(ψ₀, ℂ^χenv), ψ₀, ctm_alg) result = fixedpoint(ψ₀, H, opt_alg, env₀) @show result.E diff --git a/examples/hubbard_su.jl b/examples/hubbard_su.jl index 6478a28e0..d65512f59 100644 --- a/examples/hubbard_su.jl +++ b/examples/hubbard_su.jl @@ -43,7 +43,7 @@ Espace = Vect[fℤ₂](0 => χenv0 / 2, 1 => χenv0 / 2) envs = CTMRGEnv(randn, Float64, peps, Espace) for χ in [χenv0, χenv] ctm_alg = SequentialCTMRG(; maxiter=300, tol=1e-7) - envs = leading_boundary(envs, peps, ctm_alg) + envs, = leading_boundary(envs, peps, ctm_alg) end # Benchmark values of the ground state energy from diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl index b08a45db9..0a91ee583 100644 --- a/src/algorithms/peps_opt.jl +++ b/src/algorithms/peps_opt.jl @@ -417,7 +417,7 @@ function _rrule( state, alg::CTMRGAlgorithm, ) - envs = leading_boundary(envinit, state, alg) + envs, info = leading_boundary(envinit, state, alg) function leading_boundary_diffgauge_pullback((Δenvs′, Δinfo)) Δenvs = unthunk(Δenvs′) @@ -434,7 +434,7 @@ function _rrule( return NoTangent(), ZeroTangent(), ∂F∂envs, NoTangent() end - return envs, leading_boundary_diffgauge_pullback + return (envs, info), leading_boundary_diffgauge_pullback end # Here f is differentiated from an pre-computed SVD with fixed U, S and V diff --git a/test/boundarymps/vumps.jl b/test/boundarymps/vumps.jl index 2ec6265cf..efe544f2d 100644 --- a/test/boundarymps/vumps.jl +++ b/test/boundarymps/vumps.jl @@ -17,7 +17,7 @@ const vumps_alg = VUMPS(; alg_eigsolve=MPSKit.Defaults.alg_eigsolve(; ishermitia mps, envs, ϵ = leading_boundary(mps, T, vumps_alg) N = abs(sum(expectation_value(mps, T))) - ctm = leading_boundary( + ctm, = leading_boundary( CTMRGEnv(psi, ComplexSpace(20)), psi, SimultaneousCTMRG(; verbosity=1) ) N´ = abs(norm(psi, ctm)) @@ -33,7 +33,7 @@ end mps, envs, ϵ = leading_boundary(mps, T, vumps_alg) N = abs(prod(expectation_value(mps, T))) - ctm = leading_boundary( + ctm, = leading_boundary( CTMRGEnv(psi, ComplexSpace(20)), psi, SimultaneousCTMRG(; verbosity=1) ) N´ = abs(norm(psi, ctm)) diff --git a/test/ctmrg/fixed_iterscheme.jl b/test/ctmrg/fixed_iterscheme.jl index 208fdece0..39eccf5c0 100644 --- a/test/ctmrg/fixed_iterscheme.jl +++ b/test/ctmrg/fixed_iterscheme.jl @@ -30,7 +30,7 @@ atol = 1e-5 # initialize states Random.seed!(2394823842) psi = InfinitePEPS(2, χbond; unitcell) - env_conv1 = leading_boundary(CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg) + env_conv1, = leading_boundary(CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg) # do extra iteration to get SVD env_conv2, info = ctmrg_iteration(psi, env_conv1, ctm_alg) @@ -61,7 +61,7 @@ end Random.seed!(91283219347) psi = InfinitePEPS(2, χbond) env_init = CTMRGEnv(psi, ComplexSpace(χenv)) - env_conv1 = leading_boundary(env_init, psi, ctm_alg_iter) + env_conv1, = leading_boundary(env_init, psi, ctm_alg_iter) # do extra iteration to get SVD env_conv2_iter, info_iter = ctmrg_iteration(psi, env_conv1, ctm_alg_iter) diff --git a/test/ctmrg/flavors.jl b/test/ctmrg/flavors.jl index 5d8b440c4..1c84f7e94 100644 --- a/test/ctmrg/flavors.jl +++ b/test/ctmrg/flavors.jl @@ -19,10 +19,10 @@ projector_algs = [HalfInfiniteProjector, FullInfiniteProjector] # compute environments Random.seed!(32350283290358) psi = InfinitePEPS(2, χbond; unitcell) - env_sequential = leading_boundary( + env_sequential, = leading_boundary( CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg_sequential ) - env_simultaneous = leading_boundary( + env_simultaneous, = leading_boundary( CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg_simultaneous ) @@ -68,7 +68,7 @@ end χs = [16 17 18; 15 20 21; 14 19 22] psi = InfinitePEPS(Ds, Ds, Ds) env = CTMRGEnv(psi, rand(10:20, 3, 3), rand(10:20, 3, 3)) - env2 = leading_boundary(env, psi, ctm_alg) + env2, = leading_boundary(env, psi, ctm_alg) # check that the space is fixed @test all(space.(env.corners) .== space.(env2.corners)) diff --git a/test/ctmrg/gaugefix.jl b/test/ctmrg/gaugefix.jl index 048d875de..4bac2d163 100644 --- a/test/ctmrg/gaugefix.jl +++ b/test/ctmrg/gaugefix.jl @@ -19,7 +19,7 @@ function _pre_converge_env( Random.seed!(seed) # Seed RNG to make random environment consistent psi = InfinitePEPS(rand, T, physical_space, peps_space; unitcell) env₀ = CTMRGEnv(psi, ctm_space) - env_conv = leading_boundary(env₀, psi, SequentialCTMRG()) + env_conv, = leading_boundary(env₀, psi, SequentialCTMRG()) return env_conv, psi end @@ -44,7 +44,7 @@ end alg = ctmrg_alg(; projector_alg) env_pre, psi = preconv[(S, T, unitcell)] env_pre - env = leading_boundary(env_pre, psi, alg) + env, = leading_boundary(env_pre, psi, alg) env′, = ctmrg_iteration(psi, env, alg) env_fixed, = gauge_fix(env, env′) @test calc_elementwise_convergence(env, env_fixed) ≈ 0 atol = atol diff --git a/test/ctmrg/gradients.jl b/test/ctmrg/gradients.jl index 5659aa6d3..5d86fa386 100644 --- a/test/ctmrg/gradients.jl +++ b/test/ctmrg/gradients.jl @@ -60,7 +60,7 @@ steps = -0.01:0.005:0.01 Random.seed!(42039482030) dir = InfinitePEPS(Pspace, Vspace, Vspace) psi = InfinitePEPS(Pspace, Vspace, Vspace) - env = leading_boundary(CTMRGEnv(psi, Espace), psi, ctmrg_alg) + env, = leading_boundary(CTMRGEnv(psi, Espace), psi, ctmrg_alg) # TODO: redo this test using Manopt # alphas, fs, dfs1, dfs2 = OptimKit.optimtest( # (psi, env), diff --git a/test/heisenberg.jl b/test/heisenberg.jl index 03a748c51..d3c6722d6 100644 --- a/test/heisenberg.jl +++ b/test/heisenberg.jl @@ -36,7 +36,7 @@ end unitcell = (1, 2) H_1x2 = heisenberg_XYZ(InfiniteSquare(unitcell...)) psi_init_1x2 = InfinitePEPS(2, Dbond; unitcell) - env_init_1x2 = leading_boundary( + env_init_1x2, = leading_boundary( CTMRGEnv(psi_init_1x2, ComplexSpace(χenv)), psi_init_1x2, ctm_alg ) @@ -81,7 +81,7 @@ end # absorb weight into site tensors and CTMRG peps = InfinitePEPS(peps) envs₀ = CTMRGEnv(rand, Float64, peps, Espace) - envs = leading_boundary(envs₀, peps, SimultaneousCTMRG()) + envs, = leading_boundary(envs₀, peps, SimultaneousCTMRG()) # measure physical quantities e_site = costfun(peps, envs, ham) / (N1 * N2) diff --git a/test/j1j2_model.jl b/test/j1j2_model.jl index 3df78f881..042a2be1a 100644 --- a/test/j1j2_model.jl +++ b/test/j1j2_model.jl @@ -19,7 +19,7 @@ Random.seed!(91283219347) H = j1_j2(InfiniteSquare(); J2=0.25) psi_init = product_peps(2, χbond; noise_amp=1e-1) psi_init = symmetrize!(psi_init, RotateReflect()) -env_init = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg); +env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg); # find fixedpoint result = fixedpoint(psi_init, H, opt_alg, env_init; symmetrization=RotateReflect()) diff --git a/test/pwave.jl b/test/pwave.jl index 1007b05ac..9f82c1f69 100644 --- a/test/pwave.jl +++ b/test/pwave.jl @@ -20,7 +20,7 @@ Vspace = Vect[FermionParity](0 => χbond ÷ 2, 1 => χbond ÷ 2) Envspace = Vect[FermionParity](0 => χenv ÷ 2, 1 => χenv ÷ 2) Random.seed!(91283219347) psi_init = InfinitePEPS(Pspace, Vspace, Vspace; unitcell) -env_init = leading_boundary(CTMRGEnv(psi_init, Envspace), psi_init, ctm_alg); +env_init, = leading_boundary(CTMRGEnv(psi_init, Envspace), psi_init, ctm_alg); # find fixedpoint result = fixedpoint(psi_init, H, opt_alg, env_init) diff --git a/test/tf_ising.jl b/test/tf_ising.jl index 90f21b1b4..d8749c0f2 100644 --- a/test/tf_ising.jl +++ b/test/tf_ising.jl @@ -27,7 +27,7 @@ opt_alg = PEPSOptimize(; H = transverse_field_ising(InfiniteSquare(); g) Random.seed!(91283219347) psi_init = InfinitePEPS(2, χbond) -env_init = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg) +env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg) # find fixedpoint result = fixedpoint(psi_init, H, opt_alg, env_init) From 3b1deefc0fe4a7704844bf00f4f60e25f1c480a8 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Fri, 20 Dec 2024 15:34:57 +0100 Subject: [PATCH 13/28] Fix precompilation and deepcopy(env) in fixedpoint --- src/PEPSKit.jl | 12 +++++------- src/algorithms/ctmrg/simultaneous.jl | 2 +- src/algorithms/peps_opt.jl | 4 ++-- 3 files changed, 8 insertions(+), 10 deletions(-) diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index a32731dc5..1ddda8ef2 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -151,6 +151,7 @@ module Defaults SVDAdjoint, HalfInfiniteProjector, SimultaneousCTMRG, + RecordTruncationError, RecordConditionNumber, RecordUnitCellGradientNorm @@ -184,13 +185,10 @@ module Defaults RecordTime(; mode=:iterative) => :Time, ] const debug_group = [ - (:Iteration, "Optim %-4d "), - (:Cost, "f(x) = %.8f "), - (:GradientNorm, "‖∂f‖ = %.8f "), - (:Stepsize, "step size = %.4f "), - DebugTime(; prefix="time =", mode=:iterative), - "\n", - DebugWarnIfCostIncreases(:Always; tol=1e-10), + (:Iteration, "Optim %-4d "), + (:Cost, "f(x) = %.8f "), + (:GradientNorm, "‖∂f‖ = %.8f "), + (:Time, "time = %s"), "\n", :Stop, ] diff --git a/src/algorithms/ctmrg/simultaneous.jl b/src/algorithms/ctmrg/simultaneous.jl index e23ce83ac..8ceb714da 100644 --- a/src/algorithms/ctmrg/simultaneous.jl +++ b/src/algorithms/ctmrg/simultaneous.jl @@ -92,7 +92,7 @@ function simultaneous_projectors( P_left = map(first, projectors) P_right = map(last, projectors) S = copy(S) - truncation_error=maximum(copy(ϵ)) + truncation_error = maximum(copy(ϵ)) # TODO: This makes Zygote error on first execution? condition_number = maximum(_condition_number, S) info = (; truncation_error, condition_number, U=copy(U), S, V=copy(V)) return (P_left, P_right), info diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl index 0a91ee583..72add49b3 100644 --- a/src/algorithms/peps_opt.jl +++ b/src/algorithms/peps_opt.jl @@ -381,7 +381,7 @@ function fixedpoint( # construct cost and grad functions peps₀_vec, from_vec = to_vec(peps₀) - cache = PEPSCostFunctionCache(operator, alg, peps₀_vec, from_vec, env₀) + cache = PEPSCostFunctionCache(operator, alg, peps₀_vec, from_vec, deepcopy(env₀)) cost = cache grad = gradient_function(cache) @@ -394,7 +394,7 @@ function fixedpoint( SymmetrizeExponentialRetraction(alg.symmetrization, from_vec) end result = alg.optim_alg( - M, cost, grad, peps₀_vec; alg.optim_kwargs..., return_state=true, retraction_method + M, cost, grad, peps₀_vec; alg.optim_kwargs..., retraction_method, return_state=true ) # extract final result From 7832de589fc078420b611367217ba490a25f9db5 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Fri, 20 Dec 2024 16:49:40 +0100 Subject: [PATCH 14/28] Fix horrible typo in cost function --- src/algorithms/peps_opt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl index 72add49b3..987e69e5e 100644 --- a/src/algorithms/peps_opt.jl +++ b/src/algorithms/peps_opt.jl @@ -286,7 +286,7 @@ function cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) wh cache.alg.boundary_alg; alg_rrule=cache.alg.gradient_alg, ) - cost = expectation_value(peps, cache.operator, env) + cost = expectation_value(ψ, cache.operator, env) ignore_derivatives() do update!(cache.env, env) # update environment in-place cache.env_info = info # update environment information (truncation error, ...) From 6d3f233441de7d60a07a9ca77a45ebb8a67b83c8 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Fri, 20 Dec 2024 19:40:07 +0100 Subject: [PATCH 15/28] Improve optimization debug printing, add default stepsize --- src/PEPSKit.jl | 25 ++++++++++++------------ src/algorithms/peps_opt.jl | 39 +++++++++++++++++++++++++++++++++++++- 2 files changed, 51 insertions(+), 13 deletions(-) diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 1ddda8ef2..b2e17e624 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -114,12 +114,15 @@ Module containing default algorithm parameter values and arguments. ] ``` -- `debug_group=[...]`: Values (`DebugAction`s) that are printed during Manopt optimization +- `debug_group = [DebugPEPSOptimize(), :Stop]`: Optimization iteration info printing +- `stepsize = AdaptiveWNGradient()`: - `optim_kwargs`: All keyword arguments that are passed onto a Manopt optimization call ``` - optim_kwargs = (; stopping_criterion, record=record_group, debug=debug_group) + optim_kwargs = (; + stopping_criterion, record=record_group, debug=debug_group, stepsize + ) ``` - `fpgrad_maxiter=30`: Maximal number of iterations for computing the CTMRG fixed-point gradient @@ -153,7 +156,8 @@ module Defaults SimultaneousCTMRG, RecordTruncationError, RecordConditionNumber, - RecordUnitCellGradientNorm + RecordUnitCellGradientNorm, + DebugPEPSOptimize # CTMRG const ctmrg_tol = 1e-8 @@ -184,15 +188,11 @@ module Defaults RecordUnitCellGradientNorm() => :UnitcellGradientNorm, RecordTime(; mode=:iterative) => :Time, ] - const debug_group = [ - (:Iteration, "Optim %-4d "), - (:Cost, "f(x) = %.8f "), - (:GradientNorm, "‖∂f‖ = %.8f "), - (:Time, "time = %s"), - "\n", - :Stop, - ] - const optim_kwargs = (; stopping_criterion, record=record_group, debug=debug_group) + const debug_group = [DebugPEPSOptimize(), :Stop] + const stepsize = AdaptiveWNGradient() + const optim_kwargs = (; + stopping_criterion, record=record_group, debug=debug_group, stepsize + ) const fpgrad_maxiter = 30 const fpgrad_tol = 1e-6 const gradient_linsolver = KrylovKit.BiCGStab(; maxiter=fpgrad_maxiter, tol=fpgrad_tol) @@ -255,6 +255,7 @@ export expectation_value, costfun, product_peps, correlation_length export leading_boundary export PEPSOptimize, GeomSum, ManualIter, LinSolver export RecordTruncationError, RecordConditionNumber, RecordUnitCellGradientNorm +export DebugPEPSOptimize export fixedpoint export absorb_weight diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl index 987e69e5e..0cedebf2d 100644 --- a/src/algorithms/peps_opt.jl +++ b/src/algorithms/peps_opt.jl @@ -1,3 +1,4 @@ +using Printf using Manifolds: Manifolds using Manifolds: AbstractManifold, @@ -137,6 +138,43 @@ function (r::RecordUnitCellGradientNorm)( return Manopt.record_or_reset!(r, norm.(grad.A), i) end +""" + mutable struct DebugPEPSOptimize <: DebugAction + +Custom `DebugAction` printing for PEPS optimization runs. + +The debug info is output using `@info` and by default prints the optimization iteration, +the cost function value, the gradient norm, the last step size and the time elapsed during +the current iteration. +""" +mutable struct DebugPEPSOptimize <: DebugAction + last_time::UInt64 + DebugPEPSOptimize() = new(time_ns()) +end +function (d::DebugPEPSOptimize)( + p::AbstractManoptProblem, s::AbstractManoptSolverState, k::Int +) + time_new = time_ns() + cost = get_cost(p, get_iterate(s)) + if k == 0 + @info @sprintf("Initial f(x) = %.8f", cost) + elseif k > 0 + gradient_norm = norm(get_manifold(p), get_iterate(s), get_gradient(s)) + stepsize = get_last_stepsize(p, s, k) + time = (time_new - d.last_time) * 1e-9 + @info @sprintf( + "Optimization %d: f(x) = %.8f ‖∂f‖ = %.8f step = %.4f time = %.2f sec", + k, + cost, + gradient_norm, + stepsize, + time + ) + end + d.last_time = time_new # update time stamp + return nothing +end + """ struct PEPSOptimize{G} @@ -386,7 +424,6 @@ function fixedpoint( grad = gradient_function(cache) # optimize - # fld = scalartype(peps₀) <: Real ? Manifolds.ℝ : Manifolds.ℂ # Manopt can't optimize over ℂ? M = Euclidean(length(peps₀_vec)) retraction_method = if isnothing(alg.symmetrization) default_retraction_method(M) From 6b4a4030fc4bcfe62bf70e32bca92cd2734238ec Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Fri, 20 Dec 2024 22:37:45 +0100 Subject: [PATCH 16/28] Fix tests (mostly), fix _condition_number --- src/algorithms/ctmrg/simultaneous.jl | 7 ++++-- src/algorithms/peps_opt.jl | 36 ++++++++++++++-------------- src/utility/util.jl | 5 +--- test/ctmrg/gradients.jl | 33 ++++++++++--------------- test/heisenberg.jl | 23 +++++++++--------- test/j1j2_model.jl | 5 ++-- test/pwave.jl | 18 +++++++++----- test/tf_ising.jl | 10 ++++---- 8 files changed, 67 insertions(+), 70 deletions(-) diff --git a/src/algorithms/ctmrg/simultaneous.jl b/src/algorithms/ctmrg/simultaneous.jl index d1e7aba69..31f305678 100644 --- a/src/algorithms/ctmrg/simultaneous.jl +++ b/src/algorithms/ctmrg/simultaneous.jl @@ -50,8 +50,11 @@ end # Compute condition number σ_max / σ_min for diagonal singular value TensorMap function _condition_number(S::AbstractTensorMap) - S_diag = diag(S.data) - return maximum(S_diag) / minimum(S_diag) + # Take maximal condition number over all blocks + return maximum(blocks(S)) do (_, b) + b_diag = diag(b) + return maximum(b_diag) / minimum(b_diag) + end end @non_differentiable _condition_number(S::AbstractTensorMap) diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl index 0cedebf2d..71e082fbe 100644 --- a/src/algorithms/peps_opt.jl +++ b/src/algorithms/peps_opt.jl @@ -175,6 +175,24 @@ function (d::DebugPEPSOptimize)( return nothing end +""" + SymmetrizeExponentialRetraction <: AbstractRetractionMethod + +Exponential retraction followed by a symmetrization step. +""" +struct SymmetrizeExponentialRetraction <: AbstractRetractionMethod + symmetrization::SymmetrizationStyle + from_vec::Function +end + +function Manifolds.retract!( + M::Euclidean, p, q, X, t::Number, sr::SymmetrizeExponentialRetraction +) + v = Manifolds.retract!(M, p, q, X, t) + v_symm_peps = symmetrize!(sr.from_vec(v), sr.symmetrization) + return to_vec(v_symm_peps) +end + """ struct PEPSOptimize{G} @@ -375,24 +393,6 @@ function gradient_function(cache::PEPSCostFunctionCache{T}) where {T} end end -""" - SymmetrizeExponentialRetraction <: AbstractRetractionMethod - -Exponential retraction followed by a symmetrization step. -""" -struct SymmetrizeExponentialRetraction <: AbstractRetractionMethod - symmetrization::SymmetrizationStyle - from_vec::Function -end - -function Manifolds.retract( - M::AbstractManifold, p, X, t::Number, sr::SymmetrizeExponentialRetraction -) - q = retract(M, p, X, t, ExponentialRetraction()) - q_symm_peps = symmetrize!(sr.from_vec(q), sr.symmetrization) - return to_vec(q_symm_peps) -end - """ fixedpoint( peps₀::InfinitePEPS{T}, diff --git a/src/utility/util.jl b/src/utility/util.jl index 075f39ea2..402880add 100644 --- a/src/utility/util.jl +++ b/src/utility/util.jl @@ -56,10 +56,7 @@ function sdiag_pow(S::AbstractTensorMap, pow::Real; tol::Real=eps(scalartype(S)) tol *= norm(S, Inf) # Relative tol w.r.t. largest singular value (use norm(∘, Inf) to make differentiable) Spow = similar(S) for (k, b) in blocks(S) - copyto!( - blocks(Spow)[k], - LinearAlgebra.diagm(_safe_pow.(LinearAlgebra.diag(b), pow, tol)), - ) + copyto!(blocks(Spow)[k], diagm(_safe_pow.(diag(b), pow, tol))) end return Spow end diff --git a/test/ctmrg/gradients.jl b/test/ctmrg/gradients.jl index 5d86fa386..0fa7c419c 100644 --- a/test/ctmrg/gradients.jl +++ b/test/ctmrg/gradients.jl @@ -3,7 +3,8 @@ using Random using TensorKit using KrylovKit using PEPSKit -using Zygote +using PEPSKit: to_vec, PEPSCostFunctionCache, gradient_function +using Manopt, Manifolds ## Test models, gradmodes and CTMRG algorithm # ------------------------------------------- @@ -40,7 +41,6 @@ gradmodes = [ LinSolver(; solver=KrylovKit.BiCGStab(; tol=gradtol), iterscheme=:diffgauge), ], ] -steps = -0.01:0.005:0.01 ## Tests # ------ @@ -54,28 +54,19 @@ steps = -0.01:0.005:0.01 gms = gradmodes[i] calgs = ctmrg_algs[i] psi_init = InfinitePEPS(Pspace, Vspace, Vspace) - @testset "$ctmrg_alg and $alg_rrule" for (ctmrg_alg, alg_rrule) in - Iterators.product(calgs, gms) - # @info "optimtest of $ctmrg_alg and $alg_rrule on $(names[i])" + @testset "$ctmrg_alg and $gradient_alg" for (ctmrg_alg, gradient_alg) in + Iterators.product(calgs, gms) + @info "gradient check of $ctmrg_alg and $alg_rrule on $(names[i])" Random.seed!(42039482030) - dir = InfinitePEPS(Pspace, Vspace, Vspace) psi = InfinitePEPS(Pspace, Vspace, Vspace) env, = leading_boundary(CTMRGEnv(psi, Espace), psi, ctmrg_alg) - # TODO: redo this test using Manopt - # alphas, fs, dfs1, dfs2 = OptimKit.optimtest( - # (psi, env), - # dir; - # alpha=steps, - # retract=PEPSKit.peps_retract, - # inner=PEPSKit.real_inner, - # ) do (peps, envs) - # E, g = Zygote.withgradient(peps) do psi - # envs2 = PEPSKit.hook_pullback(leading_boundary, envs, psi, ctmrg_alg; alg_rrule) - # return costfun(psi, envs2, models[i]) - # end - # return E, only(g) - # end - # @test dfs1 ≈ dfs2 atol = 1e-2 + psi_vec, from_vec = to_vec(psi) + opt_alg = PEPSOptimize(; boundary_alg=ctmrg_alg, gradient_alg) + cache = PEPSCostFunctionCache(models[i], opt_alg, psi_vec, from_vec, deepcopy(env)) + cost = cache + grad = gradient_function(cache) + M = Euclidean(length(psi_vec)) + @test check_gradient(M, cost, grad; N=5, exactness_tol=1e-4, io=stdout) end end diff --git a/test/heisenberg.jl b/test/heisenberg.jl index d3c6722d6..55818d65d 100644 --- a/test/heisenberg.jl +++ b/test/heisenberg.jl @@ -4,7 +4,6 @@ using Accessors using TensorKit using KrylovKit using PEPSKit -using Manopt # initialize parameters Dbond = 2 @@ -23,10 +22,10 @@ E_ref = -0.6602310934799577 env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg) # optimize energy and compute correlation lengths - result = fixedpoint(psi_init, H, opt_alg, env_init) - ξ_h, ξ_v, = correlation_length(result.peps, result.env) + peps, env, E, = fixedpoint(psi_init, H, opt_alg, env_init) + ξ_h, ξ_v, = correlation_length(peps, env) - @test result.E ≈ E_ref atol = 1e-2 + @test E ≈ E_ref atol = 1e-2 @test all(@. ξ_h > 0 && ξ_v > 0) end @@ -34,18 +33,18 @@ end # initialize states Random.seed!(456) unitcell = (1, 2) - H_1x2 = heisenberg_XYZ(InfiniteSquare(unitcell...)) - psi_init_1x2 = InfinitePEPS(2, Dbond; unitcell) - env_init_1x2, = leading_boundary( - CTMRGEnv(psi_init_1x2, ComplexSpace(χenv)), psi_init_1x2, ctm_alg + H = heisenberg_XYZ(InfiniteSquare(unitcell...)) + psi_init = InfinitePEPS(2, Dbond; unitcell) + env_init, = leading_boundary( + CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg ) # optimize energy and compute correlation lengths - result_1x2 = fixedpoint(psi_init_1x2, H_1x2, opt_alg, env_init_1x2) - ξ_h_1x2, ξ_v_1x2, = correlation_length(result_1x2.peps, result_1x2.env) + peps, env, E, = fixedpoint(psi_init, H, opt_alg, env_init) + ξ_h, ξ_v, = correlation_length(peps, env) - @test result_1x2.E ≈ 2 * E_ref atol = 1e-2 - @test all(@. ξ_h_1x2 > 0 && ξ_v_1x2 > 0) + @test E ≈ 2 * E_ref atol = 1e-2 + @test all(@. ξ_h > 0 && ξ_v > 0) end @testset "Simple update into AD optimization" begin diff --git a/test/j1j2_model.jl b/test/j1j2_model.jl index 042a2be1a..f1c405ba2 100644 --- a/test/j1j2_model.jl +++ b/test/j1j2_model.jl @@ -10,8 +10,9 @@ using PEPSKit ctm_alg = SimultaneousCTMRG() opt_alg = PEPSOptimize(; boundary_alg=ctm_alg, - optimizer=LBFGS(4; gradtol=1e-3, verbosity=2), + tol=1e-3, gradient_alg=LinSolver(; iterscheme=:diffgauge), + symmetrization=RotateReflect(), ) # initialize states @@ -22,7 +23,7 @@ psi_init = symmetrize!(psi_init, RotateReflect()) env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg); # find fixedpoint -result = fixedpoint(psi_init, H, opt_alg, env_init; symmetrization=RotateReflect()) +result = fixedpoint(psi_init, H, opt_alg, env_init) ξ_h, ξ_v, = correlation_length(result.peps, result.env) # compare against Juraj Hasik's data: diff --git a/test/pwave.jl b/test/pwave.jl index 9f82c1f69..ab7439251 100644 --- a/test/pwave.jl +++ b/test/pwave.jl @@ -3,27 +3,33 @@ using Random using TensorKit using KrylovKit using PEPSKit +using Manopt +using LineSearches # Initialize parameters unitcell = (2, 2) H = pwave_superconductor(InfiniteSquare(unitcell...)) -χbond = 2 +Dbond = 2 χenv = 16 -ctm_alg = SimultaneousCTMRG(; maxiter=150) +ctm_alg = SimultaneousCTMRG() opt_alg = PEPSOptimize(; - boundary_alg=ctm_alg, optimizer=LBFGS(4; maxiter=10, gradtol=1e-3, verbosity=2) + boundary_alg=ctm_alg, + maxiter=20, + gradient_alg=LinSolver(; iterscheme=:diffgauge), + stepsize=WolfePowellLinesearch(), + memory_size=4, ) # initialize states Pspace = Vect[FermionParity](0 => 1, 1 => 1) -Vspace = Vect[FermionParity](0 => χbond ÷ 2, 1 => χbond ÷ 2) +Vspace = Vect[FermionParity](0 => Dbond ÷ 2, 1 => Dbond ÷ 2) Envspace = Vect[FermionParity](0 => χenv ÷ 2, 1 => χenv ÷ 2) Random.seed!(91283219347) psi_init = InfinitePEPS(Pspace, Vspace, Vspace; unitcell) env_init, = leading_boundary(CTMRGEnv(psi_init, Envspace), psi_init, ctm_alg); # find fixedpoint -result = fixedpoint(psi_init, H, opt_alg, env_init) +peps, env, E, = fixedpoint(psi_init, H, opt_alg, env_init) # comparison with Gaussian PEPS minimum at D=2 on 1000x1000 square lattice with aPBC -@test result.E / prod(size(psi_init)) ≈ -2.6241 atol = 5e-2 +@test E / prod(size(psi_init)) ≈ -2.6241 atol = 5e-2 diff --git a/test/tf_ising.jl b/test/tf_ising.jl index d8749c0f2..cd1d9301b 100644 --- a/test/tf_ising.jl +++ b/test/tf_ising.jl @@ -20,7 +20,7 @@ mˣ = 0.91 χenv = 16 ctm_alg = SimultaneousCTMRG() opt_alg = PEPSOptimize(; - boundary_alg=ctm_alg, optimizer=LBFGS(4; gradtol=1e-3, verbosity=2) + boundary_alg=ctm_alg, tol=1e-3, stepsize=WolfePowellBinaryLinesearch() ) # initialize states @@ -30,21 +30,21 @@ psi_init = InfinitePEPS(2, χbond) env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg) # find fixedpoint -result = fixedpoint(psi_init, H, opt_alg, env_init) -ξ_h, ξ_v, = correlation_length(result.peps, result.env) +peps, env, E, = fixedpoint(psi_init, H, opt_alg, env_init) +ξ_h, ξ_v, = correlation_length(peps, env) # compute magnetization σx = TensorMap(scalartype(psi_init)[0 1; 1 0], ℂ^2, ℂ^2) M = LocalOperator(H.lattice, (CartesianIndex(1, 1),) => σx) magn = expectation_value(result.peps, M, result.env) -@test result.E ≈ e atol = 1e-2 +@test E ≈ e atol = 1e-2 @test imag(magn) ≈ 0 atol = 1e-6 @test abs(magn) ≈ mˣ atol = 5e-2 # find fixedpoint in polarized phase and compute correlations lengths H_polar = transverse_field_ising(InfiniteSquare(); g=4.5) -result_polar = fixedpoint(psi_init, H_polar, opt_alg, env_init) +peps_polar, env_polar, E_polar, = fixedpoint(psi_init, H_polar, opt_alg, env_init) ξ_h_polar, ξ_v_polar, = correlation_length(result_polar.peps, result_polar.env) @test ξ_h_polar < ξ_h @test ξ_v_polar < ξ_v From 8872d20a2069bbb208b91c3f4906386a02c9d8c3 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Fri, 17 Jan 2025 14:42:58 +0100 Subject: [PATCH 17/28] Break up peps_opt.jl file --- src/PEPSKit.jl | 4 +- .../fixed_point_differentiation.jl | 225 +++++++ src/algorithms/optimization/manopt.jl | 116 ++++ .../optimization/peps_optimization.jl | 245 ++++++++ src/algorithms/peps_opt.jl | 587 ------------------ 5 files changed, 589 insertions(+), 588 deletions(-) create mode 100644 src/algorithms/optimization/fixed_point_differentiation.jl create mode 100644 src/algorithms/optimization/manopt.jl create mode 100644 src/algorithms/optimization/peps_optimization.jl delete mode 100644 src/algorithms/peps_opt.jl diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index b2e17e624..0ff564301 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -56,7 +56,9 @@ include("algorithms/toolbox.jl") include("utility/symmetrization.jl") -include("algorithms/peps_opt.jl") +include("algorithms/optimization/fixed_point_differentiation.jl") +include("algorithms/optimization/manopt.jl") +include("algorithms/optimization/peps_optimization.jl") """ module Defaults diff --git a/src/algorithms/optimization/fixed_point_differentiation.jl b/src/algorithms/optimization/fixed_point_differentiation.jl new file mode 100644 index 000000000..8eefc5650 --- /dev/null +++ b/src/algorithms/optimization/fixed_point_differentiation.jl @@ -0,0 +1,225 @@ +abstract type GradMode{F} end + +iterscheme(::GradMode{F}) where {F} = F + +""" + struct GeomSum(; maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol, + verbosity=0, iterscheme=Defaults.iterscheme) <: GradMode{iterscheme} + +Gradient mode for CTMRG using explicit evaluation of the geometric sum. + +With `iterscheme` the style of CTMRG iteration which is being differentiated can be chosen. +If set to `:fixed`, the differentiated CTMRG iteration is assumed to have a pre-computed +SVD of the environments with a fixed set of gauges. Alternatively, if set to `:diffgauge`, +the differentiated iteration consists of a CTMRG iteration and a subsequent gauge fixing step, +such that `gauge_fix` will also be differentiated everytime a CTMRG derivative is computed. +""" +struct GeomSum{F} <: GradMode{F} + maxiter::Int + tol::Real + verbosity::Int +end +function GeomSum(; + maxiter=Defaults.fpgrad_maxiter, + tol=Defaults.fpgrad_tol, + verbosity=0, + iterscheme=Defaults.iterscheme, +) + return GeomSum{iterscheme}(maxiter, tol, verbosity) +end + +""" + struct ManualIter(; maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol, + verbosity=0, iterscheme=Defaults.iterscheme) <: GradMode{iterscheme} + +Gradient mode for CTMRG using manual iteration to solve the linear problem. + +With `iterscheme` the style of CTMRG iteration which is being differentiated can be chosen. +If set to `:fixed`, the differentiated CTMRG iteration is assumed to have a pre-computed +SVD of the environments with a fixed set of gauges. Alternatively, if set to `:diffgauge`, +the differentiated iteration consists of a CTMRG iteration and a subsequent gauge fixing step, +such that `gauge_fix` will also be differentiated everytime a CTMRG derivative is computed. +""" +struct ManualIter{F} <: GradMode{F} + maxiter::Int + tol::Real + verbosity::Int +end +function ManualIter(; + maxiter=Defaults.fpgrad_maxiter, + tol=Defaults.fpgrad_tol, + verbosity=0, + iterscheme=Defaults.iterscheme, +) + return ManualIter{iterscheme}(maxiter, tol, verbosity) +end + +""" + struct LinSolver(; solver=KrylovKit.GMRES(), iterscheme=Defaults.iterscheme) <: GradMode{iterscheme} + +Gradient mode wrapper around `KrylovKit.LinearSolver` for solving the gradient linear +problem using iterative solvers. + +With `iterscheme` the style of CTMRG iteration which is being differentiated can be chosen. +If set to `:fixed`, the differentiated CTMRG iteration is assumed to have a pre-computed +SVD of the environments with a fixed set of gauges. Alternatively, if set to `:diffgauge`, +the differentiated iteration consists of a CTMRG iteration and a subsequent gauge fixing step, +such that `gauge_fix` will also be differentiated everytime a CTMRG derivative is computed. +""" +struct LinSolver{F} <: GradMode{F} + solver::KrylovKit.LinearSolver +end +function LinSolver(; + solver=KrylovKit.BiCGStab(; maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol), + iterscheme=Defaults.iterscheme, +) + return LinSolver{iterscheme}(solver) +end + + +#= +Evaluating the gradient of the cost function for CTMRG: +- The gradient of the cost function for CTMRG can be computed using automatic differentiation (AD) or explicit evaluation of the geometric sum. +- With AD, the gradient is computed by differentiating the cost function with respect to the PEPS tensors, including computing the environment tensors. +- With explicit evaluation of the geometric sum, the gradient is computed by differentiating the cost function with the environment kept fixed, and then manually adding the gradient contributions from the environments. +=# + +function _rrule( + gradmode::GradMode{:diffgauge}, + config::RuleConfig, + ::typeof(MPSKit.leading_boundary), + envinit, + state, + alg::CTMRGAlgorithm, +) + envs, info = leading_boundary(envinit, state, alg) + + function leading_boundary_diffgauge_pullback((Δenvs′, Δinfo)) + Δenvs = unthunk(Δenvs′) + + # find partial gradients of gauge_fixed single CTMRG iteration + f(A, x) = gauge_fix(x, ctmrg_iteration(A, x, alg)[1])[1] + _, env_vjp = rrule_via_ad(config, f, state, envs) + + # evaluate the geometric sum + ∂f∂A(x)::typeof(state) = env_vjp(x)[2] + ∂f∂x(x)::typeof(envs) = env_vjp(x)[3] + ∂F∂envs = fpgrad(Δenvs, ∂f∂x, ∂f∂A, Δenvs, gradmode) + + return NoTangent(), ZeroTangent(), ∂F∂envs, NoTangent() + end + + return (envs, info), leading_boundary_diffgauge_pullback +end + +# Here f is differentiated from an pre-computed SVD with fixed U, S and V +function _rrule( + gradmode::GradMode{:fixed}, + config::RuleConfig, + ::typeof(MPSKit.leading_boundary), + envinit, + state, + alg::SimultaneousCTMRG, +) + @assert !isnothing(alg.projector_alg.svd_alg.rrule_alg) + envs, = leading_boundary(envinit, state, alg) + envs_conv, info = ctmrg_iteration(state, envs, alg) + envs_fixed, signs = gauge_fix(envs, envs_conv) + + # Fix SVD + Ufix, Vfix = fix_relative_phases(info.U, info.V, signs) + svd_alg_fixed = SVDAdjoint(; + fwd_alg=FixedSVD(Ufix, info.S, Vfix), rrule_alg=alg.projector_alg.svd_alg.rrule_alg + ) + alg_fixed = @set alg.projector_alg.svd_alg = svd_alg_fixed + alg_fixed = @set alg_fixed.projector_alg.trscheme = notrunc() + + function leading_boundary_fixed_pullback((Δenvs′, Δinfo)) + Δenvs = unthunk(Δenvs′) + + f(A, x) = fix_global_phases(x, ctmrg_iteration(A, x, alg_fixed)[1]) + _, env_vjp = rrule_via_ad(config, f, state, envs_fixed) + + # evaluate the geometric sum + ∂f∂A(x)::typeof(state) = env_vjp(x)[2] + ∂f∂x(x)::typeof(envs) = env_vjp(x)[3] + ∂F∂envs = fpgrad(Δenvs, ∂f∂x, ∂f∂A, Δenvs, gradmode) + + return NoTangent(), ZeroTangent(), ∂F∂envs, NoTangent() + end + + return (envs_fixed, info), leading_boundary_fixed_pullback +end + +@doc """ + fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y0, alg) + +Compute the gradient of the cost function for CTMRG by solving the following equation: + +dx = ∑ₙ (∂f∂x)ⁿ ∂f∂A dA = (1 - ∂f∂x)⁻¹ ∂f∂A dA + +where `∂F∂x` is the gradient of the cost function with respect to the PEPS tensors, `∂f∂x` +is the partial gradient of the CTMRG iteration with respect to the environment tensors, +`∂f∂A` is the partial gradient of the CTMRG iteration with respect to the PEPS tensors, and +`y0` is the initial guess for the fixed-point iteration. The function returns the gradient +`dx` of the fixed-point iteration. +""" +fpgrad + +# TODO: can we construct an implementation that does not need to evaluate the vjp +# twice if both ∂f∂A and ∂f∂x are needed? +function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, _, alg::GeomSum) + g = ∂F∂x + dx = ∂f∂A(g) # n = 0 term: ∂F∂x ∂f∂A + ϵ = 2 * alg.tol + for i in 1:(alg.maxiter) + g = ∂f∂x(g) + Σₙ = ∂f∂A(g) + dx += Σₙ + ϵnew = norm(Σₙ) # TODO: normalize this error? + Δϵ = ϵ - ϵnew + alg.verbosity > 1 && + @printf("Gradient iter: %3d ‖Σₙ‖: %.2e Δ‖Σₙ‖: %.2e\n", i, ϵnew, Δϵ) + ϵ = ϵnew + + ϵ < alg.tol && break + if alg.verbosity > 0 && i == alg.maxiter + @warn "gradient fixed-point iteration reached maximal number of iterations at ‖Σₙ‖ = $ϵ" + end + end + return dx +end + +function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y₀, alg::ManualIter) + y = deepcopy(y₀) # Do not mutate y₀ + dx = ∂f∂A(y) + ϵ = 1.0 + for i in 1:(alg.maxiter) + y′ = ∂F∂x + ∂f∂x(y) + + dxnew = ∂f∂A(y′) + ϵnew = norm(dxnew - dx) + Δϵ = ϵ - ϵnew + alg.verbosity > 1 && @printf( + "Gradient iter: %3d ‖Cᵢ₊₁-Cᵢ‖/N: %.2e Δ‖Cᵢ₊₁-Cᵢ‖/N: %.2e\n", i, ϵnew, Δϵ + ) + y = y′ + dx = dxnew + ϵ = ϵnew + + ϵ < alg.tol && break + if alg.verbosity > 0 && i == alg.maxiter + @warn "gradient fixed-point iteration reached maximal number of iterations at ‖Cᵢ₊₁-Cᵢ‖ = $ϵ" + end + end + return dx +end + +function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y₀, alg::LinSolver) + y, info = linsolve(∂f∂x, ∂F∂x, y₀, alg.solver, 1, -1) + if alg.solver.verbosity > 0 && info.converged != 1 + @warn("gradient fixed-point iteration reached maximal number of iterations:", info) + end + + return ∂f∂A(y) +end \ No newline at end of file diff --git a/src/algorithms/optimization/manopt.jl b/src/algorithms/optimization/manopt.jl new file mode 100644 index 000000000..142f4640c --- /dev/null +++ b/src/algorithms/optimization/manopt.jl @@ -0,0 +1,116 @@ +using Printf +using Manifolds: Manifolds +using Manifolds: + AbstractManifold, + AbstractRetractionMethod, + Euclidean, + default_retraction_method, + retract +using Manopt + +""" + mutable struct RecordTruncationError <: RecordAction + +Record the maximal truncation error of all `boundary_alg` runs of the corresponding +optimization iteration. +""" +mutable struct RecordTruncationError <: RecordAction + recorded_values::Vector{Float64} + RecordTruncationError() = new(Vector{Float64}()) +end +function (r::RecordTruncationError)( + p::AbstractManoptProblem, ::AbstractManoptSolverState, i::Int +) + cache = Manopt.get_cost_function(get_objective(p)) + return Manopt.record_or_reset!(r, cache.env_info.truncation_error, i) +end + +""" + mutable struct RecordConditionNumber <: RecordAction + +Record the maximal condition number of all `boundary_alg` runs of the corresponding +optimization iteration. +""" +mutable struct RecordConditionNumber <: RecordAction + recorded_values::Vector{Float64} + RecordConditionNumber() = new(Vector{Float64}()) +end +function (r::RecordConditionNumber)( + p::AbstractManoptProblem, ::AbstractManoptSolverState, i::Int +) + cache = Manopt.get_cost_function(get_objective(p)) + return Manopt.record_or_reset!(r, cache.env_info.condition_number, i) +end + +""" + mutable struct RecordUnitCellGradientNorm <: RecordAction + +Record the PEPS gradient norms unit cell entry-wise, i.e. an array +of norms `norm.(peps.A)`. +""" +mutable struct RecordUnitCellGradientNorm <: RecordAction + recorded_values::Vector{Matrix{Float64}} + RecordUnitCellGradientNorm() = new(Vector{Matrix{Float64}}()) +end +function (r::RecordUnitCellGradientNorm)( + p::AbstractManoptProblem, s::AbstractManoptSolverState, i::Int +) + cache = Manopt.get_cost_function(get_objective(p)) + grad = cache.from_vec(get_gradient(s)) + return Manopt.record_or_reset!(r, norm.(grad.A), i) +end + +""" + mutable struct DebugPEPSOptimize <: DebugAction + +Custom `DebugAction` printing for PEPS optimization runs. + +The debug info is output using `@info` and by default prints the optimization iteration, +the cost function value, the gradient norm, the last step size and the time elapsed during +the current iteration. +""" +mutable struct DebugPEPSOptimize <: DebugAction + last_time::UInt64 + DebugPEPSOptimize() = new(time_ns()) +end +function (d::DebugPEPSOptimize)( + p::AbstractManoptProblem, s::AbstractManoptSolverState, k::Int +) + time_new = time_ns() + cost = get_cost(p, get_iterate(s)) + if k == 0 + @info @sprintf("Initial f(x) = %.8f", cost) + elseif k > 0 + gradient_norm = norm(get_manifold(p), get_iterate(s), get_gradient(s)) + stepsize = get_last_stepsize(p, s, k) + time = (time_new - d.last_time) * 1e-9 + @info @sprintf( + "Optimization %d: f(x) = %.8f ‖∂f‖ = %.8f step = %.4f time = %.2f sec", + k, + cost, + gradient_norm, + stepsize, + time + ) + end + d.last_time = time_new # update time stamp + return nothing +end + +""" + SymmetrizeExponentialRetraction <: AbstractRetractionMethod + +Exponential retraction followed by a symmetrization step. +""" +struct SymmetrizeExponentialRetraction <: AbstractRetractionMethod + symmetrization::SymmetrizationStyle + from_vec::Function +end + +function Manifolds.retract!( + M::Euclidean, p, q, X, t::Number, sr::SymmetrizeExponentialRetraction +) + v = Manifolds.retract!(M, p, q, X, t) + v_symm_peps = symmetrize!(sr.from_vec(v), sr.symmetrization) + return to_vec(v_symm_peps) +end diff --git a/src/algorithms/optimization/peps_optimization.jl b/src/algorithms/optimization/peps_optimization.jl new file mode 100644 index 000000000..93bdd3b85 --- /dev/null +++ b/src/algorithms/optimization/peps_optimization.jl @@ -0,0 +1,245 @@ +""" + struct PEPSOptimize{G} + +Algorithm struct for PEPS optimization using automatic differentiation. + +# Fields +- `boundary_alg::CTMRGAlgorithm`: algorithm for determining the PEPS environment +- `optim_alg::Function`: Manopt optimization algorithm +- `optim_kwargs::NamedTuple`: Keyword arguments provided to the Manopt `optim_alg` call +- `gradient_alg::G`: Algorithm computing the cost function gradient in reverse-mode +- `reuse_env::Bool`: If `true` the previous environment is used to initialize the next + `leading_boundary` call +- `symmetrization::Union{Nothing,SymmetrizationStyle}`: Symmetrize the PEPS and PEPS + gradient after each optimization iteration (does nothing if `nothing` is provided) +""" +struct PEPSOptimize{G} + boundary_alg::CTMRGAlgorithm + optim_alg::Function + optim_kwargs::NamedTuple + gradient_alg::G + reuse_env::Bool + # reuse_env_tol::Float64 # TODO: add option for reuse tolerance + symmetrization::Union{Nothing,SymmetrizationStyle} + + function PEPSOptimize( # Inner constructor to prohibit illegal setting combinations + boundary_alg::CTMRGAlgorithm, + optim_alg, + optim_kwargs, + gradient_alg::G, + reuse_env, + symmetrization, + ) where {G} + if gradient_alg isa GradMode + if boundary_alg isa SequentialCTMRG && iterscheme(gradient_alg) === :fixed + throw(ArgumentError(":sequential and :fixed are not compatible")) + end + end + return new{G}( + boundary_alg, optim_alg, optim_kwargs, gradient_alg, reuse_env, symmetrization + ) + end +end + +""" + PEPSOptimize(; + boundary_alg=Defaults.ctmrg_alg, + optim_alg=Defaults.optim_alg, + maxiter=Defaults.optim_maxiter, + tol=Defaults.optim_tol, + gradient_alg=Defaults.gradient_alg, + reuse_env=Defaults.reuse_env, + symmetrization=nothing, + kwargs..., + ) + +Convenience keyword argument constructor for `PEPSOptimize` algorithms. +Here, `maxiter` and `tol` are passed onto `StopAfterIteration` and `StopWhenGradientNormLess` +stopping criteria, respectively. Additionally, any keyword arguments can be provided which +are then stored inside `optim_kwargs` and passed to the Manopt optimization call, such that +that all arguments of the respective `optim_alg` can be used. +""" +function PEPSOptimize(; + boundary_alg=Defaults.ctmrg_alg, + optim_alg=Defaults.optim_alg, + maxiter=Defaults.optim_maxiter, + tol=Defaults.optim_tol, + gradient_alg=Defaults.gradient_alg, + reuse_env=Defaults.reuse_env, + symmetrization=nothing, + kwargs..., +) + stopping_criterion = StopAfterIteration(maxiter) | StopWhenGradientNormLess(tol) + optim_kwargs = merge(Defaults.optim_kwargs, (; stopping_criterion, kwargs...)) + return PEPSOptimize( + boundary_alg, optim_alg, optim_kwargs, gradient_alg, reuse_env, symmetrization + ) +end + +""" + mutable struct PEPSCostFunctionCache{T} + +Stores objects used for computing PEPS cost functions during optimization that are +needed apart from the PEPS that is being optimized. + +# Fields +- `operator::LocalOperator`: cost function operator +- `alg::PEPSOptimize`: optimization parameters +- `env::CTMRGEnv`: environment of the current PEPS +- `from_vec`: map which returns vectorized PEPS as an `InfinitePEPS` +- `peps_vec::Vector{T}`: current vectorized PEPS +- `grad_vec::Vector{T}`: current vectorized gradient +- `cost::Float64`: current cost function value +- `env_info::NamedTuple`: return info of `leading_boundary` used by `RecordAction`s +""" +mutable struct PEPSCostFunctionCache{T} + operator::LocalOperator + alg::PEPSOptimize + env::CTMRGEnv + from_vec + peps_vec::Vector{T} + grad_vec::Vector{T} + cost::Float64 + env_info::NamedTuple +end + +""" + PEPSCostFunctionCache( + operator::LocalOperator, alg::PEPSOptimize, peps_vec::Vector, from_vec, env::CTMRGEnv + ) + +Initialize a `PEPSCostFunctionCache` using `peps_vec` from which the vector dimension +and scalartype are derived. +""" +function PEPSCostFunctionCache( + operator::LocalOperator, alg::PEPSOptimize, peps_vec::Vector, from_vec, env::CTMRGEnv +) + return PEPSCostFunctionCache( + operator, + alg, + env, + from_vec, + similar(peps_vec), + similar(peps_vec), + 0.0, + (; truncation_error=0.0, condition_number=1.0), + ) +end + +""" + cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) where {T} + +Update the cost and gradient of the `PEPSCostFunctionCache` with respect to the new point +`peps_vec`. +""" +function cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) where {T} + cache.peps_vec .= peps_vec # update point in manifold + peps = cache.from_vec(peps_vec) # convert back to InfinitePEPS + env₀ = + cache.alg.reuse_env ? cache.env : CTMRGEnv(randn, scalartype(cache.env), cache.env) + + # compute cost and gradient + cost, grads = withgradient(peps) do ψ + env, info = hook_pullback( + leading_boundary, + env₀, + ψ, + cache.alg.boundary_alg; + alg_rrule=cache.alg.gradient_alg, + ) + cost = expectation_value(ψ, cache.operator, env) + ignore_derivatives() do + update!(cache.env, env) # update environment in-place + cache.env_info = info # update environment information (truncation error, ...) + isapprox(imag(cost), 0; atol=sqrt(eps(real(cost)))) || + @warn "Expectation value is not real: $cost." + end + return real(cost) + end + grad = only(grads) # `withgradient` returns tuple of gradients `grads` + + # symmetrize gradient + if !isnothing(cache.alg.symmetrization) + grad = symmetrize!(grad, cache.alg.symmetrization) + end + + cache.cost = cost # update cost function value + cache.grad_vec .= to_vec(grad)[1] # update vectorized gradient + return cache.cost, cache.grad_vec +end + +""" + (cache::PEPSCostFunctionCache{T})(::Euclidean, peps_vec::Vector{T}) where {T} + +Return the cost of `cache` and recompute if `peps_vec` is a new point. +""" +function (cache::PEPSCostFunctionCache{T})(::Euclidean, peps_vec::Vector{T}) where {T} + # Note that it is necessary to implement the cost function as a functor so that the + # `PEPSCostFunctionCache` is available through `get_objective(::AbstractManoptProblem)` + # to the `RecordAction`s during optimization + if !(peps_vec == cache.peps_vec) # update cache if at new point + cost_and_grad!(cache, peps_vec) + end + return cache.cost +end + +""" + gradient_function(cache::PEPSCostFunctionCache{T}) where {T} + +Get the gradient function of `cache` which returns the gradient vector +and recomputes it if the provided point is new. +""" +function gradient_function(cache::PEPSCostFunctionCache{T}) where {T} + return function gradient_function(::Euclidean, peps_vec::Vector{T}) + if !(peps_vec == cache.peps_vec) # update cache if at new point + cost_and_grad!(cache, peps_vec) + end + return cache.grad_vec + end +end + +""" + fixedpoint( + peps₀::InfinitePEPS{T}, + operator::LocalOperator, + alg::PEPSOptimize, + env₀::CTMRGEnv=CTMRGEnv(peps₀, field(T)^20); + ) where {T} + +Optimize a PEPS starting from `peps₀` and `env₀` by minimizing the cost function given +by expectation value of `operator`. All optimization parameters are provided through `alg`. +""" +function fixedpoint( + peps₀::InfinitePEPS{T}, + operator::LocalOperator, + alg::PEPSOptimize, + env₀::CTMRGEnv=CTMRGEnv(peps₀, field(T)^20); +) where {T} + if scalartype(env₀) <: Real + env₀ = complex(env₀) + @warn "the provided real environment was converted to a complex environment since \ + :fixed mode generally produces complex gauges; use :diffgauge mode instead to work \ + with purely real environments" + end + + # construct cost and grad functions + peps₀_vec, from_vec = to_vec(peps₀) + cache = PEPSCostFunctionCache(operator, alg, peps₀_vec, from_vec, deepcopy(env₀)) + cost = cache + grad = gradient_function(cache) + + # optimize + M = Euclidean(length(peps₀_vec)) + retraction_method = if isnothing(alg.symmetrization) + default_retraction_method(M) + else + SymmetrizeExponentialRetraction(alg.symmetrization, from_vec) + end + result = alg.optim_alg( + M, cost, grad, peps₀_vec; alg.optim_kwargs..., retraction_method, return_state=true + ) + + # extract final result + peps_final = from_vec(get_solver_result(result)) + return peps_final, cost.env, cache.cost, result +end diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl deleted file mode 100644 index 71e082fbe..000000000 --- a/src/algorithms/peps_opt.jl +++ /dev/null @@ -1,587 +0,0 @@ -using Printf -using Manifolds: Manifolds -using Manifolds: - AbstractManifold, - AbstractRetractionMethod, - Euclidean, - default_retraction_method, - retract -using Manopt - -abstract type GradMode{F} end - -iterscheme(::GradMode{F}) where {F} = F - -""" - struct GeomSum(; maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol, - verbosity=0, iterscheme=Defaults.iterscheme) <: GradMode{iterscheme} - -Gradient mode for CTMRG using explicit evaluation of the geometric sum. - -With `iterscheme` the style of CTMRG iteration which is being differentiated can be chosen. -If set to `:fixed`, the differentiated CTMRG iteration is assumed to have a pre-computed -SVD of the environments with a fixed set of gauges. Alternatively, if set to `:diffgauge`, -the differentiated iteration consists of a CTMRG iteration and a subsequent gauge fixing step, -such that `gauge_fix` will also be differentiated everytime a CTMRG derivative is computed. -""" -struct GeomSum{F} <: GradMode{F} - maxiter::Int - tol::Real - verbosity::Int -end -function GeomSum(; - maxiter=Defaults.fpgrad_maxiter, - tol=Defaults.fpgrad_tol, - verbosity=0, - iterscheme=Defaults.iterscheme, -) - return GeomSum{iterscheme}(maxiter, tol, verbosity) -end - -""" - struct ManualIter(; maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol, - verbosity=0, iterscheme=Defaults.iterscheme) <: GradMode{iterscheme} - -Gradient mode for CTMRG using manual iteration to solve the linear problem. - -With `iterscheme` the style of CTMRG iteration which is being differentiated can be chosen. -If set to `:fixed`, the differentiated CTMRG iteration is assumed to have a pre-computed -SVD of the environments with a fixed set of gauges. Alternatively, if set to `:diffgauge`, -the differentiated iteration consists of a CTMRG iteration and a subsequent gauge fixing step, -such that `gauge_fix` will also be differentiated everytime a CTMRG derivative is computed. -""" -struct ManualIter{F} <: GradMode{F} - maxiter::Int - tol::Real - verbosity::Int -end -function ManualIter(; - maxiter=Defaults.fpgrad_maxiter, - tol=Defaults.fpgrad_tol, - verbosity=0, - iterscheme=Defaults.iterscheme, -) - return ManualIter{iterscheme}(maxiter, tol, verbosity) -end - -""" - struct LinSolver(; solver=KrylovKit.GMRES(), iterscheme=Defaults.iterscheme) <: GradMode{iterscheme} - -Gradient mode wrapper around `KrylovKit.LinearSolver` for solving the gradient linear -problem using iterative solvers. - -With `iterscheme` the style of CTMRG iteration which is being differentiated can be chosen. -If set to `:fixed`, the differentiated CTMRG iteration is assumed to have a pre-computed -SVD of the environments with a fixed set of gauges. Alternatively, if set to `:diffgauge`, -the differentiated iteration consists of a CTMRG iteration and a subsequent gauge fixing step, -such that `gauge_fix` will also be differentiated everytime a CTMRG derivative is computed. -""" -struct LinSolver{F} <: GradMode{F} - solver::KrylovKit.LinearSolver -end -function LinSolver(; - solver=KrylovKit.BiCGStab(; maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol), - iterscheme=Defaults.iterscheme, -) - return LinSolver{iterscheme}(solver) -end - -""" - mutable struct RecordTruncationError <: RecordAction - -Record the maximal truncation error of all `boundary_alg` runs of the corresponding -optimization iteration. -""" -mutable struct RecordTruncationError <: RecordAction - recorded_values::Vector{Float64} - RecordTruncationError() = new(Vector{Float64}()) -end -function (r::RecordTruncationError)( - p::AbstractManoptProblem, ::AbstractManoptSolverState, i::Int -) - cache = Manopt.get_cost_function(get_objective(p)) - return Manopt.record_or_reset!(r, cache.env_info.truncation_error, i) -end - -""" - mutable struct RecordConditionNumber <: RecordAction - -Record the maximal condition number of all `boundary_alg` runs of the corresponding -optimization iteration. -""" -mutable struct RecordConditionNumber <: RecordAction - recorded_values::Vector{Float64} - RecordConditionNumber() = new(Vector{Float64}()) -end -function (r::RecordConditionNumber)( - p::AbstractManoptProblem, ::AbstractManoptSolverState, i::Int -) - cache = Manopt.get_cost_function(get_objective(p)) - return Manopt.record_or_reset!(r, cache.env_info.condition_number, i) -end - -""" - mutable struct RecordUnitCellGradientNorm <: RecordAction - -Record the PEPS gradient norms unit cell entry-wise, i.e. an array -of norms `norm.(peps.A)`. -""" -mutable struct RecordUnitCellGradientNorm <: RecordAction - recorded_values::Vector{Matrix{Float64}} - RecordUnitCellGradientNorm() = new(Vector{Matrix{Float64}}()) -end -function (r::RecordUnitCellGradientNorm)( - p::AbstractManoptProblem, s::AbstractManoptSolverState, i::Int -) - cache = Manopt.get_cost_function(get_objective(p)) - grad = cache.from_vec(get_gradient(s)) - return Manopt.record_or_reset!(r, norm.(grad.A), i) -end - -""" - mutable struct DebugPEPSOptimize <: DebugAction - -Custom `DebugAction` printing for PEPS optimization runs. - -The debug info is output using `@info` and by default prints the optimization iteration, -the cost function value, the gradient norm, the last step size and the time elapsed during -the current iteration. -""" -mutable struct DebugPEPSOptimize <: DebugAction - last_time::UInt64 - DebugPEPSOptimize() = new(time_ns()) -end -function (d::DebugPEPSOptimize)( - p::AbstractManoptProblem, s::AbstractManoptSolverState, k::Int -) - time_new = time_ns() - cost = get_cost(p, get_iterate(s)) - if k == 0 - @info @sprintf("Initial f(x) = %.8f", cost) - elseif k > 0 - gradient_norm = norm(get_manifold(p), get_iterate(s), get_gradient(s)) - stepsize = get_last_stepsize(p, s, k) - time = (time_new - d.last_time) * 1e-9 - @info @sprintf( - "Optimization %d: f(x) = %.8f ‖∂f‖ = %.8f step = %.4f time = %.2f sec", - k, - cost, - gradient_norm, - stepsize, - time - ) - end - d.last_time = time_new # update time stamp - return nothing -end - -""" - SymmetrizeExponentialRetraction <: AbstractRetractionMethod - -Exponential retraction followed by a symmetrization step. -""" -struct SymmetrizeExponentialRetraction <: AbstractRetractionMethod - symmetrization::SymmetrizationStyle - from_vec::Function -end - -function Manifolds.retract!( - M::Euclidean, p, q, X, t::Number, sr::SymmetrizeExponentialRetraction -) - v = Manifolds.retract!(M, p, q, X, t) - v_symm_peps = symmetrize!(sr.from_vec(v), sr.symmetrization) - return to_vec(v_symm_peps) -end - -""" - struct PEPSOptimize{G} - -Algorithm struct for PEPS optimization using automatic differentiation. - -# Fields -- `boundary_alg::CTMRGAlgorithm`: algorithm for determining the PEPS environment -- `optim_alg::Function`: Manopt optimization algorithm -- `optim_kwargs::NamedTuple`: Keyword arguments provided to the Manopt `optim_alg` call -- `gradient_alg::G`: Algorithm computing the cost function gradient in reverse-mode -- `reuse_env::Bool`: If `true` the previous environment is used to initialize the next - `leading_boundary` call -- `symmetrization::Union{Nothing,SymmetrizationStyle}`: Symmetrize the PEPS and PEPS - gradient after each optimization iteration (does nothing if `nothing` is provided) -""" -struct PEPSOptimize{G} - boundary_alg::CTMRGAlgorithm - optim_alg::Function - optim_kwargs::NamedTuple - gradient_alg::G - reuse_env::Bool - # reuse_env_tol::Float64 # TODO: add option for reuse tolerance - symmetrization::Union{Nothing,SymmetrizationStyle} - - function PEPSOptimize( # Inner constructor to prohibit illegal setting combinations - boundary_alg::CTMRGAlgorithm, - optim_alg, - optim_kwargs, - gradient_alg::G, - reuse_env, - symmetrization, - ) where {G} - if gradient_alg isa GradMode - if boundary_alg isa SequentialCTMRG && iterscheme(gradient_alg) === :fixed - throw(ArgumentError(":sequential and :fixed are not compatible")) - end - end - return new{G}( - boundary_alg, optim_alg, optim_kwargs, gradient_alg, reuse_env, symmetrization - ) - end -end - -""" - PEPSOptimize(; - boundary_alg=Defaults.ctmrg_alg, - optim_alg=Defaults.optim_alg, - maxiter=Defaults.optim_maxiter, - tol=Defaults.optim_tol, - gradient_alg=Defaults.gradient_alg, - reuse_env=Defaults.reuse_env, - symmetrization=nothing, - kwargs..., - ) - -Convenience keyword argument constructor for `PEPSOptimize` algorithms. -Here, `maxiter` and `tol` are passed onto `StopAfterIteration` and `StopWhenGradientNormLess` -stopping criteria, respectively. Additionally, any keyword arguments can be provided which -are then stored inside `optim_kwargs` and passed to the Manopt optimization call, such that -that all arguments of the respective `optim_alg` can be used. -""" -function PEPSOptimize(; - boundary_alg=Defaults.ctmrg_alg, - optim_alg=Defaults.optim_alg, - maxiter=Defaults.optim_maxiter, - tol=Defaults.optim_tol, - gradient_alg=Defaults.gradient_alg, - reuse_env=Defaults.reuse_env, - symmetrization=nothing, - kwargs..., -) - stopping_criterion = StopAfterIteration(maxiter) | StopWhenGradientNormLess(tol) - optim_kwargs = merge(Defaults.optim_kwargs, (; stopping_criterion, kwargs...)) - return PEPSOptimize( - boundary_alg, optim_alg, optim_kwargs, gradient_alg, reuse_env, symmetrization - ) -end - -""" - mutable struct PEPSCostFunctionCache{T} - -Stores objects used for computing PEPS cost functions during optimization that are -needed apart from the PEPS that is being optimized. - -# Fields -- `operator::LocalOperator`: cost function operator -- `alg::PEPSOptimize`: optimization parameters -- `env::CTMRGEnv`: environment of the current PEPS -- `from_vec`: map which returns vectorized PEPS as an `InfinitePEPS` -- `peps_vec::Vector{T}`: current vectorized PEPS -- `grad_vec::Vector{T}`: current vectorized gradient -- `cost::Float64`: current cost function value -- `env_info::NamedTuple`: return info of `leading_boundary` used by `RecordAction`s -""" -mutable struct PEPSCostFunctionCache{T} - operator::LocalOperator - alg::PEPSOptimize - env::CTMRGEnv - from_vec - peps_vec::Vector{T} - grad_vec::Vector{T} - cost::Float64 - env_info::NamedTuple -end - -""" - PEPSCostFunctionCache( - operator::LocalOperator, alg::PEPSOptimize, peps_vec::Vector, from_vec, env::CTMRGEnv - ) - -Initialize a `PEPSCostFunctionCache` using `peps_vec` from which the vector dimension -and scalartype are derived. -""" -function PEPSCostFunctionCache( - operator::LocalOperator, alg::PEPSOptimize, peps_vec::Vector, from_vec, env::CTMRGEnv -) - return PEPSCostFunctionCache( - operator, - alg, - env, - from_vec, - similar(peps_vec), - similar(peps_vec), - 0.0, - (; truncation_error=0.0, condition_number=1.0), - ) -end - -""" - cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) where {T} - -Update the cost and gradient of the `PEPSCostFunctionCache` with respect to the new point -`peps_vec`. -""" -function cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) where {T} - cache.peps_vec .= peps_vec # update point in manifold - peps = cache.from_vec(peps_vec) # convert back to InfinitePEPS - env₀ = - cache.alg.reuse_env ? cache.env : CTMRGEnv(randn, scalartype(cache.env), cache.env) - - # compute cost and gradient - cost, grads = withgradient(peps) do ψ - env, info = hook_pullback( - leading_boundary, - env₀, - ψ, - cache.alg.boundary_alg; - alg_rrule=cache.alg.gradient_alg, - ) - cost = expectation_value(ψ, cache.operator, env) - ignore_derivatives() do - update!(cache.env, env) # update environment in-place - cache.env_info = info # update environment information (truncation error, ...) - isapprox(imag(cost), 0; atol=sqrt(eps(real(cost)))) || - @warn "Expectation value is not real: $cost." - end - return real(cost) - end - grad = only(grads) # `withgradient` returns tuple of gradients `grads` - - # symmetrize gradient - if !isnothing(cache.alg.symmetrization) - grad = symmetrize!(grad, cache.alg.symmetrization) - end - - cache.cost = cost # update cost function value - cache.grad_vec .= to_vec(grad)[1] # update vectorized gradient - return cache.cost, cache.grad_vec -end - -""" - (cache::PEPSCostFunctionCache{T})(::Euclidean, peps_vec::Vector{T}) where {T} - -Return the cost of `cache` and recompute if `peps_vec` is a new point. -""" -function (cache::PEPSCostFunctionCache{T})(::Euclidean, peps_vec::Vector{T}) where {T} - # Note that it is necessary to implement the cost function as a functor so that the - # `PEPSCostFunctionCache` is available through `get_objective(::AbstractManoptProblem)` - # to the `RecordAction`s during optimization - if !(peps_vec == cache.peps_vec) # update cache if at new point - cost_and_grad!(cache, peps_vec) - end - return cache.cost -end - -""" - gradient_function(cache::PEPSCostFunctionCache{T}) where {T} - -Get the gradient function of `cache` which returns the gradient vector -and recomputes it if the provided point is new. -""" -function gradient_function(cache::PEPSCostFunctionCache{T}) where {T} - return function gradient_function(::Euclidean, peps_vec::Vector{T}) - if !(peps_vec == cache.peps_vec) # update cache if at new point - cost_and_grad!(cache, peps_vec) - end - return cache.grad_vec - end -end - -""" - fixedpoint( - peps₀::InfinitePEPS{T}, - operator::LocalOperator, - alg::PEPSOptimize, - env₀::CTMRGEnv=CTMRGEnv(peps₀, field(T)^20); - ) where {T} - -Optimize a PEPS starting from `peps₀` and `env₀` by minimizing the cost function given -by expectation value of `operator`. All optimization parameters are provided through `alg`. -""" -function fixedpoint( - peps₀::InfinitePEPS{T}, - operator::LocalOperator, - alg::PEPSOptimize, - env₀::CTMRGEnv=CTMRGEnv(peps₀, field(T)^20); -) where {T} - if scalartype(env₀) <: Real - env₀ = complex(env₀) - @warn "the provided real environment was converted to a complex environment since \ - :fixed mode generally produces complex gauges; use :diffgauge mode instead to work \ - with purely real environments" - end - - # construct cost and grad functions - peps₀_vec, from_vec = to_vec(peps₀) - cache = PEPSCostFunctionCache(operator, alg, peps₀_vec, from_vec, deepcopy(env₀)) - cost = cache - grad = gradient_function(cache) - - # optimize - M = Euclidean(length(peps₀_vec)) - retraction_method = if isnothing(alg.symmetrization) - default_retraction_method(M) - else - SymmetrizeExponentialRetraction(alg.symmetrization, from_vec) - end - result = alg.optim_alg( - M, cost, grad, peps₀_vec; alg.optim_kwargs..., retraction_method, return_state=true - ) - - # extract final result - peps_final = from_vec(get_solver_result(result)) - return peps_final, cost.env, cache.cost, result -end - -#= -Evaluating the gradient of the cost function for CTMRG: -- The gradient of the cost function for CTMRG can be computed using automatic differentiation (AD) or explicit evaluation of the geometric sum. -- With AD, the gradient is computed by differentiating the cost function with respect to the PEPS tensors, including computing the environment tensors. -- With explicit evaluation of the geometric sum, the gradient is computed by differentiating the cost function with the environment kept fixed, and then manually adding the gradient contributions from the environments. -=# - -function _rrule( - gradmode::GradMode{:diffgauge}, - config::RuleConfig, - ::typeof(MPSKit.leading_boundary), - envinit, - state, - alg::CTMRGAlgorithm, -) - envs, info = leading_boundary(envinit, state, alg) - - function leading_boundary_diffgauge_pullback((Δenvs′, Δinfo)) - Δenvs = unthunk(Δenvs′) - - # find partial gradients of gauge_fixed single CTMRG iteration - f(A, x) = gauge_fix(x, ctmrg_iteration(A, x, alg)[1])[1] - _, env_vjp = rrule_via_ad(config, f, state, envs) - - # evaluate the geometric sum - ∂f∂A(x)::typeof(state) = env_vjp(x)[2] - ∂f∂x(x)::typeof(envs) = env_vjp(x)[3] - ∂F∂envs = fpgrad(Δenvs, ∂f∂x, ∂f∂A, Δenvs, gradmode) - - return NoTangent(), ZeroTangent(), ∂F∂envs, NoTangent() - end - - return (envs, info), leading_boundary_diffgauge_pullback -end - -# Here f is differentiated from an pre-computed SVD with fixed U, S and V -function _rrule( - gradmode::GradMode{:fixed}, - config::RuleConfig, - ::typeof(MPSKit.leading_boundary), - envinit, - state, - alg::SimultaneousCTMRG, -) - @assert !isnothing(alg.projector_alg.svd_alg.rrule_alg) - envs, = leading_boundary(envinit, state, alg) - envs_conv, info = ctmrg_iteration(state, envs, alg) - envs_fixed, signs = gauge_fix(envs, envs_conv) - - # Fix SVD - Ufix, Vfix = fix_relative_phases(info.U, info.V, signs) - svd_alg_fixed = SVDAdjoint(; - fwd_alg=FixedSVD(Ufix, info.S, Vfix), rrule_alg=alg.projector_alg.svd_alg.rrule_alg - ) - alg_fixed = @set alg.projector_alg.svd_alg = svd_alg_fixed - alg_fixed = @set alg_fixed.projector_alg.trscheme = notrunc() - - function leading_boundary_fixed_pullback((Δenvs′, Δinfo)) - Δenvs = unthunk(Δenvs′) - - f(A, x) = fix_global_phases(x, ctmrg_iteration(A, x, alg_fixed)[1]) - _, env_vjp = rrule_via_ad(config, f, state, envs_fixed) - - # evaluate the geometric sum - ∂f∂A(x)::typeof(state) = env_vjp(x)[2] - ∂f∂x(x)::typeof(envs) = env_vjp(x)[3] - ∂F∂envs = fpgrad(Δenvs, ∂f∂x, ∂f∂A, Δenvs, gradmode) - - return NoTangent(), ZeroTangent(), ∂F∂envs, NoTangent() - end - - return (envs_fixed, info), leading_boundary_fixed_pullback -end - -@doc """ - fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y0, alg) - -Compute the gradient of the cost function for CTMRG by solving the following equation: - -dx = ∑ₙ (∂f∂x)ⁿ ∂f∂A dA = (1 - ∂f∂x)⁻¹ ∂f∂A dA - -where `∂F∂x` is the gradient of the cost function with respect to the PEPS tensors, `∂f∂x` -is the partial gradient of the CTMRG iteration with respect to the environment tensors, -`∂f∂A` is the partial gradient of the CTMRG iteration with respect to the PEPS tensors, and -`y0` is the initial guess for the fixed-point iteration. The function returns the gradient -`dx` of the fixed-point iteration. -""" -fpgrad - -# TODO: can we construct an implementation that does not need to evaluate the vjp -# twice if both ∂f∂A and ∂f∂x are needed? -function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, _, alg::GeomSum) - g = ∂F∂x - dx = ∂f∂A(g) # n = 0 term: ∂F∂x ∂f∂A - ϵ = 2 * alg.tol - for i in 1:(alg.maxiter) - g = ∂f∂x(g) - Σₙ = ∂f∂A(g) - dx += Σₙ - ϵnew = norm(Σₙ) # TODO: normalize this error? - Δϵ = ϵ - ϵnew - alg.verbosity > 1 && - @printf("Gradient iter: %3d ‖Σₙ‖: %.2e Δ‖Σₙ‖: %.2e\n", i, ϵnew, Δϵ) - ϵ = ϵnew - - ϵ < alg.tol && break - if alg.verbosity > 0 && i == alg.maxiter - @warn "gradient fixed-point iteration reached maximal number of iterations at ‖Σₙ‖ = $ϵ" - end - end - return dx -end - -function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y₀, alg::ManualIter) - y = deepcopy(y₀) # Do not mutate y₀ - dx = ∂f∂A(y) - ϵ = 1.0 - for i in 1:(alg.maxiter) - y′ = ∂F∂x + ∂f∂x(y) - - dxnew = ∂f∂A(y′) - ϵnew = norm(dxnew - dx) - Δϵ = ϵ - ϵnew - alg.verbosity > 1 && @printf( - "Gradient iter: %3d ‖Cᵢ₊₁-Cᵢ‖/N: %.2e Δ‖Cᵢ₊₁-Cᵢ‖/N: %.2e\n", i, ϵnew, Δϵ - ) - y = y′ - dx = dxnew - ϵ = ϵnew - - ϵ < alg.tol && break - if alg.verbosity > 0 && i == alg.maxiter - @warn "gradient fixed-point iteration reached maximal number of iterations at ‖Cᵢ₊₁-Cᵢ‖ = $ϵ" - end - end - return dx -end - -function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y₀, alg::LinSolver) - y, info = linsolve(∂f∂x, ∂F∂x, y₀, alg.solver, 1, -1) - if alg.solver.verbosity > 0 && info.converged != 1 - @warn("gradient fixed-point iteration reached maximal number of iterations:", info) - end - - return ∂f∂A(y) -end From b020eb46a4789d5730ee3bdbbd46cdc5b5e02419 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Fri, 17 Jan 2025 18:25:52 +0100 Subject: [PATCH 18/28] Update example tests --- test/pwave.jl | 4 ++-- test/tf_ising.jl | 4 +--- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/test/pwave.jl b/test/pwave.jl index ab7439251..fd33a944d 100644 --- a/test/pwave.jl +++ b/test/pwave.jl @@ -14,9 +14,9 @@ Dbond = 2 ctm_alg = SimultaneousCTMRG() opt_alg = PEPSOptimize(; boundary_alg=ctm_alg, - maxiter=20, + maxiter=10, gradient_alg=LinSolver(; iterscheme=:diffgauge), - stepsize=WolfePowellLinesearch(), + stepsize=ConstantLength(1.5), memory_size=4, ) diff --git a/test/tf_ising.jl b/test/tf_ising.jl index cd1d9301b..aa4181e87 100644 --- a/test/tf_ising.jl +++ b/test/tf_ising.jl @@ -19,9 +19,7 @@ mˣ = 0.91 χbond = 2 χenv = 16 ctm_alg = SimultaneousCTMRG() -opt_alg = PEPSOptimize(; - boundary_alg=ctm_alg, tol=1e-3, stepsize=WolfePowellBinaryLinesearch() -) +opt_alg = PEPSOptimize(; boundary_alg=ctm_alg, tol=1e-3) # initialize states H = transverse_field_ising(InfiniteSquare(); g) From 131e49bc6bc4f1eb20570871b66f4a11fe72215d Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Fri, 17 Jan 2025 18:42:37 +0100 Subject: [PATCH 19/28] Fix formatting and sequential CTMRG --- src/algorithms/ctmrg/sequential.jl | 4 ++-- src/algorithms/optimization/fixed_point_differentiation.jl | 1 - test/heisenberg.jl | 4 +--- 3 files changed, 3 insertions(+), 6 deletions(-) diff --git a/src/algorithms/ctmrg/sequential.jl b/src/algorithms/ctmrg/sequential.jl index e46c54f95..8402aad00 100644 --- a/src/algorithms/ctmrg/sequential.jl +++ b/src/algorithms/ctmrg/sequential.jl @@ -55,8 +55,8 @@ Compute CTMRG projectors in the `:sequential` scheme either for an entire column for a specific `coordinate` (where `dir=WEST` is already implied in the `:sequential` scheme). """ function sequential_projectors( - col::Int, state::InfiniteSquareNetwork, envs::CTMRGEnv, alg::ProjectorAlgorithm -) + col::Int, state::InfiniteSquareNetwork{T}, envs::CTMRGEnv, alg::ProjectorAlgorithm +) where {T} # SVD half-infinite environment column-wise ϵ = Zygote.Buffer(zeros(real(scalartype(T)), size(envs, 2))) S = Zygote.Buffer( diff --git a/src/algorithms/optimization/fixed_point_differentiation.jl b/src/algorithms/optimization/fixed_point_differentiation.jl index 8eefc5650..8d13981ef 100644 --- a/src/algorithms/optimization/fixed_point_differentiation.jl +++ b/src/algorithms/optimization/fixed_point_differentiation.jl @@ -76,7 +76,6 @@ function LinSolver(; return LinSolver{iterscheme}(solver) end - #= Evaluating the gradient of the cost function for CTMRG: - The gradient of the cost function for CTMRG can be computed using automatic differentiation (AD) or explicit evaluation of the geometric sum. diff --git a/test/heisenberg.jl b/test/heisenberg.jl index 55818d65d..a7df0db34 100644 --- a/test/heisenberg.jl +++ b/test/heisenberg.jl @@ -35,9 +35,7 @@ end unitcell = (1, 2) H = heisenberg_XYZ(InfiniteSquare(unitcell...)) psi_init = InfinitePEPS(2, Dbond; unitcell) - env_init, = leading_boundary( - CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg - ) + env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg) # optimize energy and compute correlation lengths peps, env, E, = fixedpoint(psi_init, H, opt_alg, env_init) From 7459e29e11609d103ef31b23e7ef96de61ed792c Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Mon, 20 Jan 2025 17:58:43 +0100 Subject: [PATCH 20/28] Update gradients.jl test (works except for :fixed mode) --- src/algorithms/contractions/ctmrg_contractions.jl | 1 + .../optimization/fixed_point_differentiation.jl | 2 +- test/ctmrg/gradients.jl | 10 ++++++---- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/algorithms/contractions/ctmrg_contractions.jl b/src/algorithms/contractions/ctmrg_contractions.jl index 08a405e97..b054f0892 100644 --- a/src/algorithms/contractions/ctmrg_contractions.jl +++ b/src/algorithms/contractions/ctmrg_contractions.jl @@ -352,6 +352,7 @@ Right projector: ``` """ function contract_projectors(U, S, V, Q, Q_next) + @show (space(U), space(V), space(Q), space(Q_next)) isqS = sdiag_pow(S, -0.5) P_left = Q_next * V' * isqS # use * to respect fermionic case P_right = isqS * U' * Q diff --git a/src/algorithms/optimization/fixed_point_differentiation.jl b/src/algorithms/optimization/fixed_point_differentiation.jl index 8d13981ef..bbb251f94 100644 --- a/src/algorithms/optimization/fixed_point_differentiation.jl +++ b/src/algorithms/optimization/fixed_point_differentiation.jl @@ -221,4 +221,4 @@ function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y₀, alg::LinSolver) end return ∂f∂A(y) -end \ No newline at end of file +end diff --git a/test/ctmrg/gradients.jl b/test/ctmrg/gradients.jl index 0fa7c419c..0feff686a 100644 --- a/test/ctmrg/gradients.jl +++ b/test/ctmrg/gradients.jl @@ -49,16 +49,15 @@ gradmodes = [ models ) Pspace = Pspaces[i] - Vspace = Pspaces[i] + Vspace = Vspaces[i] Espace = Espaces[i] gms = gradmodes[i] calgs = ctmrg_algs[i] - psi_init = InfinitePEPS(Pspace, Vspace, Vspace) @testset "$ctmrg_alg and $gradient_alg" for (ctmrg_alg, gradient_alg) in Iterators.product(calgs, gms) @info "gradient check of $ctmrg_alg and $alg_rrule on $(names[i])" Random.seed!(42039482030) - psi = InfinitePEPS(Pspace, Vspace, Vspace) + psi = InfinitePEPS(Pspace, Vspace) env, = leading_boundary(CTMRGEnv(psi, Espace), psi, ctmrg_alg) psi_vec, from_vec = to_vec(psi) @@ -66,7 +65,10 @@ gradmodes = [ cache = PEPSCostFunctionCache(models[i], opt_alg, psi_vec, from_vec, deepcopy(env)) cost = cache grad = gradient_function(cache) + M = Euclidean(length(psi_vec)) - @test check_gradient(M, cost, grad; N=5, exactness_tol=1e-4, io=stdout) + @test check_gradient( + M, cost, grad; N=10, exactness_tol=gradtol, limits=(-8, -3), io=stdout + ) end end From 50c76f446b1ddc1e30f5556172f788262234c2c3 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Mon, 20 Jan 2025 20:03:00 +0100 Subject: [PATCH 21/28] Replace info NamedTuple by explicit return values --- .../contractions/ctmrg_contractions.jl | 1 - src/algorithms/ctmrg/ctmrg.jl | 9 ++--- src/algorithms/ctmrg/projectors.jl | 4 +-- src/algorithms/ctmrg/sequential.jl | 35 +++++++++++++------ src/algorithms/ctmrg/simultaneous.jl | 23 ++++++------ .../fixed_point_differentiation.jl | 18 +++++----- src/algorithms/optimization/manopt.jl | 4 +-- .../optimization/peps_optimization.jl | 17 ++++----- test/ctmrg/fixed_iterscheme.jl | 8 +++-- test/ctmrg/gradients.jl | 6 ++-- 10 files changed, 68 insertions(+), 57 deletions(-) diff --git a/src/algorithms/contractions/ctmrg_contractions.jl b/src/algorithms/contractions/ctmrg_contractions.jl index b054f0892..08a405e97 100644 --- a/src/algorithms/contractions/ctmrg_contractions.jl +++ b/src/algorithms/contractions/ctmrg_contractions.jl @@ -352,7 +352,6 @@ Right projector: ``` """ function contract_projectors(U, S, V, Q, Q_next) - @show (space(U), space(V), space(Q), space(Q_next)) isqS = sdiag_pow(S, -0.5) P_left = Q_next * V' * isqS # use * to respect fermionic case P_right = isqS * U' * Q diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index ec1cd8e5b..39b25fe57 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -40,12 +40,13 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRGAlgorithm) truncation_error = 0.0 condition_number = 1.0 + U, S, V = _prealloc_svd(envinit.edges) return LoggingExtras.withlevel(; alg.verbosity) do ctmrg_loginit!(log, η, state, envinit) for iter in 1:(alg.maxiter) - env, info = ctmrg_iteration(state, env, alg) # Grow and renormalize in all 4 directions - truncation_error = info.truncation_error - condition_number = info.condition_number + env, truncation_error, condition_number, U, S, V = ctmrg_iteration( + state, env, alg + ) # Grow and renormalize in all 4 directions η, CS, TS = calc_convergence(env, CS, TS) if η ≤ alg.tol && iter ≥ alg.miniter @@ -58,7 +59,7 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRGAlgorithm) ctmrg_logiter!(log, iter, η, state, env) end end - return env, (; truncation_error, condition_number) + return env, truncation_error, condition_number, copy(U), copy(S), copy(V) end end diff --git a/src/algorithms/ctmrg/projectors.jl b/src/algorithms/ctmrg/projectors.jl index 39aa58ee7..f66db6f82 100644 --- a/src/algorithms/ctmrg/projectors.jl +++ b/src/algorithms/ctmrg/projectors.jl @@ -75,7 +75,7 @@ function compute_projector(enlarged_corners, coordinate, alg::HalfInfiniteProjec end end P_left, P_right = contract_projectors(U, S, V, enlarged_corners...) - return (P_left, P_right), (; err, U, S, V) + return (P_left, P_right), err, U, S, V end function compute_projector(enlarged_corners, coordinate, alg::FullInfiniteProjector) halfinf_left = half_infinite_environment(enlarged_corners[1], enlarged_corners[2]) @@ -93,5 +93,5 @@ function compute_projector(enlarged_corners, coordinate, alg::FullInfiniteProjec end end P_left, P_right = contract_projectors(U, S, V, halfinf_left, halfinf_right) - return (P_left, P_right), (; err, U, S, V) + return (P_left, P_right), err, U, S, V end diff --git a/src/algorithms/ctmrg/sequential.jl b/src/algorithms/ctmrg/sequential.jl index 8402aad00..7ba5af201 100644 --- a/src/algorithms/ctmrg/sequential.jl +++ b/src/algorithms/ctmrg/sequential.jl @@ -33,18 +33,26 @@ end function ctmrg_iteration(state, envs::CTMRGEnv, alg::SequentialCTMRG) truncation_error = zero(real(scalartype(state))) condition_number = zero(real(scalartype(state))) - for _ in 1:4 # rotate + U, S, V = _prealloc_svd(envs.edges) + for dir in 1:4 # rotate for col in 1:size(state, 2) # left move column-wise - projectors, info = sequential_projectors(col, state, envs, alg.projector_alg) + projectors, err, cond, U′, S′, V′ = sequential_projectors( + col, state, envs, alg.projector_alg + ) envs = renormalize_sequentially(col, projectors, state, envs) - truncation_error = max(truncation_error, info.truncation_error) - condition_number = max(condition_number, info.condition_number) + truncation_error = max(truncation_error, err) + condition_number = max(condition_number, cond) + for row in 1:size(state, 1) + U[dir, row, col] = U′[row] + S[dir, row, col] = S′[row] + V[dir, row, col] = V′[row] + end end state = rotate_north(state, EAST) envs = rotate_north(envs, EAST) end - return envs, (; truncation_error, condition_number) + return envs, truncation_error, condition_number, U, S, V end """ @@ -62,21 +70,26 @@ function sequential_projectors( S = Zygote.Buffer( zeros(size(envs, 2)), tensormaptype(spacetype(T), 1, 1, real(scalartype(T))) ) + U, S, V = _prealloc_svd(@view(envs.edges[4, :, col])) coordinates = eachcoordinate(envs)[:, col] projectors = dtmap(coordinates) do (r, c) trscheme = truncation_scheme(alg, envs.edges[WEST, _prev(r, size(envs, 2)), c]) - proj, info = sequential_projectors( + proj, err, U′, S′, V′ = sequential_projectors( (WEST, r, c), state, envs, @set(alg.trscheme = trscheme) ) - ϵ[r] = info.err / norm(info.S) - S[r] = info.S + U[r] = U′ + S[r] = S′ + V[r] = V′ + ϵ[r] = err / norm(S′) return proj end + P_left = map(first, projectors) + P_right = map(last, projectors) + S = copy(S) truncation_error = maximum(copy(ϵ)) - condition_number = maximum(_condition_number, copy(S)) - info = (; truncation_error, condition_number) - return (map(first, projectors), map(last, projectors)), info + condition_number = maximum(_condition_number, S) + return (P_left, P_right), truncation_error, condition_number, copy(U), S, copy(V) end function sequential_projectors( coordinate::NTuple{3,Int}, diff --git a/src/algorithms/ctmrg/simultaneous.jl b/src/algorithms/ctmrg/simultaneous.jl index 31f305678..5c0bc58eb 100644 --- a/src/algorithms/ctmrg/simultaneous.jl +++ b/src/algorithms/ctmrg/simultaneous.jl @@ -34,13 +34,15 @@ function ctmrg_iteration(state, envs::CTMRGEnv, alg::SimultaneousCTMRG) enlarged_corners = dtmap(eachcoordinate(state, 1:4)) do idx return TensorMap(EnlargedCorner(state, envs, idx), idx[1]) end # expand environment - projectors, info = simultaneous_projectors(enlarged_corners, envs, alg.projector_alg) # compute projectors on all coordinates + projectors, truncation_error, condition_number, U, S, V = simultaneous_projectors( + enlarged_corners, envs, alg.projector_alg + ) # compute projectors on all coordinates envs′ = renormalize_simultaneously(enlarged_corners, projectors, state, envs) # renormalize enlarged corners - return envs′, info + return envs′, truncation_error, condition_number, U, S, V end # Pre-allocate U, S, and V tensor as Zygote buffers to make it differentiable -function _prealloc_svd(edges::Array{E,N}) where {E,N} +function _prealloc_svd(edges::AbstractArray{E,N}) where {E,N} Sc = scalartype(E) U = Zygote.Buffer(map(e -> TensorMap(zeros, Sc, space(e)), edges)) V = Zygote.Buffer(map(e -> TensorMap(zeros, Sc, domain(e), codomain(e)), edges)) @@ -74,23 +76,22 @@ function simultaneous_projectors( projectors = dtmap(eachcoordinate(envs, 1:4)) do coordinate coordinate′ = _next_coordinate(coordinate, size(envs)[2:3]...) trscheme = truncation_scheme(alg, envs.edges[coordinate[1], coordinate′[2:3]...]) - proj, info = simultaneous_projectors( + proj, err, U′, S′, V′ = simultaneous_projectors( coordinate, enlarged_corners, @set(alg.trscheme = trscheme) ) - U[coordinate...] = info.U - S[coordinate...] = info.S - V[coordinate...] = info.V - ϵ[coordinate...] = info.err / norm(info.S) + U[coordinate...] = U′ + S[coordinate...] = S′ + V[coordinate...] = V′ + ϵ[coordinate...] = err / norm(S′) return proj end P_left = map(first, projectors) P_right = map(last, projectors) S = copy(S) - truncation_error = maximum(copy(ϵ)) # TODO: This makes Zygote error on first execution? + truncation_error = maximum(copy(ϵ)) condition_number = maximum(_condition_number, S) - info = (; truncation_error, condition_number, U=copy(U), S, V=copy(V)) - return (P_left, P_right), info + return (P_left, P_right), truncation_error, condition_number, copy(U), S, copy(V) end function simultaneous_projectors( coordinate, enlarged_corners::Array{E,3}, alg::HalfInfiniteProjector diff --git a/src/algorithms/optimization/fixed_point_differentiation.jl b/src/algorithms/optimization/fixed_point_differentiation.jl index bbb251f94..804101519 100644 --- a/src/algorithms/optimization/fixed_point_differentiation.jl +++ b/src/algorithms/optimization/fixed_point_differentiation.jl @@ -91,9 +91,9 @@ function _rrule( state, alg::CTMRGAlgorithm, ) - envs, info = leading_boundary(envinit, state, alg) + envs, truncation_error, condition_number = leading_boundary(envinit, state, alg) - function leading_boundary_diffgauge_pullback((Δenvs′, Δinfo)) + function leading_boundary_diffgauge_pullback((Δenvs′, Δtrunc_error, Δcond_number)) Δenvs = unthunk(Δenvs′) # find partial gradients of gauge_fixed single CTMRG iteration @@ -108,7 +108,7 @@ function _rrule( return NoTangent(), ZeroTangent(), ∂F∂envs, NoTangent() end - return (envs, info), leading_boundary_diffgauge_pullback + return (envs, truncation_error, condition_number), leading_boundary_diffgauge_pullback end # Here f is differentiated from an pre-computed SVD with fixed U, S and V @@ -122,18 +122,20 @@ function _rrule( ) @assert !isnothing(alg.projector_alg.svd_alg.rrule_alg) envs, = leading_boundary(envinit, state, alg) - envs_conv, info = ctmrg_iteration(state, envs, alg) + envs_conv, truncation_error, condition_number, U, S, V = ctmrg_iteration( + state, envs, alg + ) envs_fixed, signs = gauge_fix(envs, envs_conv) # Fix SVD - Ufix, Vfix = fix_relative_phases(info.U, info.V, signs) + Ufix, Vfix = fix_relative_phases(U, V, signs) svd_alg_fixed = SVDAdjoint(; - fwd_alg=FixedSVD(Ufix, info.S, Vfix), rrule_alg=alg.projector_alg.svd_alg.rrule_alg + fwd_alg=FixedSVD(Ufix, S, Vfix), rrule_alg=alg.projector_alg.svd_alg.rrule_alg ) alg_fixed = @set alg.projector_alg.svd_alg = svd_alg_fixed alg_fixed = @set alg_fixed.projector_alg.trscheme = notrunc() - function leading_boundary_fixed_pullback((Δenvs′, Δinfo)) + function leading_boundary_fixed_pullback((Δenvs′, Δtrunc_error, Δcond_number)) Δenvs = unthunk(Δenvs′) f(A, x) = fix_global_phases(x, ctmrg_iteration(A, x, alg_fixed)[1]) @@ -147,7 +149,7 @@ function _rrule( return NoTangent(), ZeroTangent(), ∂F∂envs, NoTangent() end - return (envs_fixed, info), leading_boundary_fixed_pullback + return (envs_fixed, truncation_error, condition_number), leading_boundary_fixed_pullback end @doc """ diff --git a/src/algorithms/optimization/manopt.jl b/src/algorithms/optimization/manopt.jl index 142f4640c..31f07ab14 100644 --- a/src/algorithms/optimization/manopt.jl +++ b/src/algorithms/optimization/manopt.jl @@ -22,7 +22,7 @@ function (r::RecordTruncationError)( p::AbstractManoptProblem, ::AbstractManoptSolverState, i::Int ) cache = Manopt.get_cost_function(get_objective(p)) - return Manopt.record_or_reset!(r, cache.env_info.truncation_error, i) + return Manopt.record_or_reset!(r, cache.truncation_error, i) end """ @@ -39,7 +39,7 @@ function (r::RecordConditionNumber)( p::AbstractManoptProblem, ::AbstractManoptSolverState, i::Int ) cache = Manopt.get_cost_function(get_objective(p)) - return Manopt.record_or_reset!(r, cache.env_info.condition_number, i) + return Manopt.record_or_reset!(r, cache.condition_number, i) end """ diff --git a/src/algorithms/optimization/peps_optimization.jl b/src/algorithms/optimization/peps_optimization.jl index 93bdd3b85..6f27fa79c 100644 --- a/src/algorithms/optimization/peps_optimization.jl +++ b/src/algorithms/optimization/peps_optimization.jl @@ -100,7 +100,8 @@ mutable struct PEPSCostFunctionCache{T} peps_vec::Vector{T} grad_vec::Vector{T} cost::Float64 - env_info::NamedTuple + truncation_error::Float64 + condition_number::Float64 end """ @@ -115,14 +116,7 @@ function PEPSCostFunctionCache( operator::LocalOperator, alg::PEPSOptimize, peps_vec::Vector, from_vec, env::CTMRGEnv ) return PEPSCostFunctionCache( - operator, - alg, - env, - from_vec, - similar(peps_vec), - similar(peps_vec), - 0.0, - (; truncation_error=0.0, condition_number=1.0), + operator, alg, env, from_vec, similar(peps_vec), similar(peps_vec), 0.0, 0.0, 1.0 ) end @@ -140,7 +134,7 @@ function cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) wh # compute cost and gradient cost, grads = withgradient(peps) do ψ - env, info = hook_pullback( + env, truncation_error, condition_number = hook_pullback( leading_boundary, env₀, ψ, @@ -150,7 +144,8 @@ function cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) wh cost = expectation_value(ψ, cache.operator, env) ignore_derivatives() do update!(cache.env, env) # update environment in-place - cache.env_info = info # update environment information (truncation error, ...) + cache.truncation_error = truncation_error # update environment information + cache.condition_number = condition_number isapprox(imag(cost), 0; atol=sqrt(eps(real(cost)))) || @warn "Expectation value is not real: $cost." end diff --git a/test/ctmrg/fixed_iterscheme.jl b/test/ctmrg/fixed_iterscheme.jl index 39eccf5c0..0feab8778 100644 --- a/test/ctmrg/fixed_iterscheme.jl +++ b/test/ctmrg/fixed_iterscheme.jl @@ -33,13 +33,15 @@ atol = 1e-5 env_conv1, = leading_boundary(CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg) # do extra iteration to get SVD - env_conv2, info = ctmrg_iteration(psi, env_conv1, ctm_alg) + env_conv2, truncation_error, condition_number, U, S, V = ctmrg_iteration( + psi, env_conv1, ctm_alg + ) env_fix, signs = gauge_fix(env_conv1, env_conv2) @test calc_elementwise_convergence(env_conv1, env_fix) ≈ 0 atol = atol # fix gauge of SVD - U_fix, V_fix = fix_relative_phases(info.U, info.V, signs) - svd_alg_fix = SVDAdjoint(; fwd_alg=FixedSVD(U_fix, info.S, V_fix)) + U_fix, V_fix = fix_relative_phases(U, V, signs) + svd_alg_fix = SVDAdjoint(; fwd_alg=FixedSVD(U_fix, S, V_fix)) ctm_alg_fix = SimultaneousCTMRG(; projector_alg, svd_alg=svd_alg_fix, trscheme=notrunc() ) diff --git a/test/ctmrg/gradients.jl b/test/ctmrg/gradients.jl index 0feff686a..c5844d937 100644 --- a/test/ctmrg/gradients.jl +++ b/test/ctmrg/gradients.jl @@ -55,7 +55,7 @@ gradmodes = [ calgs = ctmrg_algs[i] @testset "$ctmrg_alg and $gradient_alg" for (ctmrg_alg, gradient_alg) in Iterators.product(calgs, gms) - @info "gradient check of $ctmrg_alg and $alg_rrule on $(names[i])" + @info "gradient check of $ctmrg_alg and $gradient_alg on $(names[i])" Random.seed!(42039482030) psi = InfinitePEPS(Pspace, Vspace) env, = leading_boundary(CTMRGEnv(psi, Espace), psi, ctmrg_alg) @@ -67,8 +67,6 @@ gradmodes = [ grad = gradient_function(cache) M = Euclidean(length(psi_vec)) - @test check_gradient( - M, cost, grad; N=10, exactness_tol=gradtol, limits=(-8, -3), io=stdout - ) + @test check_gradient(M, cost, grad; N=10, exactness_tol=gradtol, limits=(-8, -3)) end end From a0178b54422e1c3135a7796c8134610c17e2954d Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Mon, 20 Jan 2025 20:28:33 +0100 Subject: [PATCH 22/28] Fix :sequential for rectangular unit cells, fix tf_ising test --- src/algorithms/ctmrg/sequential.jl | 7 ++++--- test/ctmrg/gaugefix.jl | 3 ++- test/tf_ising.jl | 4 ++-- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/src/algorithms/ctmrg/sequential.jl b/src/algorithms/ctmrg/sequential.jl index 7ba5af201..807f4c0e4 100644 --- a/src/algorithms/ctmrg/sequential.jl +++ b/src/algorithms/ctmrg/sequential.jl @@ -43,9 +43,10 @@ function ctmrg_iteration(state, envs::CTMRGEnv, alg::SequentialCTMRG) truncation_error = max(truncation_error, err) condition_number = max(condition_number, cond) for row in 1:size(state, 1) - U[dir, row, col] = U′[row] - S[dir, row, col] = S′[row] - V[dir, row, col] = V′[row] + rc_idx = isodd(dir) ? (row, col) : (col, row) # relevant for rectangular unit cells + U[dir, rc_idx...] = U′[row] + S[dir, rc_idx...] = S′[row] + V[dir, rc_idx...] = V′[row] end end state = rotate_north(state, EAST) diff --git a/test/ctmrg/gaugefix.jl b/test/ctmrg/gaugefix.jl index 522314983..1b053ff0c 100644 --- a/test/ctmrg/gaugefix.jl +++ b/test/ctmrg/gaugefix.jl @@ -7,7 +7,8 @@ using PEPSKit: ctmrg_iteration, gauge_fix, calc_elementwise_convergence spacetypes = [ComplexSpace, Z2Space] scalartypes = [Float64, ComplexF64] -unitcells = [(1, 1), (2, 2), (3, 2)] +# unitcells = [(1, 1), (2, 2), (3, 2)] +unitcells = [(3, 2)] ctmrg_algs = [SequentialCTMRG, SimultaneousCTMRG] projector_algs = [HalfInfiniteProjector, FullInfiniteProjector] tol = 1e-6 # large tol due to χ=6 diff --git a/test/tf_ising.jl b/test/tf_ising.jl index aa4181e87..51276f85d 100644 --- a/test/tf_ising.jl +++ b/test/tf_ising.jl @@ -34,7 +34,7 @@ peps, env, E, = fixedpoint(psi_init, H, opt_alg, env_init) # compute magnetization σx = TensorMap(scalartype(psi_init)[0 1; 1 0], ℂ^2, ℂ^2) M = LocalOperator(H.lattice, (CartesianIndex(1, 1),) => σx) -magn = expectation_value(result.peps, M, result.env) +magn = expectation_value(peps, M, env) @test E ≈ e atol = 1e-2 @test imag(magn) ≈ 0 atol = 1e-6 @@ -43,6 +43,6 @@ magn = expectation_value(result.peps, M, result.env) # find fixedpoint in polarized phase and compute correlations lengths H_polar = transverse_field_ising(InfiniteSquare(); g=4.5) peps_polar, env_polar, E_polar, = fixedpoint(psi_init, H_polar, opt_alg, env_init) -ξ_h_polar, ξ_v_polar, = correlation_length(result_polar.peps, result_polar.env) +ξ_h_polar, ξ_v_polar, = correlation_length(peps_polar, env_polar) @test ξ_h_polar < ξ_h @test ξ_v_polar < ξ_v From 279b5485e3037efe720ac096bdbf8ab2087f804c Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Tue, 21 Jan 2025 13:32:58 +0100 Subject: [PATCH 23/28] Fix fixed_iterscheme test --- test/ctmrg/fixed_iterscheme.jl | 35 ++++++++++++++++++---------------- 1 file changed, 19 insertions(+), 16 deletions(-) diff --git a/test/ctmrg/fixed_iterscheme.jl b/test/ctmrg/fixed_iterscheme.jl index 0feab8778..1954e4a8c 100644 --- a/test/ctmrg/fixed_iterscheme.jl +++ b/test/ctmrg/fixed_iterscheme.jl @@ -66,21 +66,25 @@ end env_conv1, = leading_boundary(env_init, psi, ctm_alg_iter) # do extra iteration to get SVD - env_conv2_iter, info_iter = ctmrg_iteration(psi, env_conv1, ctm_alg_iter) + env_conv2_iter, _, _, U_iter, S_iter, V_iter = ctmrg_iteration( + psi, env_conv1, ctm_alg_iter + ) env_fix_iter, signs_iter = gauge_fix(env_conv1, env_conv2_iter) @test calc_elementwise_convergence(env_conv1, env_fix_iter) ≈ 0 atol = atol - env_conv2_full, info_full = ctmrg_iteration(psi, env_conv1, ctm_alg_full) + env_conv2_full, _, _, U_full, S_full, V_full = ctmrg_iteration( + psi, env_conv1, ctm_alg_full + ) env_fix_full, signs_full = gauge_fix(env_conv1, env_conv2_full) @test calc_elementwise_convergence(env_conv1, env_fix_full) ≈ 0 atol = atol # fix gauge of SVD - U_fix_iter, V_fix_iter = fix_relative_phases(info_iter.U, info_iter.V, signs_iter) - svd_alg_fix_iter = SVDAdjoint(; fwd_alg=FixedSVD(U_fix_iter, info_iter.S, V_fix_iter)) + U_fix_iter, V_fix_iter = fix_relative_phases(U_iter, V_iter, signs_iter) + svd_alg_fix_iter = SVDAdjoint(; fwd_alg=FixedSVD(U_fix_iter, S_iter, V_fix_iter)) ctm_alg_fix_iter = SimultaneousCTMRG(; svd_alg=svd_alg_fix_iter, trscheme=notrunc()) - U_fix_full, V_fix_full = fix_relative_phases(info_full.U, info_full.V, signs_full) - svd_alg_fix_full = SVDAdjoint(; fwd_alg=FixedSVD(U_fix_full, info_full.S, V_fix_full)) + U_fix_full, V_fix_full = fix_relative_phases(U_full, V_full, signs_full) + svd_alg_fix_full = SVDAdjoint(; fwd_alg=FixedSVD(U_fix_full, S_full, V_fix_full)) ctm_alg_fix_full = SimultaneousCTMRG(; svd_alg=svd_alg_fix_full, trscheme=notrunc()) # do iteration with FixedSVD @@ -93,24 +97,23 @@ end @test calc_elementwise_convergence(env_conv1, env_fixedsvd_full) ≈ 0 atol = atol # check matching decompositions - decomposition_check = all( - zip(info_iter.U, info_iter.S, info_iter.V, info_full.U, info_full.S, info_full.V), - ) do (U_iter, S_iter, V_iter, U_full, S_full, V_full) - diff = U_iter * S_iter * V_iter - U_full * S_full * V_full - all(x -> isapprox(abs(x), 0; atol), diff.data) - end + decomposition_check = + all(zip(U_iter, S_iter, V_iter, U_full, S_full, V_full)) do (Ui, Si, Vi, Uf, Sf, Vf) + diff = Ui * Si * Vi - Uf * Sf * Vf + all(x -> isapprox(abs(x), 0; atol), diff.data) + end @test decomposition_check # check matching singular values - svalues_check = all(zip(info_iter.S, info_full.S)) do (S_iter, S_full) - diff = S_iter - S_full + svalues_check = all(zip(S_iter, S_full)) do (Si, Sf) + diff = Si - Sf all(x -> isapprox(abs(x), 0; atol), diff.data) end @test svalues_check # check normalization of U's and V's - Us = [info_iter.U, U_fix_iter, info_full.U, U_fix_full] - Vs = [info_iter.V, V_fix_iter, info_full.V, V_fix_full] + Us = [U_iter, U_fix_iter, U_full, U_fix_full] + Vs = [V_iter, V_fix_iter, V_full, V_fix_full] for (U, V) in zip(Us, Vs) U_check = all(U) do u uu = u' * u From 9c7366dbf8a8332c8bbd639ded976f930141f14e Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Tue, 21 Jan 2025 13:48:29 +0100 Subject: [PATCH 24/28] Readd costfun and rename to cost_function --- examples/hubbard_su.jl | 2 +- src/PEPSKit.jl | 2 +- src/algorithms/optimization/peps_optimization.jl | 12 +++++------- src/algorithms/toolbox.jl | 15 +++++++++++++++ test/ctmrg/flavors.jl | 4 ++-- test/heisenberg.jl | 2 +- 6 files changed, 25 insertions(+), 12 deletions(-) diff --git a/examples/hubbard_su.jl b/examples/hubbard_su.jl index d65512f59..3e0f65426 100644 --- a/examples/hubbard_su.jl +++ b/examples/hubbard_su.jl @@ -53,7 +53,7 @@ Es_exact = Dict(0 => -1.62, 2 => -0.176, 4 => 0.8603, 6 => -0.6567, 8 => -0.5243 E_exact = Es_exact[U] - U / 2 # measure energy -E = costfun(peps, envs, ham) / (N1 * N2) +E = cost_function(peps, envs, ham) / (N1 * N2) @info "Energy = $E" @info "Benchmark energy = $E_exact" @test isapprox(E, E_exact; atol=5e-2) diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 9c8945815..fe84605a9 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -256,7 +256,7 @@ export SVDAdjoint, IterSVD, NonTruncSVDAdjoint export CTMRGEnv, SequentialCTMRG, SimultaneousCTMRG export FixedSpaceTruncation, HalfInfiniteProjector, FullInfiniteProjector export LocalOperator -export expectation_value, costfun, product_peps, correlation_length +export expectation_value, cost_function, product_peps, correlation_length export leading_boundary export PEPSOptimize, GeomSum, ManualIter, LinSolver export RecordTruncationError, RecordConditionNumber, RecordUnitCellGradientNorm diff --git a/src/algorithms/optimization/peps_optimization.jl b/src/algorithms/optimization/peps_optimization.jl index 6f27fa79c..74ffbfef6 100644 --- a/src/algorithms/optimization/peps_optimization.jl +++ b/src/algorithms/optimization/peps_optimization.jl @@ -141,15 +141,13 @@ function cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) wh cache.alg.boundary_alg; alg_rrule=cache.alg.gradient_alg, ) - cost = expectation_value(ψ, cache.operator, env) - ignore_derivatives() do - update!(cache.env, env) # update environment in-place - cache.truncation_error = truncation_error # update environment information + cost = cost_function(ψ, cache.operator, env) + ignore_derivatives() do # update cache + update!(cache.env, env) + cache.truncation_error = truncation_error cache.condition_number = condition_number - isapprox(imag(cost), 0; atol=sqrt(eps(real(cost)))) || - @warn "Expectation value is not real: $cost." end - return real(cost) + return cost end grad = only(grads) # `withgradient` returns tuple of gradients `grads` diff --git a/src/algorithms/toolbox.jl b/src/algorithms/toolbox.jl index 7dffed197..2a5ed72cd 100644 --- a/src/algorithms/toolbox.jl +++ b/src/algorithms/toolbox.jl @@ -37,6 +37,21 @@ function MPSKit.expectation_value( return expectation_value(pf, CartesianIndex(op[1]) => op[2], envs) end +""" + cost_function(peps::InfinitePEPS, envs::CTMRGEnv, O::LocalOperator) + +Real part of expectation value of `O`. Prints a warning if the expectation value +yields a finite imaginary part (up to a tolerance). +""" +function cost_function(peps::InfinitePEPS, envs::CTMRGEnv, O::LocalOperator) + E = MPSKit.expectation_value(peps, O, envs) + ignore_derivatives() do + isapprox(imag(E), 0; atol=sqrt(eps(real(E)))) || + @warn "Expectation value is not real: $E." + end + return real(E) +end + function LinearAlgebra.norm(peps::InfinitePEPS, env::CTMRGEnv) total = one(scalartype(peps)) diff --git a/test/ctmrg/flavors.jl b/test/ctmrg/flavors.jl index 1c84f7e94..b3d935a54 100644 --- a/test/ctmrg/flavors.jl +++ b/test/ctmrg/flavors.jl @@ -50,8 +50,8 @@ projector_algs = [HalfInfiniteProjector, FullInfiniteProjector] # compare Heisenberg energies H = heisenberg_XYZ(InfiniteSquare(unitcell...)) - E_sequential = costfun(psi, env_sequential, H) - E_simultaneous = costfun(psi, env_simultaneous, H) + E_sequential = cost_function(psi, env_sequential, H) + E_simultaneous = cost_function(psi, env_simultaneous, H) @test E_sequential ≈ E_simultaneous rtol = 1e-3 end diff --git a/test/heisenberg.jl b/test/heisenberg.jl index a7df0db34..3fb44ca4f 100644 --- a/test/heisenberg.jl +++ b/test/heisenberg.jl @@ -81,7 +81,7 @@ end envs, = leading_boundary(envs₀, peps, SimultaneousCTMRG()) # measure physical quantities - e_site = costfun(peps, envs, ham) / (N1 * N2) + e_site = cost_function(peps, envs, ham) / (N1 * N2) @info "Simple update energy = $e_site" # benchmark data from Phys. Rev. B 94, 035133 (2016) @test isapprox(e_site, -0.6594; atol=1e-3) From acbcece3d79ce26f81aa0338016c43745473f256 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Tue, 21 Jan 2025 13:55:08 +0100 Subject: [PATCH 25/28] Fix cost_function argument order --- src/algorithms/optimization/peps_optimization.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/algorithms/optimization/peps_optimization.jl b/src/algorithms/optimization/peps_optimization.jl index 74ffbfef6..d76bd71f5 100644 --- a/src/algorithms/optimization/peps_optimization.jl +++ b/src/algorithms/optimization/peps_optimization.jl @@ -141,7 +141,7 @@ function cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) wh cache.alg.boundary_alg; alg_rrule=cache.alg.gradient_alg, ) - cost = cost_function(ψ, cache.operator, env) + cost = cost_function(ψ, env, cache.operator) ignore_derivatives() do # update cache update!(cache.env, env) cache.truncation_error = truncation_error From c011d3d6b295fb13c157be448fdf6d817aa0175c Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Tue, 21 Jan 2025 15:24:38 +0100 Subject: [PATCH 26/28] Fix partition_function test --- test/ctmrg/partition_function.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/ctmrg/partition_function.jl b/test/ctmrg/partition_function.jl index 8995d6ab8..13a20ffb1 100644 --- a/test/ctmrg/partition_function.jl +++ b/test/ctmrg/partition_function.jl @@ -98,7 +98,7 @@ projector_algs = [HalfInfiniteProjector, FullInfiniteProjector] ctm_styles, projector_algs ) ctm_alg = ctm_style(; maxiter=150, projector_alg) - env = leading_boundary(env0, Z, ctm_alg) + env, = leading_boundary(env0, Z, ctm_alg) # check observables λ = PEPSKit.value(Z, env) From b57feb7ab789df90a016083cb30602025cbfaffd Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Tue, 21 Jan 2025 15:50:22 +0100 Subject: [PATCH 27/28] Correct wrong fixedpoint return values --- README.md | 4 ++-- docs/src/index.md | 4 ++-- examples/heisenberg.jl | 4 ++-- test/heisenberg.jl | 6 +++--- test/j1j2_model.jl | 6 +++--- 5 files changed, 12 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 8dd2d139f..cdc22898b 100644 --- a/README.md +++ b/README.md @@ -55,8 +55,8 @@ opt_alg = PEPSOptimize(; # ground state search state = InfinitePEPS(2, D) ctm, = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg) -result = fixedpoint(state, H, opt_alg, ctm) +peps, env, E, = fixedpoint(state, H, opt_alg, ctm) -@show result.E # -0.6625... +@show E # -0.6625... ``` diff --git a/docs/src/index.md b/docs/src/index.md index 617753100..a75efb54a 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -37,7 +37,7 @@ opt_alg = PEPSOptimize(; # ground state search state = InfinitePEPS(2, D) ctm, = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg) -result = fixedpoint(state, H, opt_alg, ctm) +peps, env, E, = fixedpoint(state, H, opt_alg, ctm) -@show result.E # -0.6625... +@show E # -0.6625... ``` diff --git a/examples/heisenberg.jl b/examples/heisenberg.jl index 651f46aa1..14a63a353 100644 --- a/examples/heisenberg.jl +++ b/examples/heisenberg.jl @@ -27,5 +27,5 @@ opt_alg = PEPSOptimize(; # Of course there is a noticable bias for small χbond and χenv. ψ₀ = InfinitePEPS(2, χbond) env₀, = leading_boundary(CTMRGEnv(ψ₀, ℂ^χenv), ψ₀, ctm_alg) -result = fixedpoint(ψ₀, H, opt_alg, env₀) -@show result.E +peps, env, E, = fixedpoint(ψ₀, H, opt_alg, env₀) +@show E diff --git a/test/heisenberg.jl b/test/heisenberg.jl index 3fb44ca4f..ee6b2a36d 100644 --- a/test/heisenberg.jl +++ b/test/heisenberg.jl @@ -89,9 +89,9 @@ end # continue with auto differentiation svd_alg_gmres = SVDAdjoint(; rrule_alg=GMRES(; tol=1e-8)) opt_alg_gmres = @set opt_alg.boundary_alg.projector_alg.svd_alg = svd_alg_gmres - result = fixedpoint(peps, ham, opt_alg_gmres, envs) # sensitivity warnings and degeneracies due to SU(2)? - ξ_h, ξ_v, = correlation_length(result.peps, result.env) - e_site2 = result.E / (N1 * N2) + peps_final, env_final, E_final, = fixedpoint(peps, ham, opt_alg_gmres, envs) # sensitivity warnings and degeneracies due to SU(2)? + ξ_h, ξ_v, = correlation_length(peps_final, env_final) + e_site2 = E_final / (N1 * N2) @info "Auto diff energy = $e_site2" @test e_site2 ≈ E_ref atol = 1e-2 @test all(@. ξ_h > 0 && ξ_v > 0) diff --git a/test/j1j2_model.jl b/test/j1j2_model.jl index f1c405ba2..3e73e16a4 100644 --- a/test/j1j2_model.jl +++ b/test/j1j2_model.jl @@ -23,12 +23,12 @@ psi_init = symmetrize!(psi_init, RotateReflect()) env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg); # find fixedpoint -result = fixedpoint(psi_init, H, opt_alg, env_init) -ξ_h, ξ_v, = correlation_length(result.peps, result.env) +peps, env, E, = fixedpoint(psi_init, H, opt_alg, env_init) +ξ_h, ξ_v, = correlation_length(peps, env) # compare against Juraj Hasik's data: # https://github.com/jurajHasik/j1j2_ipeps_states/blob/main/single-site_pg-C4v-A1/j20.25/state_1s_A1_j20.25_D2_chi_opt48.dat ξ_ref = -1 / log(0.2723596743547324) -@test result.E ≈ -0.5618837021945925 atol = 1e-3 +@test E ≈ -0.5618837021945925 atol = 1e-3 @test all(@. isapprox(ξ_h, ξ_ref; atol=1e-1) && isapprox(ξ_v, ξ_ref; atol=1e-1)) @test ξ_h ≈ ξ_v atol = 1e-6 # Test symmetrization of optimized PEPS and environment From d0df0c1bdf9635995c6899408f804be99e17a9d3 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Mon, 3 Feb 2025 17:25:16 +0100 Subject: [PATCH 28/28] Fix jacobian_real_linear.jl test --- test/ctmrg/jacobian_real_linear.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/ctmrg/jacobian_real_linear.jl b/test/ctmrg/jacobian_real_linear.jl index d91c45294..6dc1aa16e 100644 --- a/test/ctmrg/jacobian_real_linear.jl +++ b/test/ctmrg/jacobian_real_linear.jl @@ -19,15 +19,15 @@ Dbond, χenv = 2, 16 @testset "$iterscheme and $ctm_alg" for (iterscheme, ctm_alg) in algs Random.seed!(123521938519) state = InfinitePEPS(2, Dbond) - envs = leading_boundary(CTMRGEnv(state, ComplexSpace(χenv)), state, ctm_alg) + envs, = leading_boundary(CTMRGEnv(state, ComplexSpace(χenv)), state, ctm_alg) # follow code of _rrule if iterscheme == :fixed - envs_conv, info = ctmrg_iteration(state, envs, ctm_alg) + envs_conv, _, _, U, S, V = ctmrg_iteration(state, envs, ctm_alg) envs_fixed, signs = gauge_fix(envs, envs_conv) - U_fixed, V_fixed = fix_relative_phases(info.U, info.V, signs) + U_fixed, V_fixed = fix_relative_phases(U, V, signs) svd_alg_fixed = SVDAdjoint(; - fwd_alg=PEPSKit.FixedSVD(U_fixed, info.S, V_fixed), + fwd_alg=PEPSKit.FixedSVD(U_fixed, S, V_fixed), rrule_alg=ctm_alg.projector_alg.svd_alg.rrule_alg, ) alg_fixed = @set ctm_alg.projector_alg.svd_alg = svd_alg_fixed