Add FixedSVD, refactor build_projectors, improve SVD passing

src/algorithms/ctmrg.jl

@@ -52,7 +52,7 @@ function CTMRG(;
     verbosity=0,
     svd_alg=SVDAdjoint(),
     trscheme=FixedSpaceTruncation(),
-    ctmrgscheme=:AllSides
+    ctmrgscheme=:AllSides,
 )
     return CTMRG{ctmrgscheme}(
         tol, maxiter, miniter, verbosity, ProjectorAlg(; svd_alg, trscheme, verbosity)
@@ -133,7 +133,6 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRG)
     return envfix
 end
 
-
 """
     ctmrg_iter(state, env::CTMRGEnv{C,T}, alg::CTMRG) where {C,T}
     
@@ -165,7 +164,7 @@ function left_move(state, env::CTMRGEnv{C,T}, alg::ProjectorAlg) where {C,T}
     corners::typeof(env.corners) = copy(env.corners)
     edges::typeof(env.edges) = copy(env.edges)
     ϵ = 0.0
-    Pleft, Pright = Zygote.Buffer.(projector_type(T, size(state)))  # Use Zygote.Buffer instead of @diffset to avoid ZeroTangent errors in _setindex
+    P_top, P_bottom = Zygote.Buffer.(projector_type(T, size(state)))  # Use Zygote.Buffer instead of @diffset to avoid ZeroTangent errors in _setindex
 
     for col in 1:size(state, 2)
         cprev = _prev(col, size(state, 2))
@@ -192,11 +191,10 @@ function left_move(state, env::CTMRGEnv{C,T}, alg::ProjectorAlg) where {C,T}
 
             # SVD half-infinite environment
             trscheme = if alg.trscheme isa FixedSpaceTruncation
-                truncspace(space(env.edges[WEST, row, cnext], 1))
+                truncspace(space(env.edges[WEST, row, col], 1))
             else
                 alg.trscheme
             end
-            @tensor QQ[-1 -2 -3; -4 -5 -6] := Q_sw[-1 -2 -3; 1 2 3] * Q_nw[1 2 3; -4 -5 -6]
             @autoopt @tensor QQ[χ_EB D_EBabove D_EBbelow; χ_ET D_ETabove D_ETbelow] :=
                 Q_sw[χ_EB D_EBabove D_EBbelow; χ D1 D2] *
                 Q_nw[χ D1 D2; χ_ET D_ETabove D_ETbelow]
@@ -213,18 +211,18 @@ function left_move(state, env::CTMRGEnv{C,T}, alg::ProjectorAlg) where {C,T}
             end
 
             # Compute projectors
-            Pl, Pr = build_projectors(U, S, V, Q_sw, Q_nw)
-            Pleft[row, col] = Pl
-            Pright[row, col] = Pr
+            Pt, Pb = build_projectors(U, S, V, Q_sw, Q_nw)
+            P_top[row, col] = Pt
+            P_bottom[row, col] = Pb
         end
 
         # Use projectors to grow the corners & edges
         for row in 1:size(state, 1)
             rprev = _prev(row, size(state, 1))
             C_sw, C_nw, T_w = grow_env_left(
                 state[row, col],
-                Pleft[rprev, col],
-                Pright[row, col],
+                P_left[rprev, col],
+                P_right[row, col],
                 env.corners[SOUTHWEST, row, cprev],
                 env.corners[NORTHWEST, row, cprev],
                 env.edges[SOUTH, row, col],
@@ -237,7 +235,7 @@ function left_move(state, env::CTMRGEnv{C,T}, alg::ProjectorAlg) where {C,T}
         end
     end
 
-    return CTMRGEnv(corners, edges), (; Pleft=copy(Pleft), Pright=copy(Pright), ϵ)
+    return CTMRGEnv(corners, edges), (; P_left=copy(P_top), P_right=copy(P_bottom), ϵ)
 end
 
 # Compute enlarged corners
@@ -276,12 +274,12 @@ end
 
 # Build projectors from SVD and enlarged SW & NW corners
 function build_projectors(
-    U::AbstractTensorMap{E,3,1}, S, V::AbstractTensorMap{E,1,3}, Q_SW, Q_NW
+    U::AbstractTensorMap{E,3,1}, S, V::AbstractTensorMap{E,1,3}, Q, Q_next
 ) where {E<:ElementarySpace}
     isqS = sdiag_inv_sqrt(S)
-    P_bottom = Q_NW * V' * isqS
-    P_top = isqS * U' * Q_SW
-    return P_bottom, P_top
+    P_left = Q_next * V' * isqS
+    P_right = isqS * U' * Q
+    return P_left, P_right
 end
 
 # Apply projectors to entire left half-environment to grow SW & NW corners, and W edge

src/algorithms/ctmrg_all_sides.jl

@@ -1,15 +1,20 @@
+
+struct IndexSelector{N} <: TensorKit.TruncationScheme
+    index::NTuple{N,Int}
+end
+
 # One CTMRG iteration with both-sided application of projectors
 function ctmrg_iter(state, env::CTMRGEnv, alg::CTMRG{:AllSides})
     # Compute enlarged corners
     Q = enlarge_corners_edges(state, env)
 
     # Compute projectors if none are supplied
-    Pleft, Pright, info = build_projectors(Q, env, alg)
+    P_left, P_right, info = build_projectors(Q, env, alg)
 
     # Apply projectors and normalize
-    corners, edges = renormalize_corners_edges(state, env, Q, Pleft, Pright)
+    corners, edges = renormalize_corners_edges(state, env, Q, P_left, P_right)
 
-    return CTMRGEnv(corners, edges), (; Pleft, Pright, info...)
+    return CTMRGEnv(corners, edges), (; P_left, P_right, info...)
 end
 
 # Compute enlarged corners and edges for all directions and unit cell entries
@@ -52,90 +57,66 @@ function enlarge_corners_edges(state, env::CTMRGEnv)
 end
 
 # Build projectors from SVD and enlarged corners
-function build_projectors(Q, env::CTMRGEnv, alg::ProjectorAlg)  # TODO: Add projector type annotations
-    Pleft, Pright = Zygote.Buffer.(projector_type(env.edges))
+function build_projectors(Q, env::CTMRGEnv, alg::ProjectorAlg{A,T}) where {A,T}  # TODO: Add projector type annotations
+    P_left, P_right = Zygote.Buffer.(projector_type(env.edges))
     U, V = Zygote.Buffer.(projector_type(env.edges))
     S = Zygote.Buffer(env.corners)
-    ϵ1, ϵ2, ϵ3, ϵ4 = 0, 0, 0, 0
-    for c in 1:size(env.corners, 3), r in 1:size(env.corners, 2)
-        Pl1, Pr1, info1 = build_projectors(
-            Q[1, r, c],
-            Q[2, r, _next(c, end)],
-            alg;
-            trunc=truncspace(space(env.edges[1, r, c], 1)),
-            envindex=(1, r, c),
-        )
-        Pl2, Pr2, info2 = build_projectors(
-            Q[2, r, c],
-            Q[3, _next(r, end), c],
-            alg;
-            trunc=truncspace(space(env.edges[2, r, c], 1)),
-            envindex=(2, r, c),
-        )
-        Pl3, Pr3, info3 = build_projectors(
-            Q[3, r, c],
-            Q[4, r, _prev(c, end)],
-            alg;
-            trunc=truncspace(space(env.edges[3, r, c], 1)),
-            envindex=(3, r, c),
-        )
-        Pl4, Pr4, info4 = build_projectors(
-            Q[4, r, c],
-            Q[1, _prev(r, end), c],
-            alg;
-            trunc=truncspace(space(env.edges[4, r, c], 1)),
-            envindex=(4, r, c),
-        )
-
-        Pleft[NORTH, r, c] = Pl1
-        Pright[NORTH, r, c] = Pr1
-        U[NORTH, r, c] = info1.U
-        S[NORTH, r, c] = info1.S
-        V[NORTH, r, c] = info1.V
-
-        Pleft[EAST, r, c] = Pl2
-        Pright[EAST, r, c] = Pr2
-        U[EAST, r, c] = info2.U
-        S[EAST, r, c] = info2.S
-        V[EAST, r, c] = info2.V
-
-        Pleft[SOUTH, r, c] = Pl3
-        Pright[SOUTH, r, c] = Pr3
-        U[SOUTH, r, c] = info3.U
-        S[SOUTH, r, c] = info3.S
-        V[SOUTH, r, c] = info3.V
+    ϵ = 0.0
+    rsize, csize = size(env.corners)[2:3]
+    for dir in 1:4, r in 1:rsize, c in 1:csize
+        # Row-column index of next enlarged corner
+        next_rc in if dir == 1
+            (r, _next(c, csize))
+        elseif dir == 2
+            (_next(r, rsize), c)
+        elseif dir == 3
+            (r, _prev(c, csize))
+        elseif dir == 4
+            (_prev(r, rsize), c)
+        end
 
-        Pleft[WEST, r, c] = Pl4
-        Pright[WEST, r, c] = Pr4
-        U[WEST, r, c] = info4.U
-        S[WEST, r, c] = info4.S
-        V[WEST, r, c] = info4.V
-    end
-    return copy(Pleft), copy(Pright), (; ϵ=max(ϵ1, ϵ2, ϵ3, ϵ4), U=copy(U), S=copy(S), V=copy(V))
-end
-function build_projectors(Qleft, Qright, alg::ProjectorAlg; kwargs...)
-    # SVD half-infinite environment
-    U, S, V, ϵ = PEPSKit.tsvd!(Qleft * Qright, alg.svd_alg; kwargs...)
-    ϵ /= norm(S)
+        # SVD half-infinite environment
+        trscheme = if T <: FixedSpaceTruncation
+            truncspace(space(env.edges[dir, r, c], 1))
+        else
+            alg.trscheme
+        end
+        svd_alg = if A <: SVDAdjoint{<:FixedSVD}
+            idx = (dir, r, c)
+            fwd_alg = alg.svd_alg.fwd_alg
+            fix_svd = FixedSVD(fwd_alg.U[idx...], fwd_alg.S[idx...], fwd_alg.V[idx...])
+            return SVDAdjoint(; fwd_alg=fix_svd, rrule_alg=nothing, broadening=nothing)
+        else
+            alg.svd_alg
+        end
+        @autoopt @tensor QQ[χ_EB D_EBabove D_EBbelow; χ_ET D_ETabove D_ETbelow] :=
+            Q[dir, r, c][χ_EB D_EBabove D_EBbelow; χ D1 D2] *
+            Q[_next(dir, 4), next_rc...][χ D1 D2; χ_ET D_ETabove D_ETbelow]
+        U, S, V, ϵ_local = PEPSKit.tsvd!(QQ, svd_alg; trunc=trscheme)
+        ϵ = max(ϵ, ϵ_local / norm(S))
 
-    # Compute SVD truncation error and check for degenerate singular values
-    ignore_derivatives() do
-        if alg.verbosity > 1 && is_degenerate_spectrum(S)
-            svals = TensorKit.SectorDict(c => diag(b) for (c, b) in blocks(S))
-            @warn("degenerate singular values detected: ", svals)
+        # Compute SVD truncation error and check for degenerate singular values
+        ignore_derivatives() do
+            if alg.verbosity > 0 && is_degenerate_spectrum(S)
+                svals = TensorKit.SectorDict(c => diag(b) for (c, b) in blocks(S))
+                @warn("degenerate singular values detected: ", svals)
+            end
         end
-    end
 
-    # Compute projectors
-    isqS = sdiag_inv_sqrt(S)
-    @tensor Pl[-1 -2 -3; -4] := Qright[-1 -2 -3; 1 2 3] * conj(V[4; 1 2 3]) * isqS[4; -4]
-    @tensor Pr[-1; -2 -3 -4] := isqS[-1; 1] * conj(U[2 3 4; 1]) * Qleft[2 3 4; -2 -3 -4]
+        # Compute projectors
+        Pl, Pr = build_projectors(U, S, V, Q_sw, Q_nw)
+        P_left[dir, r, c] = Pl
+        P_right[dir, r, c] = Pr
+        U[dir, r, c] = info.U
+        S[dir, r, c] = info.S
+        V[dir, r, c] = info.V
+    end
 
-    return Pl, Pr, (; ϵ, U, S, V)
+    return copy(P_left), copy(P_right), (; ϵ, U=copy(U), S=copy(S), V=copy(V))
 end
 
 # Apply projectors to renormalize corners and edges
-function renormalize_corners_edges(state, env::CTMRGEnv, Q, Pleft, Pright)
+function renormalize_corners_edges(state, env::CTMRGEnv, Q, P_left, P_right)
     corners::typeof(env.corners) = copy(env.corners)
     edges::typeof(env.edges) = copy(env.edges)
     for c in 1:size(state, 2), r in 1:size(state, 1)
@@ -144,46 +125,46 @@ function renormalize_corners_edges(state, env::CTMRGEnv, Q, Pleft, Pright)
         cprev = _prev(c, size(state, 2))
         cnext = _next(c, size(state, 2))
         @diffset @tensor corners[NORTHWEST, r, c][-1; -2] :=
-            Pright[WEST, rnext, c][-1; 1 2 3] *
+            P_right[WEST, rnext, c][-1; 1 2 3] *
             Q[NORTHWEST, r, c][1 2 3; 4 5 6] *
-            Pleft[NORTH, r, c][4 5 6; -2]
+            P_left[NORTH, r, c][4 5 6; -2]
         @diffset @tensor corners[NORTHEAST, r, c][-1; -2] :=
-            Pright[NORTH, r, cprev][-1; 1 2 3] *
+            P_right[NORTH, r, cprev][-1; 1 2 3] *
             Q[NORTHEAST, r, c][1 2 3; 4 5 6] *
-            Pleft[EAST, r, c][4 5 6; -2]
+            P_left[EAST, r, c][4 5 6; -2]
         @diffset @tensor corners[SOUTHEAST, r, c][-1; -2] :=
-            Pright[EAST, rprev, c][-1; 1 2 3] *
+            P_right[EAST, rprev, c][-1; 1 2 3] *
             Q[SOUTHEAST, r, c][1 2 3; 4 5 6] *
-            Pleft[SOUTH, r, c][4 5 6; -2]
+            P_left[SOUTH, r, c][4 5 6; -2]
         @diffset @tensor corners[SOUTHWEST, r, c][-1; -2] :=
-            Pright[SOUTH, r, cnext][-1; 1 2 3] *
+            P_right[SOUTH, r, cnext][-1; 1 2 3] *
             Q[SOUTHWEST, r, c][1 2 3; 4 5 6] *
-            Pleft[WEST, r, c][4 5 6; -2]
+            P_left[WEST, r, c][4 5 6; -2]
 
         @diffset @tensor edges[NORTH, r, c][-1 -2 -3; -4] :=
             env.edges[NORTH, rprev, c][1 2 3; 4] *
             state[r, c][9; 2 5 -2 7] *
             conj(state[r, c][9; 3 6 -3 8]) *
-            Pleft[NORTH, r, c][4 5 6; -4] *
-            Pright[NORTH, r, cprev][-1; 1 7 8]
+            P_left[NORTH, r, c][4 5 6; -4] *
+            P_right[NORTH, r, cprev][-1; 1 7 8]
         @diffset @tensor edges[EAST, r, c][-1 -2 -3; -4] :=
             env.edges[EAST, r, _next(c, end)][1 2 3; 4] *
             state[r, c][9; 7 2 5 -2] *
             conj(state[r, c][9; 8 3 6 -3]) *
-            Pleft[EAST, r, c][4 5 6; -4] *
-            Pright[EAST, rprev, c][-1; 1 7 8]
+            P_left[EAST, r, c][4 5 6; -4] *
+            P_right[EAST, rprev, c][-1; 1 7 8]
         @diffset @tensor edges[SOUTH, r, c][-1 -2 -3; -4] :=
             env.edges[SOUTH, _next(r, end), c][1 2 3; 4] *
             state[r, c][9; -2 7 2 5] *
             conj(state[r, c][9; -3 8 3 6]) *
-            Pleft[SOUTH, r, c][4 5 6; -4] *
-            Pright[SOUTH, r, cnext][-1; 1 7 8]
+            P_left[SOUTH, r, c][4 5 6; -4] *
+            P_right[SOUTH, r, cnext][-1; 1 7 8]
         @diffset @tensor edges[WEST, r, c][-1 -2 -3; -4] :=
             env.edges[WEST, r, _prev(c, end)][1 2 3; 4] *
             state[r, c][9; 5 -2 7 2] *
             conj(state[r, c][9; 6 -3 8 3]) *
-            Pleft[WEST, r, c][4 5 6; -4] *
-            Pright[WEST, rnext, c][-1; 1 7 8]
+            P_left[WEST, r, c][4 5 6; -4] *
+            P_right[WEST, rnext, c][-1; 1 7 8]
     end
 
     @diffset corners[:, :, :] ./= norm.(corners[:, :, :])

src/utility/svd.jl

@@ -168,6 +168,24 @@ function ChainRulesCore.rrule(
     return (U, S, V, ϵ), tsvd!_itersvd_pullback
 end
 
+"""
+    struct FixedSVD
+
+SVD struct containing a pre-computed decomposition or even multiple ones.
+The call to `tsvd` just returns the pre-computed U, S and V. In the reverse
+pass, the SVD adjoint is computed with these exact U, S, and V.
+"""
+struct FixedSVD{Ut,St,Vt}
+    U::Ut
+    S::St
+    V::Vt
+end
+
+# Return pre-computed SVD
+function TensorKit._tsvd!(_, alg::FixedSVD, ::NoTruncation, ::Real=2)
+    return alg.U, alg.S, alg.V, 0
+end
+
 """
     struct NonTruncAdjoint(; lorentz_broadening = 0.0)
 

