Refactor enlarged corners again, update Defaults, clean up docstrings

src/PEPSKit.jl

@@ -55,11 +55,17 @@ Module containing default values that represent typical algorithm parameters.
 - `fpgrad_tol = 1e-6`: Convergence tolerance for the fixed-point gradient iteration
 """
 module Defaults
+    using TensorKit, KrylovKit, OptimKit
     const ctmrg_maxiter = 100
     const ctmrg_miniter = 4
     const ctmrg_tol = 1e-10
     const fpgrad_maxiter = 20
     const fpgrad_tol = 1e-6
+    const ctmrgscheme = :simultaneous
+    const iterscheme = :fixed
+    const fwd_alg = TensorKit.SVD()
+    const rrule_alg = GMRES(; tol=ctmrg_tol)
+    const optimizer = LBFGS(10; maxiter=100, gradtol=1e-4, verbosity=2)
 end
 
 export SVDAdjoint, IterSVD, NonTruncSVDAdjoint

src/algorithms/ctmrg.jl

@@ -28,8 +28,8 @@ end
 """
     CTMRG(; tol=Defaults.ctmrg_tol, maxiter=Defaults.ctmrg_maxiter,
           miniter=Defaults.ctmrg_miniter, verbosity=0,
-          svd_alg=TensorKit.SVD(), trscheme=FixedSpaceTruncation(),
-          ctmrgscheme=:simultaneous)
+          svd_alg=SVDAdjoint(), trscheme=FixedSpaceTruncation(),
+          ctmrgscheme=Defaults.ctmrgscheme)
 
 Algorithm struct that represents the CTMRG algorithm for contracting infinite PEPS.
 Each CTMRG run is converged up to `tol` where the singular value convergence of the
