Add fixedpoint interface for PEPSOptimize, add NLocalOperator type wi…

…th AbstractInteractions, unroll gauge-fixing explicitly to fix AD, update ctmrg_gradient to be more efficient

examples/heisenberg.jl

@@ -16,7 +16,7 @@ function square_lattice_heisenberg(; Jx=-1, Jy=1, Jz=-1)
         Jy * σy[-1, -2] * σy[-3, -4] +
         Jz * σz[-1, -2] * σz[-3, -4]
 
-    return H
+    return NLocalOperator{NearestNeighbor}(H)
 end
 
 # Initialize InfinitePEPS with random & complex entries by default
@@ -32,15 +32,17 @@ H = square_lattice_heisenberg(; Jx=-1, Jy=1, Jz=-1)
 χbond = 2
 χenv = 20
 ctmalg = CTMRG(; trscheme=truncdim(χenv), tol=1e-10, miniter=4, maxiter=100, verbosity=2)
-optalg = PEPSOptimize{LinSolve}(;
+alg = PEPSOptimize{LinSolve}(;
+    boundary_alg=ctmalg,
     optimizer=LBFGS(4; maxiter=100, gradtol=1e-4, verbosity=2),
     fpgrad_tol=1e-6,
     fpgrad_maxiter=100,
+    reuse_env=true,
     verbosity=2,
 )
 
 # Ground state search
-ψinit = init_peps(2, χbond, 1, 1)
-envinit = leading_boundary(ψinit, ctmalg, CTMRGEnv(ψinit; Venv=ℂ^χenv))
-result = groundsearch(H, ctmalg, optalg, ψinit, envinit)
-@show result.E₀
+ψ₀ = init_peps(2, χbond, 1, 1)
+env₀ = leading_boundary(ψ₀, ctmalg, CTMRGEnv(ψ₀; Venv=ℂ^χenv))
+result = fixedpoint(ψ₀, H, alg, env₀)
+@show result.E

examples/test_gauge_fixing.jl

@@ -2,22 +2,6 @@ using LinearAlgebra
 using TensorKit, MPSKitModels, OptimKit
 using PEPSKit
 
-# Square lattice Heisenberg Hamiltonian
-function square_lattice_heisenberg(; Jx=-1.0, Jy=1.0, Jz=-1.0)
-    Sx, Sy, Sz, _ = spinmatrices(1//2)
-    Vphys = ℂ^2
-    σx = TensorMap(Sx, Vphys, Vphys)
-    σy = TensorMap(Sy, Vphys, Vphys)
-    σz = TensorMap(Sz, Vphys, Vphys)
-
-    @tensor H[-1 -3; -2 -4] :=
-        Jx * σx[-1, -2] * σx[-3, -4] +
-        Jy * σy[-1, -2] * σy[-3, -4] +
-        Jz * σz[-1, -2] * σz[-3, -4]
-
-    return H
-end
-
 # Initialize InfinitePEPS with random & complex entries by default
 function init_peps(d, D, Lx, Ly, finit=randn, dtype=ComplexF64)
     Pspaces = fill(ℂ^d, Lx, Ly)
@@ -27,17 +11,16 @@ function init_peps(d, D, Lx, Ly, finit=randn, dtype=ComplexF64)
 end
 
 # Initialize PEPS and environment
-H = square_lattice_heisenberg()
 χbond = 2
 χenv = 20
-ctmalg = CTMRG(; trscheme=truncdim(χenv), tol=1e-10, miniter=4, maxiter=50, verbosity=2)
-ψ = init_peps(2, χbond, 3, 1)
+ctmalg = CTMRG(; trscheme=truncdim(χenv), tol=1e-10, miniter=4, maxiter=100, verbosity=2)
+ψ = init_peps(2, χbond, 2, 2)
 env = leading_boundary(ψ, ctmalg, CTMRGEnv(ψ; Venv=ℂ^χenv))
 
 println("\nBefore gauge-fixing:")
 env′, = PEPSKit.ctmrg_iter(ψ, env, ctmalg)
 PEPSKit.check_elementwise_conv(env, env′)
 
 println("\nAfter gauge-fixing:")
-envfix = PEPSKit.gauge_fix(env, env′)
+envfix = PEPSKit.gauge_fix(env, env′);
 PEPSKit.check_elementwise_conv(env, envfix)

examples/test_gradients.jl

@@ -15,7 +15,7 @@ function square_lattice_heisenberg(; Jx=-1.0, Jy=1.0, Jz=-1.0)
         Jy * σy[-1, -2] * σy[-3, -4] +
         Jz * σz[-1, -2] * σz[-3, -4]
 
-    return H
+    return NLocalOperator{NearestNeighbor}(H)
 end
 
 # Initialize InfinitePEPS with random & complex entries by default
@@ -36,27 +36,26 @@ env = leading_boundary(ψ, ctmalg, CTMRGEnv(ψ; Venv=ℂ^χenv))
 
 # Compute CTM gradient in four different ways (set reuse_env=false to not mutate environment)
 println("\nFP gradient using naive AD:")
-alg_naive = PEPSOptimize{NaiveAD}(; verbosity=2, reuse_env=false)
-@time _, g_naive = PEPSKit.ctmrg_gradient((ψ, env), H, ctmalg, alg_naive)
-g_naive = InfinitePEPS(g_naive...)  # Convert NamedTuple to InfinitePEPS
+alg_naive = PEPSOptimize{NaiveAD}(; boundary_alg=ctmalg, verbosity=2, reuse_env=false)
+@time _, g_naive = PEPSKit.ctmrg_gradient((ψ, env), H, alg_naive)
 
 println("\nFP gradient using explicit evaluation of the geometric sum:")
 alg_geomsum = PEPSOptimize{GeomSum}(;
-    fpgrad_tol=1e-6, fpgrad_maxiter=100, verbosity=2, reuse_env=false
+    boundary_alg=ctmalg, fpgrad_tol=1e-6, fpgrad_maxiter=100, verbosity=2, reuse_env=false
 )
-@time _, g_geomsum = PEPSKit.ctmrg_gradient((ψ, env), H, ctmalg, alg_geomsum)
+@time _, g_geomsum = PEPSKit.ctmrg_gradient((ψ, env), H, alg_geomsum)
 
 println("\nFP gradient using manual iteration of the linear problem:")
 alg_maniter = PEPSOptimize{ManualIter}(;
-    fpgrad_tol=1e-6, fpgrad_maxiter=100, verbosity=2, reuse_env=false
+    boundary_alg=ctmalg, fpgrad_tol=1e-6, fpgrad_maxiter=100, verbosity=2, reuse_env=false
 )
-@time _, g_maniter = PEPSKit.ctmrg_gradient((ψ, env), H, ctmalg, alg_maniter)
+@time _, g_maniter = PEPSKit.ctmrg_gradient((ψ, env), H, alg_maniter)
 
 println("\nFP gradient using GMRES to solve the linear problem:")
 alg_linsolve = PEPSOptimize{LinSolve}(;
-    fpgrad_tol=1e-6, fpgrad_maxiter=100, verbosity=2, reuse_env=false
+    boundary_alg=ctmalg, fpgrad_tol=1e-6, fpgrad_maxiter=100, verbosity=2, reuse_env=false
 )
-@time _, g_linsolve = PEPSKit.ctmrg_gradient((ψ, env), H, ctmalg, alg_linsolve)
+@time _, g_linsolve = PEPSKit.ctmrg_gradient((ψ, env), H, alg_linsolve)
 
 @show norm(g_geomsum - g_naive) / norm(g_naive)
 @show norm(g_maniter - g_naive) / norm(g_naive)

src/PEPSKit.jl

@@ -22,6 +22,7 @@ include("mpskit_glue/transferpepo_environments.jl")
 
 include("environments/ctmrgenv.jl")
 include("environments/boundarympsenv.jl")
+include("operators/localoperator.jl")
 
 include("algorithms/ctmrg.jl")
 include("algorithms/peps_opt.jl")
@@ -34,14 +35,16 @@ module Defaults
     const ctmrg_maxiter = 100
     const ctmrg_miniter = 4
     const ctmrg_tol = 1e-12
-    const grad_maxiter = 100
-    const grad_tol = 1e-6
+    const fpgrad_maxiter = 100
+    const fpgrad_tol = 1e-6
 end
 
 export CTMRG, CTMRGEnv
-export leading_boundary, expectation_value
+export NLocalOperator, OnSite, NearestNeighbor
+export expectation_value, costfun
+export leading_boundary
 export PEPSOptimize, NaiveAD, GeomSum, ManualIter, LinSolve
-export groundsearch
+export fixedpoint
 export InfinitePEPS, InfiniteTransferPEPS
 export InfinitePEPO, InfiniteTransferPEPO
 export initializeMPS, initializePEPS

src/algorithms/ctmrg.jl

@@ -9,6 +9,7 @@
 end
 
 # Compute CTMRG environment for a given state
+# function MPSKit.leading_boundary(state, alg::CTMRG, envinit=CTMRGEnv(state))
 function MPSKit.leading_boundary(state, alg::CTMRG, envinit=CTMRGEnv(state))
     normold = 1.0
     CSold = map(x -> tsvd(x; alg=TensorKit.SVD())[2], envinit.corners)
@@ -22,7 +23,7 @@ function MPSKit.leading_boundary(state, alg::CTMRG, envinit=CTMRGEnv(state))
         # 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)
+        Δnorm = abs(normold - normnew) / abs(normold)
         CSnew = map(c -> tsvd(c; alg=TensorKit.SVD())[2], env.corners)
         ΔCS = maximum(norm.(CSnew - CSold))
         TSnew = map(t -> tsvd(t; alg=TensorKit.SVD())[2], env.edges)
@@ -87,103 +88,119 @@ function gauge_fix(envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C,T}) where {C,T}
 
     # Fix global phase
     cornersgfix = map(zip(envprev.corners, cornersfix)) do (Cprev, Cfix)
-        # φ = sign(sum(Cprev.data) / sum(Cfix.data))  # TODO: make sum(A.data) differentiable
-        φ = sign(tr(Cprev) / tr(Cfix))
-        φ * Cfix
+        φ = dot(Cprev, Cfix)
+        φ' * Cfix
     end
     edgesgfix = map(zip(envprev.edges, edgesfix)) do (Tprev, Tfix)
-        # φ = sign(sum(Tprev.data) / sum(Tfix.data))
-        φ = sign((@tensor Tprev[1 2 2; 1]) / (@tensor Tfix[1 2 2; 1]))
-        φ * Tfix
+        φ = dot(Tprev, Tfix)
+        φ' * Tfix
     end
     envfix = CTMRGEnv(cornersgfix, edgesgfix)
 
     return envfix
 end
 
-# Explicit version
-# function fix_relative_phases(e::CTMRGEnv, signs)
-#     cornersfix = map(Iterators.product(axes(e.corners)...)) do (dir, r, c)
-#         if dir == 1
-#             return @tensor Cfix[-1; -2] :=
+# Explicit unrolling of for loop from previous version to fix AD
+# TODO: Does not yet properly work for Lx,Ly > 2
+function fix_relative_phases(envfinal::CTMRGEnv, signs)
+    e1 = envfinal
+    σ1 = signs
+    C1 = map(Iterators.product(axes(e1.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])
+    end
+    T1 = map(Iterators.product(axes(e1.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])
+    end
+
+    e2 = rotate_north(envfinal, EAST)
+    σ2 = rotate_north(signs, EAST)
+    C2 = map(Iterators.product(axes(e2.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])
+    end
+    C2 = rotate_north(C2, WEST)
+    T2 = map(Iterators.product(axes(e2.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])
+    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)
+        @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])
+    end
+    C3 = rotate_north(C3, SOUTH)
+    T3 = map(Iterators.product(axes(e3.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])
+    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)
+        @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])
+    end
+    C4 = rotate_north(C4, EAST)
+    T4 = map(Iterators.product(axes(e4.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])
+    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] *
-#                 e.corners[NORTHWEST, r, c][1; 2] *
+#                 envfinal.corners[NORTHWEST, r, c][1; 2] *
 #                 conj(signs[NORTH, r, c][-2 2])
