Merge branch 'master' into pb-ctmrg-allsides

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

@@ -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}
     

src/algorithms/peps_opt.jl

@@ -85,38 +85,34 @@ 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(;
     boundary_alg=CTMRG(),
     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
 
 """

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