Increase test coverage for CTMRG contractions (#358)

* Try to trick coverage into covering

* Start testing more CTMRG contractions

* Test, fix, repeat

* Try to avoid LAPACK error

* Don't add unused methods when trying to increase coverage

* Remove accidental `return`s which were skipping some tests

* Remove unused `rrule`s for now

* Implement sparse environment application as proper multiplication, `env * x`

Author: leburgel

src/algorithms/contractions/ctmrg/enlarge_corner.jl

@@ -36,7 +36,7 @@ function enlarge_northwest_corner(
         E_west::CTMRG_PEPS_EdgeTensor, C_northwest::CTMRGCornerTensor,
         E_north::CTMRG_PEPS_EdgeTensor, A::PEPSSandwich,
     )
-    return @tensor begin
+    @tensor begin
         EC[χS DWt DWb; χ2] := E_west[χS DWt DWb; χ1] * C_northwest[χ1; χ2]
         # already putting χE in front here to make next permute cheaper
         ECE[χS χE DWb DNb; DWt DNt] := EC[χS DWt DWb; χ2] * E_north[χ2 DNt DNb; χE]
@@ -45,16 +45,18 @@ function enlarge_northwest_corner(
         corner[χS DSt DSb; χE DEt DEb] :=
             ECEket[χS χE DEt DSt; DWb DNb d] * conj(bra(A)[d; DNb DEb DSb DWb])
     end
+    return corner
 end
 function enlarge_northwest_corner(
         E_west::CTMRG_PF_EdgeTensor, C_northwest::CTMRGCornerTensor,
         E_north::CTMRG_PF_EdgeTensor, A::PFTensor,
     )
-    return @tensor begin
+    @tensor begin
         EC[χ_S DW; χ2] := E_west[χ_S DW; χ1] * C_northwest[χ1; χ2]
         ECE[χ_S χ_E; DW DN] := EC[χ_S DW; χ2] * E_north[χ2 DN; χ_E]
         corner[χ_S D_S; χ_E D_E] := ECE[χ_S χ_E; DW DN] * A[DW D_S; DN D_E]
     end
+    return corner
 end
 
 @generated function enlarge_northwest_corner(
@@ -103,7 +105,7 @@ function enlarge_northeast_corner(
         E_north::CTMRG_PEPS_EdgeTensor, C_northeast::CTMRGCornerTensor,
         E_east::CTMRG_PEPS_EdgeTensor, A::PEPSSandwich,
     )
-    return @tensor begin
+    @tensor begin
         EC[χW DNt DNb; χ2] := E_north[χW DNt DNb; χ1] * C_northeast[χ1; χ2]
         # already putting χE in front here to make next permute cheaper
         ECE[χW χS DNb DEb; DNt DEt] := EC[χW DNt DNb; χ2] * E_east[χ2 DEt DEb; χS]
@@ -112,16 +114,18 @@ function enlarge_northeast_corner(
         corner[χW DWt DWb; χS DSt DSb] :=
             ECEket[χW χS DSt DWt; DNb DEb d] * conj(bra(A)[d; DNb DEb DSb DWb])
     end
+    return corner
 end
 function enlarge_northeast_corner(
         E_north::CTMRG_PF_EdgeTensor, C_northeast::CTMRGCornerTensor,
         E_east::CTMRG_PF_EdgeTensor, A::PFTensor,
     )
-    return @tensor begin
+    @tensor begin
         EC[DN χ_W; χ2] := E_north[χ_W DN; χ1] * C_northeast[χ1; χ2]
         ECE[DN DE; χ_S χ_W] := EC[DN χ_W; χ2] * E_east[χ2 DE; χ_S]
         corner[χ_W D_W; χ_S D_S] := A[D_W D_S; DN DE] * ECE[DN DE; χ_S χ_W]
     end
+    return corner
 end
 
 @generated function enlarge_northeast_corner(
@@ -170,7 +174,7 @@ function enlarge_southeast_corner(
         E_east::CTMRG_PEPS_EdgeTensor, C_southeast::CTMRGCornerTensor,
         E_south::CTMRG_PEPS_EdgeTensor, A::PEPSSandwich,
     )
-    return @tensor begin
+    @tensor begin
         EC[χN DEt DEb; χ2] := E_east[χN DEt DEb; χ1] * C_southeast[χ1; χ2]
         # already putting χE in front here to make next permute cheaper
         ECE[χN χW DEb DSb; DEt DSt] := EC[χN DEt DEb; χ2] * E_south[χ2 DSt DSb; χW]
@@ -179,16 +183,18 @@ function enlarge_southeast_corner(
         corner[χN DNt DNb; χW DWt DWb] :=
             ECEket[χN χW DNt DWt; DEb DSb d] * conj(bra(A)[d; DNb DEb DSb DWb])
     end
+    return corner
 end
 function enlarge_southeast_corner(
         E_east::CTMRG_PF_EdgeTensor, C_southeast::CTMRGCornerTensor,
         E_south::CTMRG_PF_EdgeTensor, A::PFTensor,
     )
-    return @tensor begin
+    @tensor begin
         EC[χ_N D1; χ2] := E_east[χ_N D1; χ1] * C_southeast[χ1; χ2]
         ECE[χ_N χ_W; D1 D2] := EC[χ_N D1; χ2] * E_south[χ2 D2; χ_W]
         corner[χ_N D_N; χ_W D_W] := ECE[χ_N χ_W; D1 D2] * A[D_W D2; D_N D1]
     end
+    return corner
 end
 
 @generated function enlarge_southeast_corner(
@@ -237,7 +243,7 @@ function enlarge_southwest_corner(
         E_south::CTMRG_PEPS_EdgeTensor, C_southwest::CTMRGCornerTensor,
         E_west::CTMRG_PEPS_EdgeTensor, A::PEPSSandwich,
     )
-    return @tensor begin
+    @tensor begin
         EC[χE DSt DSb; χ2] := E_south[χE DSt DSb; χ1] * C_southwest[χ1; χ2]
         # already putting χE in front here to make next permute cheaper
         ECE[χE χN DSb DWb; DSt DWt] := EC[χE DSt DSb; χ2] * E_west[χ2 DWt DWb; χN]
@@ -246,16 +252,18 @@ function enlarge_southwest_corner(
         corner[χE DEt DEb; χN DNt DNb] :=
             ECEket[χE χN DNt DEt; DSb DWb d] * conj(bra(A)[d; DNb DEb DSb DWb])
     end
+    return corner
 end
 function enlarge_southwest_corner(
         E_south::CTMRG_PF_EdgeTensor, C_southwest::CTMRGCornerTensor,
         E_west::CTMRG_PF_EdgeTensor, A::PFTensor,
     )
-    return @tensor begin
+    @tensor begin
         EC[χ_E D1; χ2] := E_south[χ_E D1; χ1] * C_southwest[χ1; χ2]
         ECE[χ_E χ_N; D2 D1] := EC[χ_E D1; χ2] * E_west[χ2 D2; χ_N]
         corner[χ_E D_E; χ_N D_N] := ECE[χ_E χ_N; D2 D1] * A[D2 D1; D_N D_E]
     end
+    return corner
 end
 
 @generated function enlarge_southwest_corner(

src/algorithms/contractions/ctmrg/expr.jl

@@ -73,12 +73,12 @@ function _pepo_edge_expr(edgename, codom_label, dom_label, dir, H::Int, args...)
     return tensorexpr(
         edgename,
         (
-            envlabel(codom_label, args...),
+            envlabel(codom_label),
             virtuallabel(dir, :top, args...),
             ntuple(i -> virtuallabel(dir, :mid, i, args...), H)...,
             virtuallabel(dir, :bot, args...),
         ),
-        (envlabel(dom_label, args...),),
+        (envlabel(dom_label),),
     )
 end
 
@@ -146,9 +146,9 @@ function _pepo_codomain_projector_expr(
     )
     return tensorexpr(
         projname,
-        (envlabel(codom_label, args...),),
+        (envlabel(codom_label),),
         (
-            envlabel(dom_label, args...),
+            envlabel(dom_label),
             virtuallabel(dom_dir, :top, args...),
             ntuple(i -> virtuallabel(dom_dir, :mid, i, args...), H)...,
             virtuallabel(dom_dir, :bot, args...),
@@ -162,11 +162,158 @@ function _pepo_domain_projector_expr(
     return tensorexpr(
         projname,
         (
-            envlabel(codom_label, args...),
+            envlabel(codom_label),
             virtuallabel(codom_dir, :top, args...),
             ntuple(i -> virtuallabel(codom_dir, :mid, i, args...), H)...,
             virtuallabel(codom_dir, :bot, args...),
         ),
-        (envlabel(dom_label, args...),),
+        (envlabel(dom_label),),
+    )
+end
+
+## HalfInfiniteEnv expressions
+
+function _half_infinite_environment_expr_parts(H)
+    # site 1 (codomain)
+    C1_e = _corner_expr(:C_1, :WNW, :NNW)
+    E1_e = _pepo_edge_expr(:E_1, :SW, :WNW, :W, H, 1)
+    E2_e = _pepo_edge_expr(:E_2, :NNW, :NC, :N, H, 1)
+    ket1_e, bra1_e, pepo1_es = _pepo_sandwich_expr(:A_1, H, 1; contract_east = :NC)
+
+    # site 2 (domain)
+    C2_e = _corner_expr(:C_2, :NNE, :ENE)
+    E3_e = _pepo_edge_expr(:E_3, :NC, :NNE, :N, H, 2)
+    E4_e = _pepo_edge_expr(:E_4, :ENE, :SE, :E, H, 2)
+    ket2_e, bra2_e, pepo2_es = _pepo_sandwich_expr(:A_2, H, 2; contract_west = :NC)
+
+    return C1_e, E1_e, E2_e, ket1_e, bra1_e, pepo1_es, C2_e, E3_e, E4_e, ket2_e, bra2_e, pepo2_es
+end
+
+function _half_infinite_environment_expr(H)
+
+    C1_e, E1_e, E2_e, ket1_e, bra1_e, pepo1_es, C2_e, E3_e, E4_e, ket2_e, bra2_e, pepo2_es =
+        _half_infinite_environment_expr_parts(H)
+
+    partial_expr = Expr(
+        :call, :*,
+        E1_e, C1_e, E2_e,
+        ket1_e, Expr(:call, :conj, bra1_e), pepo1_es...,
+        E3_e, C2_e, E4_e,
+        ket2_e, Expr(:call, :conj, bra2_e), pepo2_es...,
+    )
+
+    return partial_expr
+end
+
+function _half_infinite_environment_conj_expr(H)
+
+    C1_e, E1_e, E2_e, ket1_e, bra1_e, pepo1_es, C2_e, E3_e, E4_e, ket2_e, bra2_e, pepo2_es =
+        _half_infinite_environment_expr_parts(H)
+
+    partial_expr = Expr(
+        :call, :*,
+        Expr(:call, :conj, E1_e), Expr(:call, :conj, C1_e), Expr(:call, :conj, E2_e),
+        Expr(:call, :conj, ket1_e), bra1_e, map(x -> Expr(:call, :conj, x), pepo1_es)...,
+        Expr(:call, :conj, E3_e), Expr(:call, :conj, C2_e), Expr(:call, :conj, E4_e),
+        Expr(:call, :conj, ket2_e), bra2_e, map(x -> Expr(:call, :conj, x), pepo2_es)...,
+    )
+
+    return partial_expr
+end
+
+## FullInfiniteEnv expressions
+
+function _full_infinite_environment_expr_parts(H)
+    # site 1 (codomain)
+    C1_e = _corner_expr(:C_1, :WNW, :NNW)
+    E1_e = _pepo_edge_expr(:E_1, :SW, :WNW, :W, H, 1)
+    E2_e = _pepo_edge_expr(:E_2, :NNW, :NC, :N, H, 1)
+    ket1_e, bra1_e, pepo1_es = _pepo_sandwich_expr(:A_1, H, 1; contract_east = :NC)
+
+    # site 2
+    C2_e = _corner_expr(:C_2, :NNE, :ENE)
+    E3_e = _pepo_edge_expr(:E_3, :NC, :NNE, :N, H, 2)
+    E4_e = _pepo_edge_expr(:E_4, :ENE, :EC, :E, H, 2)
+    ket2_e, bra2_e, pepo2_es = _pepo_sandwich_expr(
+        :A_2, H, 2; contract_west = :NC, contract_south = :EC
+    )
+
+    # site 3
+    C3_e = _corner_expr(:C_3, :ESE, :SSE)
+    E5_e = _pepo_edge_expr(:E_5, :EC, :ESE, :E, H, 3)
+    E6_e = _pepo_edge_expr(:E_6, :SSE, :SC, :S, H, 3)
+    ket3_e, bra3_e, pepo3_es = _pepo_sandwich_expr(
+        :A_3, H, 3; contract_north = :EC, contract_west = :SC
+    )
+
+    # site 4 (domain)
+    C4_e = _corner_expr(:C_4, :SSW, :WSW)
+    E7_e = _pepo_edge_expr(:E_7, :SC, :SSW, :S, H, 4)
+    E8_e = _pepo_edge_expr(:E_8, :WSW, :NW, :W, H, 4)
+    ket4_e, bra4_e, pepo4_es = _pepo_sandwich_expr(:A_4, H, 4; contract_east = :SC)
+
+    return (
+        E1_e, C1_e, E2_e,
+        ket1_e, bra1_e, pepo1_es,
+        E3_e, C2_e, E4_e,
+        ket2_e, bra2_e, pepo2_es,
+        E5_e, C3_e, E6_e,
+        ket3_e, bra3_e, pepo3_es,
+        E7_e, C4_e, E8_e,
+        ket4_e, bra4_e, pepo4_es,
     )
 end
+
+function _full_infinite_environment_expr(H)
+    (
+        E1_e, C1_e, E2_e,
+        ket1_e, bra1_e, pepo1_es,
+        E3_e, C2_e, E4_e,
+        ket2_e, bra2_e, pepo2_es,
+        E5_e, C3_e, E6_e,
+        ket3_e, bra3_e, pepo3_es,
+        E7_e, C4_e, E8_e,
+        ket4_e, bra4_e, pepo4_es,
+    ) = _full_infinite_environment_expr_parts(H)
+
+    partial_expr = Expr(
+        :call, :*,
+        E1_e, C1_e, E2_e,
+        ket1_e, Expr(:call, :conj, bra1_e), pepo1_es...,
+        E3_e, C2_e, E4_e,
+        ket2_e, Expr(:call, :conj, bra2_e), pepo2_es...,
+        E5_e, C3_e, E6_e,
+        ket3_e, Expr(:call, :conj, bra3_e), pepo3_es...,
+        E7_e, C4_e, E8_e,
+        ket4_e, Expr(:call, :conj, bra4_e), pepo4_es...,
+    )
+
+    return partial_expr
+end
+
+function _full_infinite_environment_conj_expr(H)
+    (
+        E1_e, C1_e, E2_e,
+        ket1_e, bra1_e, pepo1_es,
+        E3_e, C2_e, E4_e,
+        ket2_e, bra2_e, pepo2_es,
+        E5_e, C3_e, E6_e,
+        ket3_e, bra3_e, pepo3_es,
+        E7_e, C4_e, E8_e,
+        ket4_e, bra4_e, pepo4_es,
+    ) = _full_infinite_environment_expr_parts(H)
+
+    partial_expr = Expr(
+        :call, :*,
+        Expr(:call, :conj, E1_e), Expr(:call, :conj, C1_e), Expr(:call, :conj, E2_e),
+        Expr(:call, :conj, ket1_e), bra1_e, map(x -> Expr(:call, :conj, x), pepo1_es)...,
+        Expr(:call, :conj, E3_e), Expr(:call, :conj, C2_e), Expr(:call, :conj, E4_e),
+        Expr(:call, :conj, ket2_e), bra2_e, map(x -> Expr(:call, :conj, x), pepo2_es)...,
+        Expr(:call, :conj, E5_e), Expr(:call, :conj, C3_e), Expr(:call, :conj, E6_e),
+        Expr(:call, :conj, ket3_e), bra3_e, map(x -> Expr(:call, :conj, x), pepo3_es)...,
+        Expr(:call, :conj, E7_e), Expr(:call, :conj, C4_e), Expr(:call, :conj, E8_e),
+        Expr(:call, :conj, ket4_e), bra4_e, map(x -> Expr(:call, :conj, x), pepo4_es)...,
+    )
+
+    return partial_expr
+end