Logging utilities

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"

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")

src/algorithms/ctmrg.jl

@@ -34,8 +34,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). The projectors are computed from `svd_alg` SVDs
-where the truncation scheme is set via `trscheme`.
+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 
+`trscheme`.
 """
 struct CTMRG
     tol::Float64
@@ -67,70 +69,53 @@ 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)
-    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
+    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, ϵ = 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 - ϵ) / ϵ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)
+        # 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′)
 
-            Δ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
+        η = calc_elementwise_convergence(envfix, env; atol=alg.tol^(1 / 2))
+        N = norm(state, envfix)
 
-            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,
-                    ϵ,
-                    Δϵ
-                )
-            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, ϵ
+        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
-
-    # 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′)
-    check_elementwise_convergence(env, envfix; atol=alg.tol^(1 / 2)) ||
-        @warn "CTMRG did not converge elementwise."
-    return envfix
 end
 
+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}
 
@@ -275,15 +260,37 @@ function fix_relative_phases(envfinal::CTMRGEnv, signs)
     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...)
+
 """
-    check_elementwise_convergence(envfinal, envfix; atol=1e-6)
+    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 check_elementwise_convergence(
-    envfinal::CTMRGEnv, envfix::CTMRGEnv; atol::Real=1e-6
-)
+function calc_elementwise_convergence(envfinal::CTMRGEnv, envfix::CTMRGEnv; atol::Real=1e-6)
     ΔC = envfinal.corners .- envfix.corners
     ΔCmax = norm(ΔC, Inf)
     ΔCmean = norm(ΔC)
@@ -306,10 +313,10 @@ function check_elementwise_convergence(
         )
     end
 
-    return isapprox(ΔCmax, 0; atol) && isapprox(ΔTmax, 0; atol)
+    return max(ΔCmax, ΔTmax)
 end
 
-@non_differentiable check_elementwise_convergence(args...)
+@non_differentiable calc_elementwise_convergence(args...)
 
 """
     ctmrg_iter(state, env::CTMRGEnv{C,T}, alg::CTMRG) where {C,T}

src/algorithms/peps_opt.jl

@@ -40,7 +40,6 @@ 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).
 """
 @kwdef struct PEPSOptimize{G}
     boundary_alg::CTMRG = CTMRG()  # Algorithm to find boundary environment
@@ -49,7 +48,6 @@ Different levels of output verbosity can be activated using `verbosity` (0, 1 or
     )
     reuse_env::Bool = true  # Reuse environment of previous optimization as initial guess for next
     gradient_alg::G = GeomSum() # Algorithm to solve gradient linear problem
-    verbosity::Int = 0
 end
 
 """

test/ctmrg/gaugefix.jl

@@ -4,11 +4,12 @@ using PEPSKit
 using TensorKit
 using Accessors
 
-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)]
 χ = 6
+verbosity = 3
 
 function _make_symmetric(psi)
     if ==(size(psi)...)
@@ -45,7 +46,6 @@ end
     Random.seed!(987654321)  # Seed RNG to make random environment consistent
     ctm = CTMRGEnv(psi, ctm_space)
 
-    verbosity = 1
     alg = CTMRG(;
         tol=1e-10, miniter=4, maxiter=400, verbosity, trscheme=truncdim(dim(ctm_space))
     )
@@ -54,7 +54,7 @@ end
     ctm = leading_boundary(ctm, psi, alg)
     ctm2, = ctmrg_iter(psi, ctm, alg_fixed)
     ctm_fixed = gauge_fix(ctm, ctm2)
-    @test PEPSKit.check_elementwise_convergence(ctm, ctm_fixed; atol=1e-4)
+    @test PEPSKit.calc_elementwise_convergence(ctm, ctm_fixed) ≈ 0 atol = 1e-4
 end
 
 @testset "Z2 symmetry ($T) - ($unitcell)" for (T, unitcell) in
@@ -70,7 +70,6 @@ end
     Random.seed!(123456789)  # Seed RNG to make random environment consistent
     ctm = CTMRGEnv(psi, ctm_space)
 
-    verbosity = 1
     alg = CTMRG(;
         tol=1e-10, miniter=4, maxiter=400, verbosity, trscheme=truncdim(dim(ctm_space))
     )
@@ -79,5 +78,5 @@ end
     ctm = leading_boundary(ctm, psi, alg)
     ctm2, = ctmrg_iter(psi, ctm, alg_fixed)
     ctm_fixed = gauge_fix(ctm, ctm2)
-    @test PEPSKit.check_elementwise_convergence(ctm, ctm_fixed; atol=1e-4)
+    @test PEPSKit.calc_elementwise_convergence(ctm, ctm_fixed) ≈ 0 atol = 1e-4
 end

test/heisenberg.jl

@@ -14,7 +14,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

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

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