@@ -60,7 +60,7 @@ function CTMRG(;
     verbosity=1,
     svd_alg=SVDAdjoint(),
     trscheme=FixedSpaceTruncation(),
-    ctmrgscheme=:simultaneous,
+    ctmrgscheme=Defaults.ctmrgscheme,
 )
     return CTMRG{ctmrgscheme}(
         tol, maxiter, miniter, verbosity, ProjectorAlg(; svd_alg, trscheme, verbosity)
@@ -184,8 +184,8 @@ function left_move(state, env::CTMRGEnv{C,T}, alg::ProjectorAlg) where {C,T}
         # Compute projectors
         for row in 1:size(state, 1)
             # Enlarged corners
-            Q_sw = southwest_corner((_next(row, size(state, 1)), col), state, env)
-            Q_nw = northwest_corner((row, col), state, env)
+            Q_sw = southwest_corner((_next(row, size(state, 1)), col), env, state)
+            Q_nw = northwest_corner((row, col), env, state)
 
             # SVD half-infinite environment
             trscheme = if alg.trscheme isa FixedSpaceTruncation
@@ -236,48 +236,38 @@ function left_move(state, env::CTMRGEnv{C,T}, alg::ProjectorAlg) where {C,T}
     return CTMRGEnv(corners, edges), (; P_left=copy(P_top), P_right=copy(P_bottom), ϵ)
 end
 
-# Generic enlarged corner contraction
-function enlarged_corner(edge_in, corner, edge_out, peps_above, peps_below=peps_above)
-    return @autoopt @tensor Q[χ_in D_inabove D_inbelow; χ_out D_outabove D_outbelow] :=
-        edge_in[χ_in D1 D2; χ1] *
-        corner[χ1; χ2] *
-        edge_out[χ2 D3 D4; χ_out] *
-        peps_above[d; D3 D_outabove D_inabove D1] *
-        conj(peps_below[d; D4 D_outbelow D_inbelow D2])
+# Enlarged corner contractions (need direction specific methods to avoid PEPS rotations)
+function northwest_corner((row, col), env, peps_above, peps_below=peps_above)
+    return @autoopt @tensor corner[χ_S D_Sabove D_Sbelow; χ_E D_Eabove D_Ebelow] :=
+        env.edges[WEST, row, _prev(col, end)][χ_S D1 D2; χ1] *
+        env.corners[NORTHWEST, _prev(row, end), _prev(col, end)][χ1; χ2] *
+        env.edges[NORTH, _prev(row, end), col][χ2 D3 D4; χ_E] *
+        peps_above[row, col][d; D3 D_Eabove D_Sabove D1] *
+        conj(peps_below[row, col][d; D4 D_Ebelow D_Sbelow D2])
 end
-
-# Direction specific methods
-function northwest_corner((row, col), state, env)
-    return enlarged_corner(
-        env.edges[WEST, row, _prev(col, end)],
-        env.corners[NORTHWEST, _prev(row, end), _prev(col, end)],
-        env.edges[NORTH, _prev(row, end), col],
-        state[row, col],
-    )
-end
-function northeast_corner((row, col), state, env)
-    return enlarged_corner(
-        env.edges[NORTH, _prev(row, end), col],
-        env.corners[NORTHEAST, _prev(row, end), _next(col, end)],
-        env.edges[EAST, row, _next(col, end)],
-        rotate_north(state[row, col], EAST),
-    )
+function northeast_corner((row, col), env, peps_above, peps_below=peps_above)
+    return @autoopt @tensor corner[χ_W D_Wabove D_Wbelow; χ_S D_Sabove D_Sbelow] :=
+        env.edges[NORTH, _prev(row, end), col][χ_W D1 D2; χ1] *
+        env.corners[NORTHEAST, _prev(row, end), _next(col, end)][χ1; χ2] *
+        env.edges[EAST, row, _next(col, end)][χ2 D3 D4; χ_S] *
+        peps_above[row, col][d; D1 D3 D_Sabove D_Wabove] *
+        conj(peps_below[row, col][d; D2 D4 D_Sbelow D_Wbelow])
 end
-function southeast_corner((row, col), state, env)
-    return enlarged_corner(
-        env.edges[EAST, row, _next(col, end)],
-        env.corners[SOUTHEAST, _next(row, end), _next(col, end)],
-        env.edges[SOUTH, _next(row, end), col],
-        rotate_north(state[row, col], SOUTH),
-    )
+function southeast_corner((row, col), env, peps_above, peps_below=peps_above)
+    return @autoopt @tensor corner[χ_N D_Nabove D_Nbelow; χ_W D_Wabove D_Wbelow] :=
+        env.edges[EAST, row, _next(col, end)][χ_N D1 D2; χ1] *
+        env.corners[SOUTHEAST, _next(row, end), _next(col, end)][χ1; χ2] *
+        env.edges[SOUTH, _next(row, end), col][χ2 D3 D4; χ_W] *
+        peps_above[row, col][d; D_Nabove D1 D3 D_Wabove] *
+        conj(peps_below[row, col][d; D_Nbelow D2 D4 D_Wbelow])
 end
-function southwest_corner((row, col), state, env)
-    return enlarged_corner(
-        env.edges[SOUTH, _next(row, end), col],
-        env.corners[SOUTHWEST, _next(row, end), _prev(col, end)],
-        env.edges[WEST, row, _prev(col, end)],
-        rotate_north(state[row, col], WEST),
-    )
+function southwest_corner((row, col), env, peps_above, peps_below=peps_above)
+    return @autoopt @tensor corner[χ_E D_Eabove D_Ebelow; χ_N D_Nabove D_Nbelow] :=
+        env.edges[SOUTH, _next(row, end), col][χ_E D1 D2; χ1] *
+        env.corners[SOUTHWEST, _next(row, end), _prev(col, end)][χ1; χ2] *
+        env.edges[WEST, row, _prev(col, end)][χ2 D3 D4; χ_N] *
+        peps_above[row, col][d; D_Nabove D_Eabove D1 D3] *
+        conj(peps_below[row, col][d; D_Nbelow D_Ebelow D2 D4])
 end
 
 # Build projectors from SVD and enlarged SW & NW corners

src/algorithms/ctmrg_all_sides.jl

@@ -19,13 +19,13 @@ function enlarge_corners_edges(state, env::CTMRGEnv{C,T}) where {C,T}
     drc_combinations = collect(Iterators.product(axes(env.corners)...))
     @fwdthreads for (dir, r, c) in drc_combinations
         Q[dir, r, c] = if dir == NORTHWEST
-            northwest_corner((r, c), state, env)
+            northwest_corner((r, c), env, state)
         elseif dir == NORTHEAST
-            northeast_corner((r, c), state, env)
-        elseif dir == SOUTHEAST
-            southeast_corner((r, c), state, env)
-        elseif dir == SOUTHWEST
-            southwest_corner((r, c), state, env)
+            northeast_corner((r, c), env, state)
+        elseif dir == SOUTHEAST         
+            southeast_corner((r, c), env, state)
+        elseif dir == SOUTHWEST         
+            southwest_corner((r, c), env, state)
         end
     end
 

src/algorithms/peps_opt.jl

@@ -2,7 +2,7 @@ abstract type GradMode{F} end
 
 """
     struct GeomSum(; maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol,
-                   verbosity=0, iterscheme=:fixed) <: GradMode{iterscheme}
+                   verbosity=0, iterscheme=Defaults.iterscheme) <: GradMode{iterscheme}
 
 Gradient mode for CTMRG using explicit evaluation of the geometric sum.
 
@@ -18,14 +18,17 @@ struct GeomSum{F} <: GradMode{F}
     verbosity::Int
 end
 function GeomSum(;
-    maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol, verbosity=0, iterscheme=:fixed
+    maxiter=Defaults.fpgrad_maxiter,
+    tol=Defaults.fpgrad_tol,
+    verbosity=0,
+    iterscheme=Defaults.iterscheme,
 )
     return GeomSum{iterscheme}(maxiter, tol, verbosity)
 end
 
 """
     struct ManualIter(; maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol,
-                      verbosity=0, iterscheme=:fixed) <: GradMode{iterscheme}
+                      verbosity=0, iterscheme=Defaults.iterscheme) <: GradMode{iterscheme}
 
 Gradient mode for CTMRG using manual iteration to solve the linear problem.
 
