diff --git a/Project.toml b/Project.toml index c4abe2b8..b5be59ea 100644 --- a/Project.toml +++ b/Project.toml @@ -9,6 +9,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36" MPSKit = "bb1c41ca-d63c-52ed-829e-0820dda26502" OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index a30251c5..1a3c941b 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -7,6 +7,8 @@ using Accessors using VectorInterface using TensorKit, KrylovKit, MPSKit, OptimKit, TensorOperations using ChainRulesCore, Zygote +using LoggingExtras +using MPSKit: loginit!, logiter!, logfinish!, logcancel! include("utility/util.jl") include("utility/svd.jl") diff --git a/src/algorithms/ctmrg.jl b/src/algorithms/ctmrg.jl index 65b10029..deb46990 100644 --- a/src/algorithms/ctmrg.jl +++ b/src/algorithms/ctmrg.jl @@ -35,9 +35,10 @@ Algorithm struct that represents the CTMRG algorithm for contracting infinite PE Each CTMRG run is converged up to `tol` where the singular value convergence of the corners as well as the norm is checked. The maximal and minimal number of CTMRG iterations is set with `maxiter` and `miniter`. Different levels of output information are printed -depending on `verbosity` (0, 1 or 2). +depending on `verbosity`, where `0` suppresses all output, `1` only prints warnings, `2` +gives information at the start and end, and `3` prints information every iteration. -The projectors are computed from `svd_alg` SVDs where the truncation scheme is set via +The projectors are computed from `svd_alg` SVDs where the truncation scheme is set via `trscheme`. In general, two different schemes can be selected with `ctmrgscheme` which determine how @@ -76,75 +77,256 @@ Per default, a random initial environment is used. function MPSKit.leading_boundary(state, alg::CTMRG) return MPSKit.leading_boundary(CTMRGEnv(state, oneunit(spacetype(state))), state, alg) end -function MPSKit.leading_boundary(envinit, state, alg::CTMRG{S}) where {S} - normold = 1.0 - CSold = map(x -> tsvd(x; alg=TensorKit.SVD())[2], envinit.corners) - TSold = map(x -> tsvd(x; alg=TensorKit.SVD())[2], envinit.edges) - ϵold = 1.0 +function MPSKit.leading_boundary(envinit, state, alg::CTMRG) + CS = map(x -> tsvd(x; alg=TensorKit.SVD())[2], envinit.corners) + TS = map(x -> tsvd(x; alg=TensorKit.SVD())[2], envinit.edges) + + η = one(real(scalartype(state))) + N = norm(state, envinit) env = deepcopy(envinit) + log = ignore_derivatives(() -> MPSKit.IterLog("CTMRG")) + + return LoggingExtras.withlevel(; alg.verbosity) do + ctmrg_loginit!(log, η, N) + local iter + for outer iter in 1:(alg.maxiter) + env, = ctmrg_iter(state, env, alg) # Grow and renormalize in all 4 directions + η, CS, TS = calc_convergence(env, CS, TS) + N = norm(state, env) + ctmrg_logiter!(log, iter, η, N) + + (iter > alg.miniter && η <= alg.tol) && break + end - for i in 1:(alg.maxiter) - env, info = ctmrg_iter(state, env, alg) # Grow and renormalize in all 4 directions - - conv_condition, normold, CSold, TSold, ϵ = ignore_derivatives() do - # Compute convergence criteria and take max (TODO: How should we handle logging all of this?) - Δϵ = abs((ϵold - info.ϵ) / ϵold) - normnew = norm(state, env) - Δnorm = abs(normold - normnew) / abs(normold) - CSnew = map(c -> tsvd(c; alg=TensorKit.SVD())[2], env.corners) - ΔCS = maximum(zip(CSold, CSnew)) do (c_old, c_new) - # only compute the difference on the smallest part of the spaces - smallest = infimum(MPSKit._firstspace(c_old), MPSKit._firstspace(c_new)) - e_old = isometry(MPSKit._firstspace(c_old), smallest) - e_new = isometry(MPSKit._firstspace(c_new), smallest) - return norm(e_new' * c_new * e_new - e_old' * c_old * e_old) - end - TSnew = map(t -> tsvd(t; alg=TensorKit.SVD())[2], env.edges) - ΔTS = maximum(zip(TSold, TSnew)) do (t_old, t_new) - MPSKit._firstspace(t_old) == MPSKit._firstspace(t_new) || - return scalartype(t_old)(Inf) - # TODO: implement when spaces aren't the same - return norm(t_new - t_old) - end + # Do one final iteration that does not change the spaces + alg_fixed = @set alg.projector_alg.trscheme = FixedSpaceTruncation() + env′, = ctmrg_iter(state, env, alg_fixed) + envfix = gauge_fix(env, env′) - conv_condition = max(Δnorm, ΔCS, ΔTS) < alg.tol && i > alg.miniter - - if alg.verbosity > 1 || (alg.verbosity == 1 && (i == 1 || conv_condition)) - @printf( - "CTMRG iter: %3d norm: %.2e Δnorm: %.2e ΔCS: %.2e ΔTS: %.2e ϵ: %.2e Δϵ: %.2e\n", - i, - abs(normnew), - Δnorm, - ΔCS, - ΔTS, - info.ϵ, - Δϵ - ) - end - alg.verbosity > 0 && - i == alg.maxiter && - @warn( - "CTMRG reached maximal number of iterations at (Δnorm=$Δnorm, ΔCS=$ΔCS, ΔTS=$ΔTS)" - ) - return conv_condition, normnew, CSnew, TSnew, info.ϵ + η = calc_elementwise_convergence(envfix, env; atol=alg.tol^(1 / 2)) + N = norm(state, envfix) + + if η < alg.tol^(1 / 2) + ctmrg_logfinish!(log, iter, η, N) + else + ctmrg_logcancel!(log, iter, η, N) end - conv_condition && break # Converge if maximal Δ falls below tolerance + return envfix end +end - # Do one final iteration that does not change the spaces - alg_fixed = CTMRG(; - verbosity=alg.verbosity, - svd_alg=alg.projector_alg.svd_alg, - trscheme=FixedSpaceTruncation(), - ctmrgscheme=S, - ) - env′, = ctmrg_iter(state, env, alg_fixed) - envfix, = gauge_fix(env, env′) - check_elementwise_convergence(env, envfix; atol=alg.tol^(1 / 2)) || - @warn "CTMRG did not converge elementwise." - return envfix +ctmrg_loginit!(log, η, N) = @infov 2 loginit!(log, η, N) +ctmrg_logiter!(log, iter, η, N) = @infov 3 logiter!(log, iter, η, N) +ctmrg_logfinish!(log, iter, η, N) = @infov 2 logfinish!(log, iter, η, N) +ctmrg_logcancel!(log, iter, η, N) = @warnv 1 logcancel!(log, iter, η, N) + +@non_differentiable ctmrg_loginit!(args...) +@non_differentiable ctmrg_logiter!(args...) +@non_differentiable ctmrg_logfinish!(args...) +@non_differentiable ctmrg_logcancel!(args...) + +""" + gauge_fix(envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C,T}) where {C,T} + +Fix the gauge of `envfinal` based on the previous environment `envprev`. +This assumes that the `envfinal` is the result of one CTMRG iteration on `envprev`. +Given that the CTMRG run is converged, the returned environment will be +element-wise converged to `envprev`. +""" +function gauge_fix(envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C,T}) where {C,T} + # Check if spaces in envprev and envfinal are the same + same_spaces = map(Iterators.product(axes(envfinal.edges)...)) do (dir, r, c) + space(envfinal.edges[dir, r, c]) == space(envprev.edges[dir, r, c]) && + space(envfinal.corners[dir, r, c]) == space(envprev.corners[dir, r, c]) + end + @assert all(same_spaces) "Spaces of envprev and envfinal are not the same" + + # Try the "general" algorithm from https://arxiv.org/abs/2311.11894 + signs = map(Iterators.product(axes(envfinal.edges)...)) do (dir, r, c) + # Gather edge tensors and pretend they're InfiniteMPSs + if dir == NORTH + Tsprev = circshift(envprev.edges[dir, r, :], 1 - c) + Tsfinal = circshift(envfinal.edges[dir, r, :], 1 - c) + elseif dir == EAST + Tsprev = circshift(envprev.edges[dir, :, c], 1 - r) + Tsfinal = circshift(envfinal.edges[dir, :, c], 1 - r) + elseif dir == SOUTH + Tsprev = circshift(reverse(envprev.edges[dir, r, :]), c) + Tsfinal = circshift(reverse(envfinal.edges[dir, r, :]), c) + elseif dir == WEST + Tsprev = circshift(reverse(envprev.edges[dir, :, c]), r) + Tsfinal = circshift(reverse(envfinal.edges[dir, :, c]), r) + end + + # Random MPS of same bond dimension + M = map(Tsfinal) do t + TensorMap(randn, scalartype(t), codomain(t) ← domain(t)) + end + + # Find right fixed points of mixed transfer matrices + ρinit = TensorMap( + randn, + scalartype(T), + MPSKit._lastspace(Tsfinal[end])' ← MPSKit._lastspace(M[end])', + ) + ρprev = transfermatrix_fixedpoint(Tsprev, M, ρinit) + ρfinal = transfermatrix_fixedpoint(Tsfinal, M, ρinit) + + # Decompose and multiply + Qprev, = leftorth(ρprev) + Qfinal, = leftorth(ρfinal) + + return Qprev * Qfinal' + end + + cornersfix, edgesfix = fix_relative_phases(envfinal, signs) + + # Fix global phase + cornersgfix = map(envprev.corners, cornersfix) do Cprev, Cfix + return dot(Cfix, Cprev) * Cfix + end + edgesgfix = map(envprev.edges, edgesfix) do Tprev, Tfix + return dot(Tfix, Tprev) * Tfix + end + return CTMRGEnv(cornersgfix, edgesgfix) end +# this is a bit of a hack to get the fixed point of the mixed transfer matrix +# because MPSKit is not compatible with AD +function transfermatrix_fixedpoint(tops, bottoms, ρinit) + _, vecs, info = eigsolve(ρinit, 1, :LM, Arnoldi()) do ρ + return foldr(zip(tops, bottoms); init=ρ) do (top, bottom), ρ + return @tensor ρ′[-1; -2] := top[-1 4 3; 1] * conj(bottom[-2 4 3; 2]) * ρ[1; 2] + end + end + info.converged > 0 || @warn "eigsolve did not converge" + return first(vecs) +end + +# Explicit fixing of relative phases (doing this compactly in a loop is annoying) +function _contract_gauge_corner(corner, σ_in, σ_out) + @autoopt @tensor corner_fix[χ_in; χ_out] := + σ_in[χ_in; χ1] * corner[χ1; χ2] * conj(σ_out[χ_out; χ2]) +end +function _contract_gauge_edge(edge, σ_in, σ_out) + @autoopt @tensor edge_fix[χ_in D_above D_below; χ_out] := + σ_in[χ_in; χ1] * edge[χ1 D_above D_below; χ2] * conj(σ_out[χ_out; χ2]) +end +function fix_relative_phases(envfinal::CTMRGEnv, signs) + C1 = map(Iterators.product(axes(envfinal.corners)[2:3]...)) do (r, c) + _contract_gauge_corner( + envfinal.corners[NORTHWEST, r, c], + signs[WEST, r, c], + signs[NORTH, r, _next(c, end)], + ) + end + T1 = map(Iterators.product(axes(envfinal.edges)[2:3]...)) do (r, c) + _contract_gauge_edge( + envfinal.edges[NORTH, r, c], + signs[NORTH, r, c], + signs[NORTH, r, _next(c, end)], + ) + end + C2 = map(Iterators.product(axes(envfinal.corners)[2:3]...)) do (r, c) + _contract_gauge_corner( + envfinal.corners[NORTHEAST, r, c], + signs[NORTH, r, c], + signs[EAST, _next(r, end), c], + ) + end + T2 = map(Iterators.product(axes(envfinal.edges)[2:3]...)) do (r, c) + _contract_gauge_edge( + envfinal.edges[EAST, r, c], signs[EAST, r, c], signs[EAST, _next(r, end), c] + ) + end + C3 = map(Iterators.product(axes(envfinal.corners)[2:3]...)) do (r, c) + _contract_gauge_corner( + envfinal.corners[SOUTHEAST, r, c], + signs[EAST, r, c], + signs[SOUTH, r, _prev(c, end)], + ) + end + T3 = map(Iterators.product(axes(envfinal.edges)[2:3]...)) do (r, c) + _contract_gauge_edge( + envfinal.edges[SOUTH, r, c], + signs[SOUTH, r, c], + signs[SOUTH, r, _prev(c, end)], + ) + end + C4 = map(Iterators.product(axes(envfinal.corners)[2:3]...)) do (r, c) + _contract_gauge_corner( + envfinal.corners[SOUTHWEST, r, c], + signs[SOUTH, r, c], + signs[WEST, _prev(r, end), c], + ) + end + T4 = map(Iterators.product(axes(envfinal.edges)[2:3]...)) do (r, c) + _contract_gauge_edge( + envfinal.edges[WEST, r, c], signs[WEST, r, c], signs[WEST, _prev(r, end), c] + ) + end + + return stack([C1, C2, C3, C4]; dims=1), stack([T1, T2, T3, T4]; dims=1) +end + +function calc_convergence(envs, CSold, TSold) + CSnew = map(x -> tsvd(x; alg=TensorKit.SVD())[2], envs.corners) + ΔCS = maximum(zip(CSold, CSnew)) do (c_old, c_new) + # only compute the difference on the smallest part of the spaces + smallest = infimum(MPSKit._firstspace(c_old), MPSKit._firstspace(c_new)) + e_old = isometry(MPSKit._firstspace(c_old), smallest) + e_new = isometry(MPSKit._firstspace(c_new), smallest) + return norm(e_new' * c_new * e_new - e_old' * c_old * e_old) + end + + TSnew = map(x -> tsvd(x; alg=TensorKit.SVD())[2], envs.edges) + ΔTS = maximum(zip(TSold, TSnew)) do (t_old, t_new) + MPSKit._firstspace(t_old) == MPSKit._firstspace(t_new) || + return scalartype(t_old)(Inf) + return norm(t_new - t_old) + end + + @debug "maxᵢ|Cⁿ⁺¹ - Cⁿ|ᵢ = $ΔCS maxᵢ|Tⁿ⁺¹ - Tⁿ|ᵢ = $ΔTS" + + return max(ΔCS, ΔTS), CSnew, TSnew +end + +@non_differentiable calc_convergence(args...) + +""" + calc_elementwise_convergence(envfinal, envfix; atol=1e-6) + +Check if the element-wise difference of the corner and edge tensors of the final and fixed +CTMRG environments are below some tolerance. +""" +function calc_elementwise_convergence(envfinal::CTMRGEnv, envfix::CTMRGEnv; atol::Real=1e-6) + ΔC = envfinal.corners .- envfix.corners + ΔCmax = norm(ΔC, Inf) + ΔCmean = norm(ΔC) + @debug "maxᵢⱼ|Cⁿ⁺¹ - Cⁿ|ᵢⱼ = $ΔCmax mean |Cⁿ⁺¹ - Cⁿ|ᵢⱼ = $ΔCmean" + + ΔT = envfinal.edges .- envfix.edges + ΔTmax = norm(ΔT, Inf) + ΔTmean = norm(ΔT) + @debug "maxᵢⱼ|Tⁿ⁺¹ - Tⁿ|ᵢⱼ = $ΔTmax mean |Tⁿ⁺¹ - Tⁿ|ᵢⱼ = $ΔTmean" + + # Check differences for all tensors in unit cell to debug properly + for (dir, r, c) in Iterators.product(axes(envfinal.edges)...) + @debug( + "$((dir, r, c)): all |Cⁿ⁺¹ - Cⁿ|ᵢⱼ < ϵ: ", + all(x -> abs(x) < atol, convert(Array, ΔC[dir, r, c])), + ) + @debug( + "$((dir, r, c)): all |Tⁿ⁺¹ - Tⁿ|ᵢⱼ < ϵ: ", + all(x -> abs(x) < atol, convert(Array, ΔT[dir, r, c])), + ) + end + + return max(ΔCmax, ΔTmax) +end + +@non_differentiable calc_elementwise_convergence(args...) + """ ctmrg_iter(state, env::CTMRGEnv{C,T}, alg::CTMRG) where {C,T} diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl index 6ac567a8..c517f6ff 100644 --- a/src/algorithms/peps_opt.jl +++ b/src/algorithms/peps_opt.jl @@ -85,28 +85,25 @@ based on the CTMRG gradient and updates the PEPS parameters. In this optimizatio 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. -Different levels of output verbosity can be activated using `verbosity` (0, 1 or 2). """ struct PEPSOptimize{G} boundary_alg::CTMRG optimizer::OptimKit.OptimizationAlgorithm reuse_env::Bool gradient_alg::G - verbosity::Int function PEPSOptimize( # Inner constructor to prohibit illegal setting combinations boundary_alg::CTMRG{S}, optimizer, reuse_env, gradient_alg::G, - verbosity, ) where {S,G} if gradient_alg isa GradMode if S == :sequential && G.parameters[1] == :fixed throw(ArgumentError(":sequential and :fixed are not compatible")) end end - return new{G}(boundary_alg, optimizer, reuse_env, gradient_alg, verbosity) + return new{G}(boundary_alg, optimizer, reuse_env, gradient_alg) end end function PEPSOptimize(; @@ -114,9 +111,8 @@ function PEPSOptimize(; optimizer=Defaults.optimizer, reuse_env=true, gradient_alg=LinSolver(), - verbosity=0, ) - return PEPSOptimize(boundary_alg, optimizer, reuse_env, gradient_alg, verbosity) + return PEPSOptimize(boundary_alg, optimizer, reuse_env, gradient_alg) end """ diff --git a/test/ctmrg/gaugefix.jl b/test/ctmrg/gaugefix.jl index 359bbfd2..71644ffc 100644 --- a/test/ctmrg/gaugefix.jl +++ b/test/ctmrg/gaugefix.jl @@ -3,7 +3,7 @@ using Random using PEPSKit using TensorKit -using PEPSKit: ctmrg_iter, gauge_fix, check_elementwise_convergence +using PEPSKit: ctmrg_iter, gauge_fix, calc_elementwise_convergence scalartypes = [Float64, ComplexF64] unitcells = [(1, 1), (2, 2), (3, 2)] diff --git a/test/heisenberg.jl b/test/heisenberg.jl index d72b5568..e97f7de1 100644 --- a/test/heisenberg.jl +++ b/test/heisenberg.jl @@ -21,7 +21,6 @@ opt_alg = PEPSOptimize(; optimizer=LBFGS(4; maxiter=100, gradtol=1e-3, verbosity=2), gradient_alg=LinSolver(; solver=GMRES(; tol=1e-6, maxiter=100), iterscheme=:fixed), reuse_env=true, - verbosity=2, ) # initialize states diff --git a/test/pwave.jl b/test/pwave.jl index d2f7d3aa..092aa1d2 100644 --- a/test/pwave.jl +++ b/test/pwave.jl @@ -16,7 +16,6 @@ opt_alg = PEPSOptimize(; optimizer=LBFGS(4; maxiter=10, gradtol=1e-3, verbosity=2), gradient_alg=GMRES(; tol=1e-3, maxiter=2, krylovdim=50), reuse_env=true, - verbosity=2, ) # initialize states diff --git a/test/tf_ising.jl b/test/tf_ising.jl index c40b17ba..1a3a16dd 100644 --- a/test/tf_ising.jl +++ b/test/tf_ising.jl @@ -25,7 +25,6 @@ opt_alg = PEPSOptimize(; optimizer=LBFGS(4; maxiter=100, gradtol=1e-3, verbosity=2), gradient_alg=GMRES(; tol=1e-6, maxiter=100), reuse_env=true, - verbosity=2, ) # initialize states