Start working on separating contractions

src/PEPSKit.jl

@@ -32,6 +32,7 @@ include("environments/transferpeps_environments.jl")
 include("environments/transferpepo_environments.jl")
 
 include("algorithms/contractions/localoperator.jl")
+include("algorithms/contractions/ctmrg_contractions.jl")
 
 include("algorithms/ctmrg/ctmrg.jl")
 include("algorithms/ctmrg/gaugefix.jl")

src/algorithms/contractions/ctmrg_contractions.jl

@@ -0,0 +1,168 @@
+const CTMRGEdgeTensor{S} = AbstractTensorMap{S,3,1}
+const CTMRGCornerTensor{S} = AbstractTensorMap{S,1,1}
+
+# Enlarged corner contractions
+# ----------------------------
+
+"""
+    enlarge_northwest_corner((row, col), envs, ket, bra)
+    enlarge_northwest_corner(E_west, C_northwest, E_north, ket, bra)
+
+Contract the enlarged northwest corner of the CTMRG environment, either by specifying the
+coordinates, environments and state, or by directly providing the tensors.
+
+```
+    C_northwest -- E_north --
+         |            ||
+      E_west    == ket-bra ==
+         |            ||
+```
+"""
+function enlarge_northwest_corner(
+    (row, col), envs::CTMRGEnv, ket::InfinitePEPS, bra::InfinitePEPS=ket
+)
+    E_west = envs.edges[WEST, row, _prev(col, end)]
+    C_northwest = envs.corners[NORTHWEST, _prev(row, end), _prev(col, end)]
+    E_north = envs.edges[NORTH, _prev(row, end), col]
+    return enlarge_northwest_corner(
+        E_west, C_northwest, E_north, ket[row, col], bra[row, col]
+    )
+end
+function enlarge_northwest_corner(
+    E_west::CTMRGEdgeTensor,
+    C_northwest::CTMRGCornerTensor,
+    E_north::CTMRGEdgeTensor,
+    ket::PEPSTensor,
+    bra::PEPSTensor=ket,
+)
+    return @autoopt @tensor corner[χ_S D_Sabove D_Sbelow; χ_E D_Eabove D_Ebelow] :=
+        E_west[χ_S D1 D2; χ1] *
+        C_northwest[χ1; χ2] *
+        E_north[χ2 D3 D4; χ_E] *
+        ket[d; D3 D_Eabove D_Sabove D1] *
+        conj(bra[d; D4 D_Ebelow D_Sbelow D2])
+end
+
+"""
+    enlarge_northeast_corner((row, col), envs, ket, bra)
+    enlarge_northeast_corner(E_north, C_northeast, E_east, ket, bra)
+
+Contract the enlarged northeast corner of the CTMRG environment, either by specifying the
+coordinates, environments and state, or by directly providing the tensors.
+
+```
+    -- E_north -- C_northeast
+          ||            |
+    == ket-bra ==    E_east
+          ||            |
+```
+"""
+function enlarge_northeast_corner(
+    (row, col), envs::CTMRGEnv, ket::InfinitePEPS, bra::InfinitePEPS=ket
+)
+    E_north = envs.edges[NORTH, _prev(row, end), col]
+    C_northeast = envs.corners[NORTHEAST, _prev(row, end), _next(col, end)]
+    E_east = envs.edges[EAST, row, _next(col, end)]
+    return enlarge_northeast_corner(
+        E_north, C_northeast, E_east, ket[row, col], bra[row, col]
+    )
+end
+function enlarge_northeast_corner(
+    E_north::CTMRGEdgeTensor,
+    C_northeast::CTMRGCornerTensor,
+    E_east::CTMRGEdgeTensor,
+    ket::PEPSTensor,
+    bra::PEPSTensor=ket,
+)
+    return @autoopt @tensor corner[χ_W D_Wabove D_Wbelow; χ_S D_Sabove D_Sbelow] :=
+        E_north[χ_W D1 D2; χ1] *
+        C_northeast[χ1; χ2] *
+        E_east[χ2 D3 D4; χ_S] *
+        ket[d; D1 D3 D_Sabove D_Wabove] *
+        conj(bra[d; D2 D4 D_Sbelow D_Wbelow])
+end
+
+# TODO: also bring other corners in same form
+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), 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
+
+# Projector contractions
+# ----------------------
+
+function halfinfinite_environment(quadrant1::AbstractTensorMap{S,3,3}, quadrant2::AbstractTensorMap{S,3,3})
+    return @autoopt @tensor half[χ_in D_inabove D_inbelow; χ_out D_outabove D_outbelow] :=
+            quadrant1[χ_in D_inabove D_inbelow; χ D1 D2] *
+            quadrant2[χ D1 D2; χ_out D_outabove D_outbelow]
+end
+
+# Renormalization contractions
+# ----------------------------
+
+# corner normalizations are the same contractions everywhere
+function renormalize_corner(quadrant, P_left, P_right)
+    return @autoopt @tensor corner[χ_in; χ_out] :=
+        P_right[χ_in; χ1 D1 D2] * quadrant[χ1 D1 D2; χ2 D3 D4] * P_left[χ2 D3 D4; χ_out]
+end
+function rightrenormalize_corner(C, E, P)
+    return @autoopt @tensor corner[χ_in; χ_out] :=
+        E[χ_in D1 D2; χ1] * C[χ1; χ2] * P[χ2 D1 D2; χ_out]
+end
+function leftrenormalize_corner(C, E, P)
+    return @autoopt @tensor corner[χ_in; χ_out] :=
+        P[χ_in D1 D2; χ1] * C[χ1; χ2] * E[χ2 D1 D2; χ_out]
+end
+
+"""
+    renormalize_west_edge((row, col), envs, P_top, P_bottom, ket, bra)
+    renormalize_west_edge(E_west, P_top, P_bottom, ket, bra)
+
+Absorb a bra-ket pair into the west edge using the given projectors and environment tensors.
+
+```
+    TODO: diagram
+```
+"""
+function renormalize_west_edge((row, col), envs, P_top, P_bottom, ket, bra)
+    return renormalize_west_edge(envs.edges[WEST, row, _prev(col, end)],
+            P_top[WEST, row, col], P_bottom[WEST, _next(row, end), col],
+            ket[row, col], bra[row, col])
+end
+function renormalize_west_edge(E_west, P_top, P_bottom, ket, bra)
+    return @autoopt @tensor edge[χ_S D_Eab D_Ebe; χ_N] :=
+        E_west[χ1 D1 D2; χ2] *
+        ket[d; D3 D_Eab D5 D1] *
+        conj(bra[d; D4 D_Ebe D6 D2]) *
+        P_bottom[χ2 D3 D4; χ_N] *
+        P_top[χ_S; χ1 D5 D6]
+end
+
+# TODO: docstring
+function renormalize_north_edge((row, col), envs, P_left, P_right, ket, bra)
+    return renormalize_north_edge(envs.edges[NORTH, _prev(row, end), col],
+            P_left[NORTH, row, col], P_right[NORTH, row, _prev(col, end)],
+            ket[row, col], bra[row, col])
+end
+function renormalize_north_edge(E_north, P_left, P_right, ket, bra)
+    return @autoopt @tensor edge[χ_W D_Eab D_Ebe; χ_E] :=
+        E_north[χ1 D1 D2; χ2] *
+        ket[d; D3 D_Eab D5 D1] *
+        conj(bra[d; D4 D_Ebe D6 D2]) *
+        P_right[χ2 D3 D4; χ_E] *
+        P_left[χ_W; χ1 D5 D6]
+end
+
+# TODO: add other contractions