@@ -41,13 +44,16 @@ struct ManualIter{F} <: GradMode{F}
     verbosity::Int
 end
 function ManualIter(;
-    maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol, verbosity=0, iterscheme=:fixed
+    maxiter=Defaults.fpgrad_maxiter,
+    tol=Defaults.fpgrad_tol,
+    verbosity=0,
+    iterscheme=Defaults.iterscheme,
 )
     return ManualIter{iterscheme}(maxiter, tol, verbosity)
 end
 
 """
-    struct LinSolver(; solver=KrylovKit.GMRES(), iterscheme=:fixed) <: GradMode{iterscheme}
+    struct LinSolver(; solver=KrylovKit.GMRES(), iterscheme=Defaults.iterscheme) <: GradMode{iterscheme}
 
 Gradient mode wrapper around `KrylovKit.LinearSolver` for solving the gradient linear
 problem using iterative solvers.
@@ -63,7 +69,7 @@ struct LinSolver{F} <: GradMode{F}
 end
 function LinSolver(;
     solver=KrylovKit.GMRES(; maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol),
-    iterscheme=:fixed,
+    iterscheme=Defaults.iterscheme,
 )
     return LinSolver{iterscheme}(solver)
 end
@@ -105,7 +111,7 @@ struct PEPSOptimize{G}
 end
 function PEPSOptimize(;
     boundary_alg=CTMRG(),
-    optimizer=LBFGS(4; maxiter=100, gradtol=1e-4, verbosity=2),
+    optimizer=Defaults.optimizer,
     reuse_env=true,
     gradient_alg=LinSolver(),
     verbosity=0,
@@ -204,6 +210,7 @@ function _rrule(
     alg::CTMRG{C},
 ) where {C}
     @assert C == :simultaneous
+    @assert alg.projector_alg.svd_alg.rrule_alg isa Union{KrylovKit.LinearSolver,Arnoldi}
     envs = leading_boundary(envinit, state, alg)
     envsconv, info = ctmrg_iter(state, envs, alg)
     envsfix, signs = gauge_fix(envs, envsconv)

src/utility/svd.jl

@@ -10,7 +10,7 @@ using TensorKit:
 const CRCExt = Base.get_extension(KrylovKit, :KrylovKitChainRulesCoreExt)
 
 """