-#         elseif dir == 2
-#             return @tensor Cfix[-1; -2] :=
-#                 signs[NORTH, _prev(r, end), _prev(c, end)][-1 1] *
-#                 e.corners[NORTHEAST, r, c][1; 2] *
-#                 conj(signs[EAST, r, c][-2 2])
-#         elseif dir == 3
-#             return @tensor Cfix[-1; -2] :=
-#                 signs[EAST, _prev(r, end), c][-1 1] *
-#                 e.corners[SOUTHEAST, r, c][1; 2] *
-#                 conj(signs[SOUTH, r, c][-2 2])
-#         elseif dir == 4
-#             return @tensor Cfix[-1; -2] :=
-#                 signs[SOUTH, _prev(r, end), c][-1 1] *
-#                 e.corners[SOUTHWEST, r, c][1; 2] *
-#                 conj(signs[WEST, r, c][-2 2])
 #         end
-#     end
-
-#     edgesfix = map(Iterators.product(axes(e.edges)...)) do (dir, r, c)
-#         if dir == 1
-#             return @tensor Tfix[-1 -2 -3; -4] :=
+#         @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] *
-#                 e.edges[NORTH, r, c][1 -2 -3; 2] *
+#                 envfinal.edges[NORTH, r, c][1 -2 -3; 2] *
 #                 conj(signs[NORTH, r, _next(c, end)][-4 2])