src/utility/util.jl

@@ -51,16 +51,24 @@ end
 
 """
     projector_type(T::DataType, size)
+    projector_type(edges::Array{<:AbstractTensorMap})
 
 Create two arrays of specified `size` that contain undefined tensors representing
 left and right acting projectors, respectively. The projector types are inferred
 from the TensorMap type `T` which avoids having to recompute transpose tensors.
+Alternatively, supply an array of edge tensors from which left and right projectors
+are intialized explicitly with zeros.
 """
 function projector_type(T::DataType, size)
-    Pleft = Array{T,length(size)}(undef, size)
+    P_left = Array{T,length(size)}(undef, size)
     Prtype = tensormaptype(spacetype(T), numin(T), numout(T), storagetype(T))
-    Pright = Array{Prtype,length(size)}(undef, size)
-    return Pleft, Pright
+    P_right = Array{Prtype,length(size)}(undef, size)
+    return P_left, P_right
+end
+function projector_type(edges::Array{<:AbstractTensorMap})
+    P_left = map(e -> TensorMap(zeros, scalartype(e), space(e)), edges)
+    P_right = map(e -> TensorMap(zeros, scalartype(e), domain(e), codomain(e)), edges)
+    return P_left, P_right
 end
 
 # There are no rrules for rotl90 and rotr90 in ChainRules.jl