Fix gauge-fixing for larger unit cells, decrease default tol for elem…

…ent-wise convergence check

examples/heisenberg.jl

@@ -27,7 +27,8 @@ ctmalg = CTMRG(; trscheme=truncdim(χenv), tol=1e-10, miniter=4, maxiter=100, ve
 alg = PEPSOptimize(;
     boundary_alg=ctmalg,
     optimizer=LBFGS(4; maxiter=100, gradtol=1e-4, verbosity=2),
-    gradient_alg=GMRES(; tol=1e-6, maxiter=100),
+    # gradient_alg=GMRES(; tol=1e-6, maxiter=100),
+    gradient_alg=ManualIter(; tol=1e-6, maxiter=100),
     reuse_env=true,
     verbosity=2,
 )

examples/test_gauge_fixing.jl

@@ -34,6 +34,7 @@ E = ℂ^χenv # environment virtual space
 
 test_gauge_fixing(P, V, E; χenv, unitcell=(1, 1))
 test_gauge_fixing(P, V, E; χenv, unitcell=(2, 2))
+test_gauge_fixing(P, V, E; χenv, unitcell=(3, 4)) # check gauge-fixing for unit cells > (2, 2)
 
 # Z2
 

examples/test_gradients.jl

@@ -2,12 +2,8 @@ using LinearAlgebra
 using TensorKit, MPSKitModels, OptimKit
 using PEPSKit, KrylovKit
 
-
 using Zygote
 
-
-
-
 # Square lattice Heisenberg Hamiltonian
 function square_lattice_heisenberg(; Jx=-1.0, Jy=1.0, Jz=-1.0)
     Sx, Sy, Sz, _ = spinmatrices(1//2)

src/algorithms/ctmrg.jl

@@ -70,7 +70,7 @@ function MPSKit.leading_boundary(state, alg::CTMRG, envinit=CTMRGEnv(state))
 
     env′, = ctmrg_iter(state, env, alg)
     envfix = gauge_fix(env, env′)
-    check_elementwise_convergence(env, envfix; atol=alg.tol) ||
+    check_elementwise_convergence(env, envfix; atol=alg.tol^(3 / 4)) ||
         @warn "CTMRG did not converge elementwise."
     return envfix
 end
@@ -97,13 +97,19 @@ function gauge_fix(envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C,T}) where {C,T}
 
     # 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 InfinitMPSs
-        if isodd(dir)
+        # 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)
-        else
+        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
@@ -117,13 +123,16 @@ function gauge_fix(envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C,T}) where {C,T}
             scalartype(T),
             MPSKit._lastspace(Tsfinal[end])' ← MPSKit._lastspace(M[end])',
         )
+
         rhoprev = eigsolve(TransferMatrix(Tsprev, M), rhoinit, 1, :LM)[2][1]
         rhofinal = eigsolve(TransferMatrix(Tsfinal, M), rhoinit, 1, :LM)[2][1]
         # rhoprev = @showgrad rhoprev
+
         # Decompose and multiply
         Qprev, = leftorth(rhoprev, ((1,), (2,)))
         Qfinal, = leftorth(rhofinal, ((1,), (2,)))
         σ = @checkgrad Qprev * Qfinal'
+        # σ = Qprev * Qfinal'
 
         return σ
     end
@@ -144,107 +153,63 @@ function gauge_fix(envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C,T}) where {C,T}
     return envfix
 end
 
-# Explicit unrolling of for loop from previous version to fix AD
-# TODO: Does not yet properly work for Lx,Ly > 2
+# Explicit fixing of relative phases (doing this compactly in a loop is annoying)
 function fix_relative_phases(envfinal::CTMRGEnv, signs)
-    e1 = envfinal
-    σ1 = signs
-    C1 = map(Iterators.product(axes(e1.corners)[2:3]...)) do (r, c)
+    C1 = map(Iterators.product(axes(envfinal.corners)[2:3]...)) do (r, c)
         @tensor Cfix[-1; -2] :=
