CTMRG support for PEPS-PEPO-PEPS networks

src/algorithms/contractions/localoperator.jl

@@ -262,38 +262,3 @@ end
     end
     return macroexpand(@__MODULE__, returnex)
 end
-
-# Partition function contractions
-
-"""
-    contract_local_tensor(inds, O, env)
-
-Contract a local tensor `O` inserted into a partition function `pf` at position `inds`,
-using the environment `env`.
-"""
-function contract_local_tensor(
-    inds::Tuple{Int,Int},
-    O::AbstractTensorMap{T,S,2,2},
-    env::CTMRGEnv{C,<:CTMRG_PF_EdgeTensor},
-) where {T,S,C}
-    r, c = inds
-    return @autoopt @tensor env.corners[NORTHWEST, _prev(r, end), _prev(c, end)][
-            χ_WNW
-            χ_NNW
-        ] *
-        env.edges[NORTH, _prev(r, end), c][χ_NNW D_N; χ_NNE] *
-        env.corners[NORTHEAST, _prev(r, end), _next(c, end)][χ_NNE; χ_ENE] *
-        env.edges[EAST, r, _next(c, end)][χ_ENE D_E; χ_ESE] *
-        env.corners[SOUTHEAST, _next(r, end), _next(c, end)][χ_ESE; χ_SSE] *
-        env.edges[SOUTH, _next(r, end), c][χ_SSE D_S; χ_SSW] *
-        env.corners[SOUTHWEST, _next(r, end), _prev(c, end)][χ_SSW; χ_WSW] *
-        env.edges[WEST, r, _prev(c, end)][χ_WSW D_W; χ_WNW] *
-        O[D_W D_S; D_N D_E]
-end
-function contract_local_tensor(
-    inds::CartesianIndex{2},
-    O::AbstractTensorMap{T,S,2,2},
-    env::CTMRGEnv{C,<:CTMRG_PF_EdgeTensor},
-) where {T,S,C}
-    return contract_local_tensor(Tuple(inds), O, env)
-end

src/algorithms/contractions/vumps_contractions.jl

@@ -36,95 +36,6 @@
 
 ## PEPO
 
