Change argument order leading_boundary with CTMRG

examples/heisenberg.jl

@@ -36,6 +36,6 @@ alg = PEPSOptimize(;
 # E/N = −0.6694421, which is a QMC estimate from https://arxiv.org/abs/1101.3281.
 # Of course there is a noticable bias for small χbond and χenv.
 ψ₀ = InfinitePEPS(2, χbond)
-env₀ = leading_boundary(ψ₀, ctmalg, CTMRGEnv(ψ₀; Venv=ℂ^χenv))
+env₀ = leading_boundary(CTMRGEnv(ψ₀; Venv=ℂ^χenv), ψ₀, ctmalg)
 result = fixedpoint(ψ₀, H, alg, env₀)
 @show result.E

src/algorithms/ctmrg.jl

@@ -22,12 +22,15 @@ Regardless of the truncation scheme, the space can be kept fixed with `fixedspac
 end
 
 """
-    MPSKit.leading_boundary(state, alg::CTMRG, [envinit])
+    MPSKit.leading_boundary([envinit], state, alg::CTMRG)
 
 Contract `state` using CTMRG and return the CTM environment.
 Per default, a random initial environment is used.
 """
-function MPSKit.leading_boundary(state, alg::CTMRG, envinit=CTMRGEnv(state))
+function MPSKit.leading_boundary(state, alg::CTMRG)
+    return MPSKit.leading_boundary(CTMRGEnv(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)

src/algorithms/peps_opt.jl

@@ -68,9 +68,9 @@ function fixedpoint(
         E, g = withgradient(peps) do ψ
             envs´ = hook_pullback(
                 leading_boundary,
+                envs,
                 ψ,
-                alg.boundary_alg,
-                envs;
+                alg.boundary_alg;
                 alg_rrule=alg.gradient_alg,
             )
             ignore_derivatives() do
@@ -107,11 +107,11 @@ function _rrule(
     gradmode::Union{GradMode,KrylovKit.LinearSolver},
     ::RuleConfig,
     ::typeof(MPSKit.leading_boundary),
+    envinit,
     state,
     alg::CTMRG,
-    envinit,
 )
-    envs = leading_boundary(state, alg, envinit)
+    envs = leading_boundary(envinit, state, alg)
 
     function leading_boundary_pullback(Δenvs′)
         Δenvs = unthunk(Δenvs′)
@@ -127,7 +127,7 @@ function _rrule(
         ∂f∂x(x)::typeof(envs) = env_vjp(x)[2]
         ∂F∂envs = fpgrad(Δenvs, ∂f∂x, ∂f∂A, Δenvs, gradmode)
 
-        return NoTangent(), ∂F∂envs, NoTangent(), ZeroTangent()
+        return NoTangent(), ZeroTangent(), ∂F∂envs, NoTangent()
     end
 
     return envs, leading_boundary_pullback

test/boundarymps/vumps.jl

@@ -6,17 +6,16 @@ using MPSKit
 using LinearAlgebra
 
 Random.seed!(29384293742893)
-psi = InfinitePEPS(ComplexSpace(2), ComplexSpace(2))
 
+psi = InfinitePEPS(ComplexSpace(2), ComplexSpace(2))
 T = PEPSKit.InfiniteTransferPEPS(psi, 1, 1)
-mps = PEPSKit.initializeMPS(T, [ComplexSpace(20)])
 
+mps = PEPSKit.initializeMPS(T, [ComplexSpace(20)])
 mps, envs, ϵ = leading_boundary(mps, T, VUMPS())
 N = abs(sum(expectation_value(mps, T)))
 
-ctm = leading_boundary(
-    psi, CTMRG(; verbosity=1, fixedspace=true), CTMRGEnv(psi; Venv=ComplexSpace(20))
-)
+ctm = CTMRGEnv(psi; Venv=ComplexSpace(20))
+ctm = leading_boundary(ctm, psi, CTMRG(; verbosity=1, fixedspace=true))
 N2 = abs(norm(psi, ctm))
 
 @test N ≈ N2 atol = 1e-3

test/ctmrg/gaugefix.jl

@@ -50,7 +50,7 @@ end
     )
     alg_fixed = CTMRG(; trscheme=truncdim(dim(ctm_space)), verbosity, fixedspace=true)
 
-    ctm = leading_boundary(psi, alg, ctm)
+    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)
@@ -75,7 +75,7 @@ end
     )
     alg_fixed = CTMRG(; trscheme=truncspace(ctm_space), verbosity, fixedspace=true)
 
-    ctm = leading_boundary(psi, alg, ctm)
+    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/ctmrg/gradients.jl

@@ -38,13 +38,13 @@ gradmodes = [nothing, GeomSum(; tol), ManualIter(; tol), KrylovKit.GMRES(; tol)]
     Random.seed!(42039482038)
     dir = InfinitePEPS(2, χbond)
     psi = InfinitePEPS(2, χbond)
-    env = leading_boundary(psi, boundary_alg, CTMRGEnv(psi; Venv=ℂ^χenv))
+    env = leading_boundary(CTMRGEnv(psi; Venv=ℂ^χenv), psi, boundary_alg)
 
     alphas, fs, dfs1, dfs2 = OptimKit.optimtest(
         (psi, env), dir; alpha=steps, retract=PEPSKit.my_retract, inner=PEPSKit.my_inner
     ) do (peps, envs)
         E, g = Zygote.withgradient(peps) do psi
-            envs2 = PEPSKit.hook_pullback(leading_boundary, psi, boundary_alg, envs; alg_rrule)
+            envs2 = PEPSKit.hook_pullback(leading_boundary, envs, psi, boundary_alg; alg_rrule)
             return costfun(psi, envs2, H)
         end
         return E, only(g)

test/heisenberg.jl

@@ -37,7 +37,7 @@ opt_alg = PEPSOptimize(;
 # initialize states
 Random.seed!(91283219347)
 psi_init = InfinitePEPS(2, χbond)
-env_init = leading_boundary(psi_init, ctm_alg, CTMRGEnv(psi_init; Venv=ComplexSpace(χenv)))
+env_init = leading_boundary(CTMRGEnv(psi_init; Venv=ComplexSpace(χenv)), psi_init, ctm_alg)
 
 # find fixedpoint
 result = fixedpoint(psi_init, H, opt_alg, env_init)