Add gauge-fixing and element-wise convergence check, add function to …

…get empty projectors

Project.toml

@@ -13,6 +13,7 @@ MPSKitModels = "ca635005-6f8c-4cd1-b51d-8491250ef2ab"
 OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

examples/heisenberg.jl

@@ -42,6 +42,6 @@ optalg = PEPSOptimize{NaiveAD}(;
 
 # Ground state search
 ψinit = init_peps(2, χbond, 1, 1)
-envinit = leading_boundary(ψinit, ctmalg, CTMRGEnv(ψinit; Venv=ℂ^χenv))
+envinit, = leading_boundary(ψinit, ctmalg, CTMRGEnv(ψinit; Venv=ℂ^χenv));
 result = groundsearch(H, ctmalg, optalg, ψinit, envinit)
 @show result.E₀

src/PEPSKit.jl

@@ -1,6 +1,6 @@
 module PEPSKit
 
-using LinearAlgebra, Base.Threads, Base.Iterators, Printf
+using LinearAlgebra, Statistics, Base.Threads, Base.Iterators, Printf
 using Accessors
 using VectorInterface
 using TensorKit, KrylovKit, MPSKit, OptimKit

src/algorithms/ctmrg.jl

@@ -13,10 +13,11 @@ function MPSKit.leading_boundary(state, alg::CTMRG, envinit=CTMRGEnv(state))
     CSold = tsvd(envinit.corners[NORTHWEST]; alg=TensorKit.SVD())[2]
     TSold = tsvd(envinit.edges[NORTH]; alg=TensorKit.SVD())[2]
     ϵold = 1.0
-
     env = deepcopy(envinit)
+    Pleft, Pright = empty_projectors(eltype(env.edges), (4, size(state)...))
+
     for i in 1:(alg.maxiter)
-        env, _, _, ϵ = ctmrg_iter(state, env, alg)  # Grow and renormalize in all 4 directions
+        env, Pleft, Pright, ϵ = ctmrg_iter(state, env, alg)  # Grow and renormalize in all 4 directions
 
         # Compute convergence criteria and take max (TODO: How should we handle logging all of this?)
         Δϵ = abs((ϵold - ϵ) / ϵold)
@@ -54,36 +55,115 @@ function MPSKit.leading_boundary(state, alg::CTMRG, envinit=CTMRGEnv(state))
         ϵold = ϵ
     end
 
-    return env
+    env′, = ctmrg_iter(state, env, alg)
+    envfix, Pleftfix, Prightfix = gauge_fix(env, env′, Pleft, Pright)
+    @ignore_derivatives check_elementwise_conv(env, envfix)
+    return envfix, Pleftfix, Prightfix
+end
+
+# Fix gauge of corner end edge tensors from last and second last CTMRG iteration
+function gauge_fix(
+    envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C,T}, Pleft=nothing, Pright=nothing
+) where {C,T}
+    # Compute gauge tensors by comparing signs
+    signs = map(zip(envprev.edges, envfinal.edges)) do (Tprev, Tfinal)
+        σ = sign.(convert(Array, Tprev)[1, 1, 1, :] ./ convert(Array, Tfinal)[1, 1, 1, :])
+        Tensor(diagm(σ), space(Tfinal, 1) * space(Tfinal, 1)')
+    end
+
+    # Correct relative phases
+    cornersfix = map(Iterators.product(axes(envfinal.corners)...)) do (dir, r, c)
+        @tensor Cfix[-1; -2] :=
+            conj(signs[_prev(dir, 4), r, c][1 -1]) *
+            envfinal.corners[dir, r, c][1; 2] *
+            signs[dir, r, c][2, -2]
+    end
+    edgesfix = map(zip(signs, envfinal.edges)) do (σ, edge)
+        @tensor Tfix[-1 -2 -3; -4] := conj(σ[1 -1]) * edge[1 -2 -3; 2] * σ[2, -4]
+    end
+
+    # Fix global phase
+    cornersgfix = map(zip(envprev.corners, cornersfix)) do (Cprev, Cfix)
+        φ = tr(Cprev) / tr(Cfix)  # Extract phase via trace to make it differentiable
+        return φ * Cfix
+    end
+    envfix = CTMRGEnv(cornersgfix, edgesfix)
+
+    # Gauge projectors for correct backpropagation
+    if !isnothing(Pleft) && !isnothing(Pright)
+        Pleftfix = map(zip(signs, Pleft)) do (σ, P)
+            @tensor Plfix[-1 -2 -3; -4] := P[-1 -2 -3; 1] * σ[1; -4]
+        end
+        Prightfix = map(zip(signs, Pright)) do (σ, P)
+            @tensor Prfix[-1; -2 -3 -4] := σ[-1; 1] * P[1; -2 -3 -4]
+        end
+        return envfix, Pleftfix, Prightfix
+    else
+        return envfix
+    end
+end
+
+# Explicitly check if element-wise difference of fixed and final environment tensors are below some tolerance
+function check_elementwise_conv(
+    envfinal::CTMRGEnv, envfix::CTMRGEnv; atol::Real=1e-6, print_conv=true
+)
+    ΔC = map(zip(envfinal.corners, envfix.corners)) do (Cfin, Cfix)
+        return abs.(convert(Array, Cfix - Cfin))
+    end
+    ΔCmax = maximum(maximum, ΔC)
+    ΔCmean = maximum(mean, ΔC)
+    Cerr = map(δ -> all(x -> x < atol, δ), ΔC)
+
+    ΔT = map(zip(envfinal.edges, envfix.edges)) do (Tfin, Tfix)
+        return abs.(convert(Array, Tfix - Tfin))
+    end
+    ΔTmax = maximum(maximum, ΔT)
+    ΔTmean = maximum(mean, ΔT)
+    Terr = map(δ -> all(x -> x < atol, δ), ΔT)
+
+    if print_conv
+        if all(Cerr) && all(Terr)
+            println("{Cᵢ,Tᵢ} converged elementwise up to ϵ < $atol")
+        else
+            @warn "no elementwise convergence up to ϵ < $atol:"
+            for i in 1:4
+                println("$i: all |Cⁿ⁺¹ - Cⁿ|ᵢⱼ < ϵ: ", Cerr[i])
+                println("$i: all |Tⁿ⁺¹ - Tⁿ|ᵢⱼ < ϵ: ", Terr[i])
+            end
+        end
+        println("maxᵢⱼ|Cⁿ⁺¹ - Cⁿ|ᵢⱼ = $ΔCmax")
+        println("mean |Cⁿ⁺¹ - Cⁿ|ᵢⱼ = $ΔCmean")
+        println("maxᵢⱼ|Tⁿ⁺¹ - Tⁿ|ᵢⱼ = $ΔTmax")
+        println("mean |Tⁿ⁺¹ - Tⁿ|ᵢⱼ = $ΔTmean")
+    end
+
+    return Cerr, Terr
 end
 
 # One CTMRG iteration x′ = f(A, x)
 function ctmrg_iter(state, env::CTMRGEnv{C,T}, alg::CTMRG) where {C,T}
     ϵ = 0.0
-    Prtype = tensormaptype(spacetype(T), numin(T), numout(T), storagetype(T))
-    Pleft = Vector{Matrix{T}}(undef, 4)
-    Pright = Vector{Matrix{Prtype}}(undef, 4)
+    Pleft, Pright = empty_projectors(T, (4, size(state)...))
 
     for i in 1:4
         env, Pl, Pr, ϵ₀ = left_move(state, env, alg)
         state = rotate_north(state, EAST)
         env = rotate_north(env, EAST)
         ϵ = max(ϵ, ϵ₀)
-        @diffset Pleft[i] = Pl
-        @diffset Pright[i] = Pr
+        @diffset Pleft[i, :, :] .= Pl
+        @diffset Pright[i, :, :] .= Pr
     end
 
     return env, Pleft, Pright, ϵ
 end
 
-#  row environment, compute projectors and renormalize
-function left_move(state, env::CTMRGEnv, alg::CTMRG)
+# Grow environment, compute projectors and renormalize
+function left_move(state, env::CTMRGEnv{C,T}, alg::CTMRG) where {C,T}
     corners::typeof(env.corners) = copy(env.corners)
     edges::typeof(env.edges) = copy(env.edges)
     ϵ = 0.0
+    Pleft, Pright = empty_projectors(T, size(state))
 
-    Pleft = similar(state.A, typeof(env.edges[1]))
-    Pright = similar(state.A, typeof(transpose(env.edges[1])))
     for col in 1:size(state, 2)
         cnext = _next(col, size(state, 2))
 
@@ -198,20 +278,6 @@ end
 
 # Apply projectors to entire left half-environment to grow SW & NW corners, and W edge
 function grow_env_left(peps, Pl, Pr, C_sw, C_nw, T_s, T_w, T_n)
-    # @diffset @tensor corners[NORTHWEST, rop, cop][-1; -2] :=
-    #     envs.corners[NORTHWEST, rop, col][1, 2] *
-    #     envs.edges[NORTH, rop, col][2, 3, 4, -2] *
-    #     Q[-1; 1 3 4]
-    # @diffset @tensor corners[SOUTHWEST, rom, cop][-1; -2] :=
-    #     envs.corners[SOUTHWEST, rom, col][1, 4] *
-    #     envs.edges[SOUTH, rom, col][-1, 2, 3, 1] *
-    #     P[4 2 3; -2]
-    # @diffset @tensor edges[WEST, row, cop][-1 -2 -3; -4] :=
-    #     envs.edges[WEST, row, col][1 2 3; 4] *
-    #     peps_above[row, col][9; 5 -2 7 2] *
-    #     conj(peps_below[row, col][9; 6 -3 8 3]) *
-    #     P[4 5 6; -4] *
-    #     Q[-1; 1 7 8]
     @tensor C_sw′[-1; -2] := C_sw[1; 4] * T_s[-1 2 3; 1] * Pl[4 2 3; -2]
     @tensor C_nw′[-1; -2] := C_nw[1; 2] * T_n[2 3 4; -2] * Pr[-1; 1 3 4]
     @tensor T_w′[-1 -2 -3; -4] :=

src/algorithms/peps_opt.jl

@@ -47,14 +47,14 @@ end
 # Energy cost function with proper backwards rule depending only on final CTMRG fixed-point
 # Mutates environment to reuse previous environments in optimization
 function costfun!(peps, env, H, ctmalg::CTMRG, optalg::PEPSOptimize)
-    env′ = leading_boundary(peps, ctmalg, env)
+    env′, = leading_boundary(peps, ctmalg, env)
     @diffset env = env′
     return optalg.energyfun(peps, env′, H)
 end
 
 # Non-mutating version, recomputing environment from random initial guess in every optimization step
 function costfun(peps, env, H, ctmalg::CTMRG, optalg::PEPSOptimize)
-    env′ = deepcopy(env)  # Create copy to make non-mutating
+    env′, = deepcopy(env)  # Create copy to make non-mutating
     return costfun!(peps, env′, H, ctmalg, optalg)
 end
 
@@ -80,7 +80,7 @@ end
 
 # Contraction of CTMRGEnv and PEPS tensors with open physical bonds
 function one_site_rho(peps::InfinitePEPS, env::CTMRGEnv{C,T}) where {C,T}
-    ρunitcell = map(Iterators.product(axes(env.corners, 2), axes(env.corners, 3))) do (r, c)
+    return map(Iterators.product(axes(env.corners, 2), axes(env.corners, 3))) do (r, c)
         @tensor ρ[-1; -2] :=
             env.corners[NORTHWEST, r, c][1; 2] *
             env.edges[NORTH, r, c][2 3 4; 5] *
@@ -92,14 +92,12 @@ function one_site_rho(peps::InfinitePEPS, env::CTMRGEnv{C,T}) where {C,T}
             env.edges[WEST, r, c][14 15 16; 1] *
             peps[r, c][-1; 3 7 11 15] *
             conj(peps[r, c][-2; 4 8 12 16])
-        return ρ
     end
-    return ρunitcell
 end
 
 # Horizontally extended contraction of CTMRGEnv and PEPS tensors with open physical bonds
 function two_site_rho(peps::InfinitePEPS, env::CTMRGEnv{C,T}) where {C,T}
-    ρunitcell = map(Iterators.product(axes(env.corners, 2), axes(env.corners, 3))) do (r, c)
+    return map(Iterators.product(axes(env.corners, 2), axes(env.corners, 3))) do (r, c)
         cnext = _next(c, size(peps, 2))
         @tensor ρ[-11 -20; -12 -18] :=
             env.corners[NORTHWEST, r, c][1; 3] *
@@ -116,9 +114,7 @@ function two_site_rho(peps::InfinitePEPS, env::CTMRGEnv{C,T}) where {C,T}
             conj(peps[r, c][-11; 8 19 10 9]) *
             peps[r, cnext][-18; 16 25 17 15] *
             conj(peps[r, cnext][-20; 22 26 21 19])
-        return ρ
     end
-    return ρunitcell
 end
 
 # 1-site operator expectation values on unit cell

src/utility/util.jl

@@ -47,6 +47,14 @@ function is_degenerate_spectrum(
     return false
 end
 
+# Create empty projectors for given state without recomputing transpose
+function empty_projectors(T::DataType, size)
+    Pleft = Array{T,length(size)}(undef, size)
+    Prtype = tensormaptype(spacetype(T), numin(T), numout(T), storagetype(T))
+    Pright = Array{Prtype,length(size)}(undef, size)
+    return Pleft, Pright
+end
+
 # rotl90 is set to non_differentiable in ChainRules
 function ChainRulesCore.rrule(::typeof(rotl90), a::AbstractMatrix)
     function pb(x)