src/algorithms/ctmrg/ctmrg.jl

@@ -202,7 +202,7 @@ function ctmrg_expand(state, envs::CTMRGEnv{C,T}, ::SequentialCTMRG) where {C,T}
     directions = collect(Iterators.product(axes(state)...))
     @fwdthreads for (r, c) in directions
         r′ = _next(r, size(state, 1))
-        Q_nw[r, c] = northwest_corner((r, c), envs, state)
+        Q_nw[r, c] = enlarge_northwest_corner((r, c), envs, state)
         Q_sw[r, c] = southwest_corner((r′, c), envs, state)
     end
 
@@ -214,9 +214,9 @@ function ctmrg_expand(state, envs::CTMRGEnv{C,T}, ::SimultaneousCTMRG) where {C,
     drc_combinations = collect(Iterators.product(axes(envs.corners)...))
     @fwdthreads for (dir, r, c) in drc_combinations
         Q[dir, r, c] = if dir == NORTHWEST
-            northwest_corner((r, c), envs, state)
+            enlarge_northwest_corner((r, c), envs, state)
         elseif dir == NORTHEAST
-            northeast_corner((r, c), envs, state)
+            enlarge_northeast_corner((r, c), envs, state)
         elseif dir == SOUTHEAST
             southeast_corner((r, c), envs, state)
         elseif dir == SOUTHWEST
@@ -247,9 +247,7 @@ function ctmrg_projectors(
     directions = collect(Iterators.product(axes(envs.corners, 2), axes(envs.corners, 3)))
     @fwdthreads for (r, c) in directions
         # SVD half-infinite environment
-        @autoopt @tensor QQ[χ_EB D_EBabove D_EBbelow; χ_ET D_ETabove D_ETbelow] :=
-            enlarged_envs[2][r, c][χ_EB D_EBabove D_EBbelow; χ D1 D2] *
-            enlarged_envs[1][r, c][χ D1 D2; χ_ET D_ETabove D_ETbelow]
+        QQ = halfinfinite_environment(enlarged_envs[2][r, c], enlarged_envs[1][r, c])
 
         trscheme = truncation_scheme(projector_alg, envs.edges[WEST, r, c])
         svd_alg = svd_algorithm(projector_alg, (WEST, r, c))
@@ -297,9 +295,9 @@ function ctmrg_projectors(
         end
 
         # SVD half-infinite environment
-        @autoopt @tensor QQ[χ_in D_inabove D_inbelow; χ_out D_outabove D_outbelow] :=
-            enlarged_envs[dir, r, c][χ_in D_inabove D_inbelow; χ D1 D2] *
-            enlarged_envs[_next(dir, 4), next_rc...][χ D1 D2; χ_out D_outabove D_outbelow]
+        QQ = halfinfinite_environment(
+            enlarged_envs[dir, r, c], enlarged_envs[_next(dir, 4), next_rc...]
+        )
 
         trscheme = truncation_scheme(projector_alg, envs.edges[dir, r, c])
         svd_alg = svd_algorithm(projector_alg, (dir, r, c))
@@ -360,6 +358,7 @@ function ctmrg_renormalize(enlarged_envs, projectors, state, envs, ::SequentialC
     @fwdthreads for (r, c) in coordinates
         r′ = _prev(r, size(state, 1))
         c′ = _prev(c, size(state, 2))
+        # TODO: switch this to use contractions
         C_sw, C_nw, T_w = grow_env_left(
             state[r, c],
             projectors[1][r′, c],
@@ -389,6 +388,7 @@ function ctmrg_renormalize(enlarged_envs, projectors, state, envs, ::Simultaneou
         cprev = _prev(c, size(state, 2))
         cnext = _next(c, size(state, 2))
 
+        # TODO: switch this to use contractions
         C_northwest = _contract_new_corner(
             P_right[WEST, rnext, c], enlarged_envs[NORTHWEST, r, c], P_left[NORTH, r, c]
         )
@@ -441,6 +441,8 @@ function ctmrg_renormalize(enlarged_envs, projectors, state, envs, ::Simultaneou
             P_right[WEST, rnext, c][χ_S; χ1 D5 D6]
         edges[WEST, r, c] = E_west / norm(E_west)
     end
+
+    # TODO: here we do not normalize, is this on purpose?
 
     return CTMRGEnv(copy(corners), copy(edges))
 end
@@ -469,40 +471,6 @@ function enlarge_corners_edges(state, env::CTMRGEnv{C,T}) where {C,T}
     return copy(Q)
 end
 
-# 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
-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), 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), 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
-
 function _contract_new_corner(P_right, Q, P_left)
     return @autoopt @tensor corner[χ_in; χ_out] :=
         P_right[χ_in; χ1 D1 D2] * Q[χ1 D1 D2; χ2 D3 D4] * P_left[χ2 D3 D4; χ_out]
@@ -522,6 +490,7 @@ end
 function grow_env_left(
     peps, P_bottom, P_top, corners_SW, corners_NW, edge_S, edge_W, edge_N
 )
+    # TODO: switch to use contractions: renormalize_west_edge, leftrenornalize...
     @autoopt @tensor corner_SW′[χ_E; χ_N] :=
         corners_SW[χ1; χ2] * edge_S[χ_E D1 D2; χ1] * P_bottom[χ2 D1 D2; χ_N]
     @autoopt @tensor corner_NW′[χ_S; χ_E] :=