-#         elseif dir == 2
-#             return @tensor Tfix[-1 -2 -3; -4] :=
-#                 signs[EAST, r, c][-1 1] *
-#                 e.edges[EAST, r, c][1 -2 -3; 2] *
-#                 conj(signs[EAST, _next(r, end), c][-4 2])
-#         elseif dir == 3
-#             return @tensor Tfix[-1 -2 -3; -4] :=
-#                 signs[SOUTH, r, c][-1 1] *
-#                 e.edges[SOUTH, r, c][1 -2 -3; 2] *
-#                 conj(signs[SOUTH, _next(r, end), _next(c, end)][-4 2])
-#         elseif dir == 4
-#             return @tensor Tfix[-1 -2 -3; -4] :=
-#                 signs[WEST, r, c][-1 1] *
-#                 e.edges[WEST, r, c][1 -2 -3; 2] *
-#                 conj(signs[WEST, _prev(r, end), c][-4 2])
 #         end
-#     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
 
-# 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
-
-
 # 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
@@ -208,14 +225,8 @@ function check_elementwise_conv(
         else
             @warn "no elementwise convergence up to ϵ < $atol:"
             for dir in 1:4
-                println(
-                    "dir=$dir: all |Cⁿ⁺¹ - Cⁿ|ᵢⱼ < ϵ: ",
-                    Cerr[dir],# maximum.(ΔC[dir, :, :])
-                )
-                println(
-                    "dir=$dir: all |Tⁿ⁺¹ - Tⁿ|ᵢⱼ < ϵ: ",
-                    Terr[dir],# maximum.(ΔT[dir, :, :])
-                )
+                println("dir=$dir: all |Cⁿ⁺¹ - Cⁿ|ᵢⱼ < ϵ: ", Cerr[dir])
+                println("dir=$dir: all |Tⁿ⁺¹ - Tⁿ|ᵢⱼ < ϵ: ", Terr[dir])
             end
         end
         println("maxᵢⱼ|Cⁿ⁺¹ - Cⁿ|ᵢⱼ = $ΔCmax")
@@ -374,7 +385,7 @@ function grow_env_left(peps, Pl, Pr, C_sw, C_nw, T_s, T_w, T_n)
 end
 
 # Compute norm of the entire CMTRG enviroment
-function LinearAlgebra.norm(peps, env::CTMRGEnv)
+function LinearAlgebra.norm(peps::InfinitePEPS, env::CTMRGEnv)
     total = one(scalartype(peps))
 
     for r in 1:size(peps, 1), c in 1:size(peps, 2)