-            σ1[WEST, _prev(r, end), c][-1 1] *
-            e1.corners[NORTHWEST, r, c][1; 2] *
-            conj(σ1[NORTH, r, c][-2 2])
+            signs[WEST, _prev(r, end), c][-1 1] *
+            envfinal.corners[NORTHWEST, r, c][1; 2] *
+            conj(signs[NORTH, r, c][-2 2])
     end
-    T1 = map(Iterators.product(axes(e1.edges)[2:3]...)) do (r, c)
+    T1 = map(Iterators.product(axes(envfinal.edges)[2:3]...)) do (r, c)
         @tensor Tfix[-1 -2 -3; -4] :=
-            σ1[NORTH, r, c][-1 1] *
-            e1.edges[NORTH, r, c][1 -2 -3; 2] *
-            conj(σ1[NORTH, r, _next(c, end)][-4 2])
+            signs[NORTH, r, c][-1 1] *
+            envfinal.edges[NORTH, r, c][1 -2 -3; 2] *
+            conj(signs[NORTH, r, _next(c, end)][-4 2])
     end
 
-    e2 = rotate_north(envfinal, EAST)
-    σ2 = rotate_north(signs, EAST)
-    C2 = map(Iterators.product(axes(e2.corners)[2:3]...)) do (r, c)
+    C2 = map(Iterators.product(axes(envfinal.corners)[2:3]...)) do (r, c)
         @tensor Cfix[-1; -2] :=
-            σ2[WEST, _prev(r, end), c][-1 1] *
-            e2.corners[NORTHWEST, r, c][1; 2] *
-            conj(σ2[NORTH, r, c][-2 2])
+            signs[NORTH, r, _next(c, end)][-1 1] *
+            envfinal.corners[NORTHEAST, r, c][1; 2] *
+            conj(signs[EAST, r, c][-2 2])
     end
-    C2 = rotate_north(C2, WEST)
-    T2 = map(Iterators.product(axes(e2.edges)[2:3]...)) do (r, c)
+    T2 = map(Iterators.product(axes(envfinal.edges)[2:3]...)) do (r, c)
         @tensor Tfix[-1 -2 -3; -4] :=
-            σ2[NORTH, r, c][-1 1] *
-            e2.edges[NORTH, r, c][1 -2 -3; 2] *
-            conj(σ2[NORTH, r, _next(c, end)][-4 2])
+            signs[EAST, r, c][-1 1] *
+            envfinal.edges[EAST, r, c][1 -2 -3; 2] *
+            conj(signs[EAST, _next(r, end), c][-4 2])
     end
-    T2 = rotate_north(T2, WEST)
 
-    e3 = rotate_north(envfinal, SOUTH)
-    σ3 = rotate_north(signs, SOUTH)
-    C3 = map(Iterators.product(axes(e3.corners)[2:3]...)) do (r, c)
+    C3 = map(Iterators.product(axes(envfinal.corners)[2:3]...)) do (r, c)
         @tensor Cfix[-1; -2] :=
-            σ3[WEST, _prev(r, end), c][-1 1] *
-            e3.corners[NORTHWEST, r, c][1; 2] *
-            conj(σ3[NORTH, r, c][-2 2])
+            signs[EAST, _next(r, end), c][-1 1] *
+            envfinal.corners[SOUTHEAST, r, c][1; 2] *
+            conj(signs[SOUTH, r, c][-2 2])
     end
-    C3 = rotate_north(C3, SOUTH)
-    T3 = map(Iterators.product(axes(e3.edges)[2:3]...)) do (r, c)
+    T3 = map(Iterators.product(axes(envfinal.edges)[2:3]...)) do (r, c)
         @tensor Tfix[-1 -2 -3; -4] :=
-            σ3[NORTH, r, c][-1 1] *
-            e3.edges[NORTH, r, c][1 -2 -3; 2] *
-            conj(σ3[NORTH, r, _next(c, end)][-4 2])
+            signs[SOUTH, r, c][-1 1] *
+            envfinal.edges[SOUTH, r, c][1 -2 -3; 2] *
+            conj(signs[SOUTH, r, _prev(c, end)][-4 2])
     end
-    T3 = rotate_north(T3, SOUTH)
 
