Implement sparse SVD with function handles

src/PEPSKit.jl

@@ -37,6 +37,7 @@ include("algorithms/contractions/localoperator.jl")
 include("algorithms/contractions/ctmrg_contractions.jl")
 
 include("algorithms/ctmrg/ctmrg.jl")
+include("algorithms/ctmrg/sparse_environments.jl")
 include("algorithms/ctmrg/gaugefix.jl")
 
 include("algorithms/toolbox.jl")
@@ -71,6 +72,7 @@ module Defaults
     const fpgrad_maxiter = 30
     const fpgrad_tol = 1e-6
     const ctmrgscheme = :simultaneous
+    const sparse_env = false
     const reuse_env = true
     const trscheme = FixedSpaceTruncation()
     const iterscheme = :fixed

src/algorithms/contractions/ctmrg_contractions.jl

@@ -163,8 +163,58 @@ end
 # Projector contractions
 # ----------------------
 
+"""
+    left_projector(E_1, C, E_2, V, isqS, ket::PEPSTensor, bra::PEPSTensor=ket)
+
+Contract the CTMRG left projector with the higher-dimensional subspace facing to the left.
+
+```
+     C  --  E_2    -- |~~|
+     |       ||       |V'| -- isqS --
+    E_1 == ket-bra == |~~|
+     |       ||
+```
+"""
+function left_projector(E_1, C, E_2, V, isqS, ket::PEPSTensor, bra::PEPSTensor=ket)
+    return @autoopt @tensor P_left[χ_in D_inabove D_inbelow; χ_out] :=
+        E_1[χ_in D1 D2; χ1] *
+        C[χ1; χ2] *
+        E_2[χ2 D3 D4; χ3] *
+        ket[d; D3 D5 D_inabove D1] *
+        conj(bra[d; D4 D6 D_inbelow D2]) *
+        conj(V[χ4; χ3 D5 D6]) *
+        isqS[χ4; χ_out]
+end
+
+"""
+    right_projector(E_1, C, E_2, U, isqS, ket::PEPSTensor, bra::PEPSTensor=ket)
+
+Contract the CTMRG right projector with the higher-dimensional subspace facing to the right.
+
+```
+               |~~| --   E_2   --  C
+    -- isqS -- |U'|      ||        |
+               |~~| == ket-bra == E_1
+                         ||        |
+```
+"""
+function right_projector(E_1, C, E_2, U, isqS, ket::PEPSTensor, bra::PEPSTensor=ket)
+    return @autoopt @tensor P_right[χ_in; χ_out D_outabove D_outbelow] :=
+        isqS[χ_in; χ1] *
+        conj(U[χ1; χ2 D1 D2]) *
+        ket[d; D1 D5 D_outabove D1] *
+        conj(bra[d; D2 D6 D_outbelow D2]) *
+        E_2[χ2 D3 D4; χ3] *
+        C[χ3; χ4] *
+        E_1[χ4 D5 D6; χ_out]
+end
+
 """
     halfinfinite_environment(quadrant1::AbstractTensorMap{S,3,3}, quadrant2::AbstractTensorMap{S,3,3})
+    function halfinfinite_environment(C_1, C_2, E_1, E_2, E_3, E_4,
+                                      ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2) where {P<:PEPSTensor}
+    function halfinfinite_environment(C_1, C_2, E_1, E_2, E_3, E_4, x,
+                                      ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2) where {P<:PEPSTensor}
 
 Contract two quadrants (enlarged corners) to form a half-infinite environment.
 
@@ -174,22 +224,110 @@ Contract two quadrants (enlarged corners) to form a half-infinite environment.
     |~~~~~~~~~| == |~~~~~~~~~|
       |    ||        ||    |
 ```
+
+The environment can also be contracted directly from all its constituent tensors.
+
+```
+    C_1 --  E_2      --  E_3      -- C_2
+     |       ||          ||           | 
+    E_1 == ket_bra_1 == ket_bra_2 == E_4
+     |       ||          ||           |
+```
+
+Alternatively, contract environment with a vector `x` acting on it, e.g. as needed for iterative solvers.
+
+```
+    C_1 --  E_2      --  E_3      -- C_2
+     |       ||          ||           | 
+    E_1 == ket_bra_1 == ket_bra_2 == E_4
+     |       ||          ||           |
+                         [~~~~~~x~~~~~~]
+                         ||           |
+```
 """
 function halfinfinite_environment(
     quadrant1::AbstractTensorMap{S,3,3}, quadrant2::AbstractTensorMap{S,3,3}
 ) where {S}