-    struct SVDAdjoint(; fwd_alg = TensorKit.SVD(), rrule_alg = nothing,
+    struct SVDAdjoint(; fwd_alg = Defaults.fwd_alg, rrule_alg = Defaults.rrule_alg,
                       broadening = nothing)
 
 Wrapper for a SVD algorithm `fwd_alg` with a defined reverse rule `rrule_alg`.
@@ -19,8 +19,8 @@ In case of degenerate singular values, one might need a `broadening` scheme whic
 removes the divergences from the adjoint.
 """
 @kwdef struct SVDAdjoint{F,R,B}
-    fwd_alg::F = TensorKit.SVD()
-    rrule_alg::R = nothing
+    fwd_alg::F = Defaults.fwd_alg
+    rrule_alg::R = Defaults.rrule_alg
     broadening::B = nothing
 end  # Keep truncation algorithm separate to be able to specify CTMRG dependent information
 
@@ -149,9 +149,10 @@ function ChainRulesCore.rrule(
             n_vals = length(Sdc)
             lvecs = Vector{Vector{scalartype(t)}}(eachcol(Uc))
             rvecs = Vector{Vector{scalartype(t)}}(eachcol(Vc'))
-            minimal_info = KrylovKit.ConvergenceInfo(length(Sdc), nothing, nothing, -1, -1)  # Only num. converged is used
-            minimal_alg = GKL(; tol=1e-6)  # Only tolerance is used for gauge sensitivity
-            # TODO: How do we not hard-code this tolerance?
+
+            # Dummy objects only used for warnings
+            minimal_info = KrylovKit.ConvergenceInfo(n_vals, nothing, nothing, -1, -1)  # Only num. converged is used
+            minimal_alg = GKL(; tol=1e-6)  # Only tolerance is used for gauge sensitivity (# TODO: How do we not hard-code this tolerance?)
 
             if ΔUc isa AbstractZero && ΔVc isa AbstractZero  # Handle ZeroTangent singular vectors
                 Δlvecs = fill(ZeroTangent(), n_vals)

test/ctmrg/fixed_iterscheme.jl

@@ -48,31 +48,31 @@ end
 
 # TODO: Why doesn't fixed work with IterSVD?
 ##
-# ctm_alg = CTMRG(;
-#     tol=1e-12,
-#     miniter=4,
-#     maxiter=100,
-#     verbosity=1,
-#     ctmrgscheme=:simultaneous,
-#     svd_alg=SVDAdjoint(; fwd_alg=IterSVD()),
-# )
+ctm_alg = CTMRG(;
+    tol=1e-12,
+    miniter=4,
+    maxiter=100,
+    verbosity=1,
+    ctmrgscheme=:simultaneous,
+    svd_alg=SVDAdjoint(; fwd_alg=IterSVD()),
+)
 
-# # initialize states
-# Random.seed!(91283219347)
-# psi = InfinitePEPS(2, χbond)
-# env_conv1 = leading_boundary(CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg);
+# initialize states
+Random.seed!(91283219347)
+psi = InfinitePEPS(2, χbond)
+env_conv1 = leading_boundary(CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg);
 
 # # do extra iteration to get SVD
-# env_conv2, info = ctmrg_iter(psi, env_conv1, ctm_alg);
-# env_fix, signs = gauge_fix(env_conv1, env_conv2);
-# @test check_elementwise_convergence(env_conv1, env_fix)
+env_conv2, info = ctmrg_iter(psi, env_conv1, ctm_alg);
+env_fix, signs = gauge_fix(env_conv1, env_conv2);
+@test check_elementwise_convergence(env_conv1, env_fix)
 
-# # fix gauge of SVD
-# U_fix, V_fix = fix_relative_phases(info.U, info.V, signs);
-# svd_alg_fix = SVDAdjoint(; fwd_alg=FixedSVD(U_fix, info.S, V_fix));
-# ctm_alg_fix = CTMRG(; svd_alg=svd_alg_fix, trscheme=notrunc(), ctmrgscheme=:simultaneous);
+# fix gauge of SVD
+U_fix, V_fix = fix_relative_phases(info.U, info.V, signs);
+svd_alg_fix = SVDAdjoint(; fwd_alg=FixedSVD(U_fix, info.S, V_fix));
+ctm_alg_fix = CTMRG(; svd_alg=svd_alg_fix, trscheme=notrunc(), ctmrgscheme=:simultaneous);
 
-# # do iteration with FixedSVD
-# env_fixedsvd, = ctmrg_iter(psi, env_conv1, ctm_alg_fix);
-# env_fixedsvd = fix_global_phases(env_conv1, env_fixedsvd);
-# @test check_elementwise_convergence(env_conv1, env_fixedsvd)
+# do iteration with FixedSVD
+env_fixedsvd, = ctmrg_iter(psi, env_conv1, ctm_alg_fix);
+env_fixedsvd = fix_global_phases(env_conv1, env_fixedsvd);
+@test check_elementwise_convergence(env_conv1, env_fixedsvd)

test/heisenberg.jl

@@ -13,8 +13,7 @@ ctm_alg = CTMRG(;
     miniter=4,
     maxiter=100,
     verbosity=1,
-    trscheme=truncdim(χenv),
-    svd_alg=SVDAdjoint(; fwd_alg=TensorKit.SVD(), rrule_alg=GMRES(; tol=1e-10)),
+    svd_alg=SVDAdjoint(; fwd_alg=TensorKit.SVD(), rrule_alg=GMRES(tol=1e-10)),
     ctmrgscheme=:simultaneous,
 )
 opt_alg = PEPSOptimize(;