-    e4 = rotate_north(envfinal, WEST)
-    σ4 = rotate_north(signs, WEST)
-    C4 = map(Iterators.product(axes(e4.corners)[2:3]...)) do (r, c)
+    C4 = map(Iterators.product(axes(envfinal.corners)[2:3]...)) do (r, c)
         @tensor Cfix[-1; -2] :=
-            σ4[WEST, _prev(r, end), c][-1 1] *
-            e4.corners[NORTHWEST, r, c][1; 2] *
-            conj(σ4[NORTH, r, c][-2 2])
+            signs[SOUTH, r, _prev(c, end)][-1 1] *
+            envfinal.corners[SOUTHWEST, r, c][1; 2] *
+            conj(signs[WEST, r, c][-2 2])
     end
-    C4 = rotate_north(C4, EAST)
-    T4 = map(Iterators.product(axes(e4.edges)[2:3]...)) do (r, c)
+    T4 = map(Iterators.product(axes(envfinal.edges)[2:3]...)) do (r, c)
         @tensor Tfix[-1 -2 -3; -4] :=
-            σ4[NORTH, r, c][-1 1] *
-            e4.edges[NORTH, r, c][1 -2 -3; 2] *
-            conj(σ4[NORTH, r, _next(c, end)][-4 2])
+            signs[WEST, r, c][-1 1] *
+            envfinal.edges[WEST, r, c][1 -2 -3; 2] *
+            conj(signs[WEST, _prev(r, end), c][-4 2])
     end
-    T4 = rotate_north(T4, EAST)
 
     return stack([C1, C2, C3, C4]; dims=1), stack([T1, T2, T3, T4]; dims=1)
 end
 
-# Semi-working version analogous to left_move with rotations
-# function fix_relative_phases(envfinal::CTMRGEnv, signs)
-#     cornersfix = deepcopy(envfinal.corners)
-#     edgesfix = deepcopy(envfinal.edges)
-#     for _ in 1:4
-#         corners = map(Iterators.product(axes(envfinal.corners)[2:3]...)) do (r, c)
-#             @tensor Cfix[-1; -2] :=
-#                 signs[WEST, _prev(r, end), c][-1 1] *
-#                 envfinal.corners[NORTHWEST, r, c][1; 2] *
-#                 conj(signs[NORTH, r, c][-2 2])
-#         end
-#         @diffset cornersfix[NORTHWEST, :, :] .= corners
-#         edges = map(Iterators.product(axes(envfinal.edges)[2:3]...)) do (r, c)
-#             @tensor Tfix[-1 -2 -3; -4] :=
-#                 signs[NORTH, r, c][-1 1] *
-#                 envfinal.edges[NORTH, r, c][1 -2 -3; 2] *
-#                 conj(signs[NORTH, r, _next(c, end)][-4 2])
-#         end
-#         @diffset edgesfix[NORTH, :, :] .= edges
-
-#         # Rotate east-wards
-#         envfinal = rotate_north(envfinal, EAST)
-#         cornersfix = rotate_north(cornersfix, EAST)
-#         edgesfix = rotate_north(edgesfix, EAST)
-#         signs = rotate_north(signs, EAST)  # TODO: Fix AD problem here
-#     end
-#     return cornersfix, edgesfix
-# end
-
 """
     check_elementwise_convergence(envfinal, envfix; atol=1e-6)
 
@@ -254,7 +219,6 @@ CTMRG environments are below some tolerance.
 function check_elementwise_convergence(
     envfinal::CTMRGEnv, envfix::CTMRGEnv; atol::Real=1e-6
 )
-    # TODO: do we need both max and mean?
     ΔC = envfinal.corners .- envfix.corners
     ΔCmax = norm(ΔC, Inf)
     ΔCmean = norm(ΔC)
@@ -265,6 +229,18 @@ function check_elementwise_convergence(
     Δ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, ΔC[dir, r, c].data),
+        )
+        @debug(
+            "$((dir, r, c)): all |Tⁿ⁺¹ - Tⁿ|ᵢⱼ < ϵ: ",
+            all(x -> abs(x) < atol, ΔT[dir, r, c].data),
+        )
+    end
+
     return isapprox(ΔCmax, 0; atol) && isapprox(ΔTmax, 0; atol)
 end