-    return @autoopt @tensor half[χ_in D_inabove D_inbelow; χ_out D_outabove D_outbelow] :=
+    return @autoopt @tensor env[χ_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
+function halfinfinite_environment(
+    C_1, C_2, E_1, E_2, E_3, E_4, ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2
+) where {P<:PEPSTensor}
+    return @autoopt @tensor env[χ_in D_inabove D_inbelow; χ_out D_outabove D_outbelow] :=
+        E_1[χ_in D1 D2; χ1] *
+        C_1[χ1; χ2] *
+        E_2[χ2 D3 D4; χ3] *
+        ket_1[d1; D3 D9 D_inabove D1] *
+        conj(bra_1[d1; D4 D10 D_inbelow D2]) *
+        ket_2[d2; D5 D7 D_outabove D9] *
+        conj(bra_2[d2; D6 D8 D_outbelow D10]) *
+        E_3[χ3 D5 D6; χ4] *
+        C_2[χ4; χ5] *
+        E_4[χ5 D7 D8; χ_out]
+end
+function halfinfinite_environment(
+    C_1,
+    C_2,
+    E_1,
+    E_2,
+    E_3,
+    E_4,
+    x::AbstractTensor{S,3},
+    ket_1::P,
+    ket_2::P,
+    bra_1::P=ket_1,
+    bra_2::P=ket_2,
+) where {S,P<:PEPSTensor}
+    return @autoopt @tensor env_x[χ_in D_inabove D_inbelow] :=
+        E_1[χ_in D1 D2; χ1] *
+        C_1[χ1; χ2] *
+        E_2[χ2 D3 D4; χ3] *
+        ket_1[d1; D3 D9 D_inabove D1] *
+        conj(bra_1[d1; D4 D10 D_inbelow D2]) *
+        ket_2[d2; D5 D7 D11 D9] *
+        conj(bra_2[d2; D6 D8 D12 D10]) *
+        E_3[χ3 D5 D6; χ4] *
+        C_2[χ4; χ5] *
+        E_4[χ5 D7 D8; χ6] *
+        x[χ6 D11 D12]
+end
+function halfinfinite_environment(
+    x::AbstractTensor{S,3},
+    C_1,
+    C_2,
+    E_1,
+    E_2,
+    E_3,
+    E_4,
+    ket_1::P,
+    ket_2::P,
+    bra_1::P=ket_1,
+    bra_2::P=ket_2,
+) where {S,P<:PEPSTensor}
+    return @autoopt @tensor env_x[χ_in D_inabove D_inbelow] :=
+        x[χ1 D1 D2] *
+        conj(E_1[χ1 D3 D4; χ2]) *
+        conj(C_1[χ2; χ3]) *
+        conj(E_2[χ3 D5 D6; χ4]) *
+        conj(ket_1[d1; D5 D11 D1 D3]) *
+        bra_1[d1; D6 D12 D2 D4] *
+        conj(ket_2[d2; D7 D9 D_inabove D11]) *
+        bra_2[d2; D8 D10 D_inbelow D12] *
+        conj(E_3[χ4 D7 D8; χ5]) *
+        conj(C_2[χ5; χ6]) *
+        conj(E_4[χ6 D9 D10; χ_in])
+end
 
 # Renormalization contractions
 # ----------------------------
 
 # corners
 
 """
-    renormalize_corner(quadrant, P_left, P_right)
+    renormalize_corner(quadrant::AbstractTensorMap{S,3,3}, P_left, P_right)
+    renormalize_corner(E_1, C, E_2, P_left, P_right, ket::PEPSTensor, bra::PEPSTensor=ket)
 
 Apply projectors to each side of a quadrant.
 
@@ -201,11 +339,34 @@ Apply projectors to each side of a quadrant.
     [P_right]
         |
 ```
+
+Alternatively, provide the constituent tensors and perform the complete contraction.
+
+```
+     C  --   E_2   -- |~~~~~~|
+     |        |       |P_left| --
+    E_1 == ket-bra == |~~~~~~|
+     |        ||
+    [~~P_right~~]
+         |
+```
 """
 function renormalize_corner(quadrant::AbstractTensorMap{S,3,3}, P_left, P_right) where {S}
     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 renormalize_corner(
+    E_1, C, E_2, P_left, P_right, ket::PEPSTensor, bra::PEPSTensor=ket
+)
+    return @autoopt @tensor corner[χ_in; χ_out] :=
+        P_right[χ_in; χ1 D1 D2] *
+        E_1[χ1 D3 D4; χ2] *
+        C[χ2; χ3] *
+        E_2[χ3 D5 D6; χ4] *
+        ket[d; D5 D7 D1 D3] *
+        conj(bra[d; D6 D8 D2 D4]) *
+        P_left[χ4 D7 D8; χ_out]
+end
 
 """
     renormalize_northwest_corner((row, col), enlarged_envs::CTMRGEnv, P_left, P_right)

src/algorithms/ctmrg/ctmrg.jl

@@ -46,7 +46,7 @@ end
     CTMRG(; tol=Defaults.ctmrg_tol, maxiter=Defaults.ctmrg_maxiter,
           miniter=Defaults.ctmrg_miniter, verbosity=0,
           svd_alg=SVDAdjoint(), trscheme=FixedSpaceTruncation(),
-          ctmrgscheme=Defaults.ctmrgscheme)
+          ctmrgscheme=Defaults.ctmrgscheme, sparse_env=Defaults.sparse_env)
 
 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
@@ -63,8 +63,12 @@ CTMRG is implemented. It can either be `:sequential`, where the projectors are s
 computed on the western side, and then applied and rotated. Or with `:simultaneous` all projectors
 are computed and applied simultaneously on all sides, where in particular the corners get
 contracted with two projectors at the same time.
+
+The contraction and SVD of the CTMRG environments can be performed densely, or sparsely
+if `sparse_env` is `true`. If sparsity is enabled, then the SVD will use only function handles
+and the full enlarged corners and environments are never explicitly constructed.
 """
-struct CTMRG{S}
+struct CTMRG{M,S}
     tol::Float64
     maxiter::Int
     miniter::Int
@@ -79,13 +83,15 @@ function CTMRG(;
     svd_alg=Defaults.svd_alg,
     trscheme=Defaults.trscheme,
     ctmrgscheme::Symbol=Defaults.ctmrgscheme,
+    sparse_env::Bool=Defaults.sparse_env,
 )
-    return CTMRG{ctmrgscheme}(
+    return CTMRG{ctmrgscheme,sparse_env}(
         tol, maxiter, miniter, verbosity, ProjectorAlg(; svd_alg, trscheme, verbosity)
     )
 end
 
-ctmrgscheme(::CTMRG{S}) where {S} = S
+ctmrgscheme(::CTMRG{M,S}) where {M,S} = M
+sparse_env(::CTMRG{M,S}) where {M,S} = S
 
 # aliases for the different CTMRG schemes
 const SequentialCTMRG = CTMRG{:sequential}
@@ -183,45 +189,28 @@ There are two modes of expansion: `M = :sequential` and `M = :simultaneous`.
 The first mode expands the environment in one direction at a time, for convenience towards
 the left. The second mode expands the environment in all four directions simultaneously.
 """
-function ctmrg_expand(state, envs::CTMRGEnv{C,T}, ::SequentialCTMRG) where {C,T}
-    Qtype = tensormaptype(spacetype(C), 3, 3, storagetype(C))
-    Q_sw = Zygote.Buffer(envs.corners, Qtype, axes(state)...)
-    Q_nw = Zygote.Buffer(envs.corners, Qtype, axes(state)...)
-
-    directions = collect(Iterators.product(axes(state)...))
-    # @fwdthreads for (r, c) in directions
-    for (r, c) in directions
-        Q_sw[r, c] = enlarge_southwest_corner((r, c), envs, state)
-        Q_nw[r, c] = enlarge_northwest_corner((r, c), envs, state)
+function ctmrg_expand(state, envs::CTMRGEnv, alg::SequentialCTMRG)
+    drc_combinations = collect(Iterators.product([4, 1], axes(state)...))
+    return map(drc_combinations) do (dir, r, c)
+        enlarge_corner(Val(dir), (r, c), envs, state, alg)
     end
-
-    return copy(Q_sw), copy(Q_nw)
 end
-function ctmrg_expand(state, envs::CTMRGEnv{C,T}, ::SimultaneousCTMRG) where {C,T}
-    Qtype = tensormaptype(spacetype(C), 3, 3, storagetype(C))
-    Q = Zygote.Buffer(Array{Qtype,3}(undef, size(envs.corners)))
-    drc_combinations = collect(Iterators.product(axes(envs.corners)...))
-    @fwdthreads for (dir, r, c) in drc_combinations
-        Q[dir, r, c] = if dir == NORTHWEST
-            enlarge_northwest_corner((r, c), envs, state)
-        elseif dir == NORTHEAST
-            enlarge_northeast_corner((r, c), envs, state)
-        elseif dir == SOUTHEAST
-            enlarge_southeast_corner((r, c), envs, state)
-        elseif dir == SOUTHWEST
-            enlarge_southwest_corner((r, c), envs, state)
-        end
+function ctmrg_expand(state, envs::CTMRGEnv, alg::SimultaneousCTMRG)
+    drc_combinations = collect(Iterators.product(1:4, axes(state)...))
+    return map(drc_combinations) do (dir, r, c)
+        enlarge_corner(Val(dir), (r, c), envs, state, alg)
     end
-
-    return copy(Q)
 end
 
 # ======================================================================================== #
 # Projector step
 # ======================================================================================== #
 
 """
-    ctmrg_projectors(Q, env, alg::CTMRG{M})
+    ctmrg_projectors(enlarged_envs, env, alg::CTMRG{M})
+
+Compute the CTMRG projectors based from enlarged environments.
+In the `:simultaneous` mode, the environment SVD is run in parallel.
 """
 function ctmrg_projectors(
     enlarged_envs, envs::CTMRGEnv{C,E}, alg::SequentialCTMRG
@@ -238,7 +227,7 @@ function ctmrg_projectors(
     for (r, c) in directions
         # SVD half-infinite environment
         r′ = _prev(r, size(envs.corners, 2))
-        QQ = halfinfinite_environment(enlarged_envs[1][r, c], enlarged_envs[2][r′, c])
+        QQ = halfinfinite_environment(enlarged_envs[1, r, c], enlarged_envs[2, r′, c])
 
         trscheme = truncation_scheme(projector_alg, envs.edges[WEST, r′, c])
         svd_alg = svd_algorithm(projector_alg, (WEST, r, c))
@@ -255,7 +244,7 @@ function ctmrg_projectors(
 
         # Compute projectors
         P_bottom[r, c], P_top[r, c] = build_projectors(
-            U, S, V, enlarged_envs[1][r, c], enlarged_envs[2][r′, c]
+            U, S, V, enlarged_envs[1, r, c], enlarged_envs[2, r′, c]
         )
     end
 
@@ -391,9 +380,33 @@ end
 # Auxiliary routines
 # ======================================================================================== #
 
-# Build projectors from SVD and enlarged SW & NW corners
+"""
+    enlarge_corner(::Val{<:Int}, (r, c), envs, state, alg::CTMRG)
+
+Enlarge a CTMRG corner into the one of four directions `Val(1)` (NW), `Val(2)` (NE),
+`Val(3)` (SE) or `Val(4)` (SW). This serves as a wrapper that can dispatch on the
+directional index and different `CTMRG` algorithms.
+"""
+function enlarge_corner(::Val{1}, (r, c), envs, state, alg::CTMRG)
+    return enlarge_northwest_corner((r, c), envs, state)
+end
+function enlarge_corner(::Val{2}, (r, c), envs, state, alg::CTMRG)
+    return enlarge_northeast_corner((r, c), envs, state)
+end
+function enlarge_corner(::Val{3}, (r, c), envs, state, alg::CTMRG)
+    return enlarge_southeast_corner((r, c), envs, state)
+end
+function enlarge_corner(::Val{4}, (r, c), envs, state, alg::CTMRG)
+    return enlarge_southwest_corner((r, c), envs, state)
+end
+
+# Build projectors from SVD and dense enlarged corners
 function build_projectors(
-    U::AbstractTensorMap{E,3,1}, S, V::AbstractTensorMap{E,1,3}, Q, Q_next
+    U::AbstractTensorMap{E,3,1},
+    S::AbstractTensorMap{E,1,1},
+    V::AbstractTensorMap{E,1,3},
+    Q::AbstractTensorMap{E,3,3},
+    Q_next::AbstractTensorMap{E,3,3},
 ) where {E<:ElementarySpace}
     isqS = sdiag_inv_sqrt(S)
     P_left = Q_next * V' * isqS