-# some plumbing for generic expressions...
-
-# side=:W for argument, side=:E for output, PEPO height H
-function _pepo_leftenv_expr(envname, side::Symbol, H::Int)
-    return tensorexpr(
-        envname,
-        (
-            envlabel(:S, side),
-            virtuallabel(side, :top),
-            ntuple(i -> virtuallabel(side, :mid, i), H)...,
-            virtuallabel(side, :bot),
-        ),
-        (envlabel(:N, side),),
-    )
-end
-
-# side=:E for argument, side=:W for output, PEPO height H
-function _pepo_rightenv_expr(envname, side::Symbol, H::Int)
-    return tensorexpr(
-        envname,
-        (
-            envlabel(:N, side),
-            virtuallabel(side, :top),
-            ntuple(i -> virtuallabel(side, :mid, i), H)...,
-            virtuallabel(side, :bot),
-        ),
-        (envlabel(:S, side),),
-    )
-end
-
-# side=:N for ket MPS, side=:S for bra MPS, PEPO height H
-function _pepo_mpstensor_expr(tensorname, side::Symbol, H::Int)
-    return tensorexpr(
-        tensorname,
-        (
-            envlabel(side, :W),
-            virtuallabel(side, :top),
-            ntuple(i -> virtuallabel(side, :mid, i), H)...,
-            virtuallabel(side, :bot),
-        ),
-        (envlabel(side, :E),),
-    )
-end
-
-# layer=:top for ket PEPS, layer=:bot for bra PEPS, connects to PEPO slice H
-function _pepo_pepstensor_expr(tensorname, layer::Symbol, h::Int)
-    return tensorexpr(
-        tensorname,
-        (physicallabel(h),),
-        (
-            virtuallabel(:N, layer),
-            virtuallabel(:E, layer),
-            virtuallabel(:S, layer),
-            virtuallabel(:W, layer),
-        ),
-    )
-end
-
-# PEPO slice h
-function _pepo_pepotensor_expr(tensorname, h::Int)
-    return tensorexpr(
-        tensorname,
-        (physicallabel(h + 1), physicallabel(h)),
-        (
-            virtuallabel(:N, :mid, h),
-            virtuallabel(:E, :mid, h),
-            virtuallabel(:S, :mid, h),
-            virtuallabel(:W, :mid, h),
-        ),
-    )
-end
-
-# specialize simple case
-function MPSKit.transfer_left(
-    GL::GenericMPSTensor{S,4},
-    O::PEPOSandwich{1},
-    A::GenericMPSTensor{S,4},
-    Ā::GenericMPSTensor{S,4},
-) where {S}
-    return @autoopt @tensor GL′[χ_SE D_E_above D_E_mid D_E_below; χ_NE] :=
-        GL[χ_SW D_W_above D_W_mid D_W_below; χ_NW] *
-        conj(Ā[χ_SW D_S_above D_S_mid D_S_below; χ_SE]) *
-        ket(O)[d_in; D_N_above D_E_above D_S_above D_W_above] *
-        only(pepo(O))[d_out d_in; D_N_mid D_E_mid D_S_mid D_W_mid] *
-        conj(bra(O)[d_out; D_N_below D_E_below D_S_below D_W_below]) *
-        A[χ_NW D_N_above D_N_mid D_N_below; χ_NE]
-end
-
-# general case
 @generated function MPSKit.transfer_left(
     GL::GenericMPSTensor{S,N},
     O::PEPOSandwich{H},
@@ -134,15 +45,11 @@
     # sanity check
     @assert H == N - 3
 
-    GL´_e = _pepo_leftenv_expr(:GL´, :E, H)
-    GL_e = _pepo_leftenv_expr(:GL, :W, H)
-    A_e = _pepo_mpstensor_expr(:A, :N, H)
-    Ā_e = _pepo_mpstensor_expr(:Ā, :S, H)
-    ket_e = _pepo_pepstensor_expr(:(ket(O)), :top, 1)
-    bra_e = _pepo_pepstensor_expr(:(bra(O)), :bot, H + 1)
-    pepo_es = map(1:H) do h
-        return _pepo_pepotensor_expr(:(pepo(O)[$h]), h)
-    end
+    GL´_e = _pepo_edge_expr(:GL´, :SE, :NE, :E, H)
+    GL_e = _pepo_edge_expr(:GL, :SW, :NW, :W, H)
+    A_e = _pepo_edge_expr(:A, :NW, :NE, :N, H)
+    Ā_e = _pepo_edge_expr(:Ā, :SW, :SE, :S, H)
+    ket_e, bra_e, pepo_es = _pepo_sandwich_expr(:O, H)
 
     rhs = Expr(
         :call,
@@ -158,23 +65,6 @@
     return macroexpand(@__MODULE__, :(return @autoopt @tensor $GL´_e := $rhs))
 end
 
-# specialize simple case
-function MPSKit.transfer_right(
-    GR::GenericMPSTensor{S,4},
-    O::PEPOSandwich{1},
-    A::GenericMPSTensor{S,4},
-    Ā::GenericMPSTensor{S,4},
-) where {S}
-    return @tensor GR′[χ_NW D_W_above D_W_mid D_W_below; χ_SW] :=
-        GR[χ_NE D_E_above D_E_mid D_E_below; χ_SE] *
-        conj(Ā[χ_SW D_S_above D_S_mid D_S_below; χ_SE]) *
-        ket(O)[d_in; D_N_above D_E_above D_S_above D_W_above] *
-        only(pepo(O))[d_out d_in; D_N_mid D_E_mid D_S_mid D_W_mid] *
-        conj(bra(O)[d_out; D_N_below D_E_below D_S_below D_W_below]) *
-        A[χ_NW D_N_above D_N_mid D_N_below; χ_NE]
-end
-
-# general case
 @generated function MPSKit.transfer_right(
     GR::GenericMPSTensor{S,N},
     O::PEPOSandwich{H},
@@ -184,15 +74,11 @@
     # sanity check
     @assert H == N - 3
 
-    GR´_e = _pepo_rightenv_expr(:GR´, :W, H)
-    GR_e = _pepo_rightenv_expr(:GR, :E, H)
-    A_e = _pepo_mpstensor_expr(:A, :N, H)
-    Ā_e = _pepo_mpstensor_expr(:Ā, :S, H)
-    ket_e = _pepo_pepstensor_expr(:(ket(O)), :top, 1)
-    bra_e = _pepo_pepstensor_expr(:(bra(O)), :bot, H + 1)
-    pepo_es = map(1:H) do h
-        return _pepo_pepotensor_expr(:(pepo(O)[$h]), h)
-    end
+    GR´_e = _pepo_edge_expr(:GR´, :NW, :SW, :W, H)
+    GR_e = _pepo_edge_expr(:GR, :NE, :SE, :E, H)
+    A_e = _pepo_edge_expr(:A, :NW, :NE, :N, H)
+    Ā_e = _pepo_edge_expr(:Ā, :SW, :SE, :S, H)
+    ket_e, bra_e, pepo_es = _pepo_sandwich_expr(:O, H)
 
     rhs = Expr(
         :call,
@@ -276,22 +162,6 @@
 
 ## PEPO
 
-# specialize simple case
-function MPSKit.∂AC(
-    AC::GenericMPSTensor{S,4},
-    O::PEPOSandwich{1},
-    GL::GenericMPSTensor{S,4},
-    GR::GenericMPSTensor{S,4},
-) where {S}
-    return @tensor AC′[χ_SW D_S_above D_S_mid D_S_below; χ_SE] :=
-        GL[χ_SW D_W_above D_W_mid D_W_below; χ_NW] *
-        AC[χ_NW D_N_above D_N_mid D_N_below; χ_NE] *
-        GR[χ_NE D_E_above D_E_mid D_E_below; χ_SE] *
-        ket(O)[d_in; D_N_above D_E_above D_S_above D_W_above] *
-        only(pepo(O))[d_out d_in; D_N_mid D_E_mid D_S_mid D_W_mid] *
-        conj(bra(O)[d_out; D_N_below D_E_below D_S_below D_W_below])
-end
-
 @generated function MPSKit.∂AC(
     AC::GenericMPSTensor{S,N},
     O::PEPOSandwich{H},
@@ -301,15 +171,11 @@
     # sanity check
     @assert H == N - 3
 
-    AC´_e = _pepo_mpstensor_expr(:AC´, :S, H)
-    AC_e = _pepo_mpstensor_expr(:AC, :N, H)
-    GL_e = _pepo_leftenv_expr(:GL, :W, H)
-    GR_e = _pepo_rightenv_expr(:GR, :E, H)
-    ket_e = _pepo_pepstensor_expr(:(ket(O)), :top, 1)
-    bra_e = _pepo_pepstensor_expr(:(bra(O)), :bot, H + 1)
-    pepo_es = map(1:H) do h
-        return _pepo_pepotensor_expr(:(pepo(O)[$h]), h)
-    end
+    AC´_e = _pepo_edge_expr(:AC´, :SW, :SE, :S, H)
+    AC_e = _pepo_edge_expr(:AC, :NW, :NE, :N, H)
+    GL_e = _pepo_edge_expr(:GL, :SW, :NW, :W, H)
+    GR_e = _pepo_edge_expr(:GR, :NE, :SE, :E, H)
+    ket_e, bra_e, pepo_es = _pepo_sandwich_expr(:O, H)
 
     rhs = Expr(:call, :*, AC_e, GL_e, GR_e, ket_e, Expr(:call, :conj, bra_e), pepo_es...)
 
@@ -324,23 +190,6 @@
 pepo(p::∂PEPOSandwich) = p[2:end]
 pepo(p::∂PEPOSandwich, i::Int) = p[1 + i]
 
-# specialize simple case
-function ∂peps(
-    AC::GenericMPSTensor{S,4},
-    ĀC::GenericMPSTensor{S,4},
-    O::∂PEPOSandwich{1},
-    GL::GenericMPSTensor{S,4},
-    GR::GenericMPSTensor{S,4},
-) where {S}
-    return @tensor ∂p[d_out; D_N_below D_E_below D_S_below D_W_below] :=
-        GL[χ_SW D_W_above D_W_mid D_W_below; χ_NW] *
-        AC[χ_NW D_N_above D_N_mid D_N_below; χ_NE] *
-        ket(O)[d_in; D_N_above D_E_above D_S_above D_W_above] *
-        only(pepo(O))[d_out d_in; D_N_mid D_E_mid D_S_mid D_W_mid] *
-        GR[χ_NE D_E_above D_E_mid D_E_below; χ_SE] *
-        conj(ĀC[χ_SW D_S_above D_S_mid D_S_below; χ_SE])
-end
-
 @generated function ∂peps(
     AC::GenericMPSTensor{S,N},
     ĀC::GenericMPSTensor{S,N},
@@ -352,14 +201,11 @@
     @assert H == N - 3
 
     ∂p_e = _pepo_pepstensor_expr(:∂p, :bot, H + 1)
-    AC_e = _pepo_mpstensor_expr(:AC, :N, H)
-    ĀC_e = _pepo_mpstensor_expr(:ĀC, :S, H)
-    GL_e = _pepo_leftenv_expr(:GL, :W, H)
-    GR_e = _pepo_rightenv_expr(:GR, :E, H)
-    ket_e = _pepo_pepstensor_expr(:(ket(O)), :top, 1)
-    pepo_es = map(1:H) do h
-        return _pepo_pepotensor_expr(:(pepo(O, $h)), h)
-    end
+    AC_e = _pepo_edge_expr(:AC, :NW, :NE, :N, H)
+    ĀC_e = _pepo_edge_expr(:ĀC, :SW, :SE, :S, H)
+    GL_e = _pepo_edge_expr(:GL, :SW, :NW, :W, H)
+    GR_e = _pepo_edge_expr(:GR, :NE, :SE, :E, H)
+    ket_e, bra_e, pepo_es = _pepo_sandwich_expr(:O, H)
 
     rhs = Expr(:call, :*, AC_e, Expr(:call, :conj, ĀC_e), GL_e, GR_e, ket_e, pepo_es...)
 

src/algorithms/ctmrg/gaugefix.jl

@@ -56,10 +56,23 @@
 
 # this is a bit of a hack to get the fixed point of the mixed transfer matrix
 # because MPSKit is not compatible with AD
+@generated function _transfer_right(
+    v::AbstractTensorMap{<:Any,S,1,N₁},
+    A::GenericMPSTensor{S,N₂},
+    Abar::GenericMPSTensor{S,N₂},
+) where {S,N₁,N₂}
+    t_out = tensorexpr(:v, -1, -(2:(N₁ + 1)))
+    t_top = tensorexpr(:A, (-1, reverse(3:(N₂ + 1))...), 1)
+    t_bot = tensorexpr(:Abar, (-(N₁ + 1), reverse(3:(N₂ + 1))...), 2)
+    t_in = tensorexpr(:v, 1, (-(2:N₁)..., 2))
+    return macroexpand(
+        @__MODULE__, :(return @tensor $t_out := $t_top * conj($t_bot) * $t_in)
+    )
+end
 function transfermatrix_fixedpoint(tops, bottoms, ρinit)
     _, vecs, info = eigsolve(ρinit, 1, :LM, Arnoldi()) do ρ
         return foldr(zip(tops, bottoms); init=ρ) do (top, bottom), ρ
-            return @tensor ρ′[-1; -2] := top[-1 4 3; 1] * conj(bottom[-2 4 3; 2]) * ρ[1; 2]
+            return _transfer_right(ρ, top, bottom)
         end
     end
     info.converged > 0 || @warn "eigsolve did not converge"