working randomized SVD

Author: ogauthe

src/PEPSKit.jl

@@ -5,6 +5,7 @@ using Compat
 using Accessors: @set, @reset
 using VectorInterface
 import VectorInterface as VI
+import Random
 
 using TensorKit
 using TensorKit: TruncationScheme

src/algorithms/contractions/ctmrg_contractions.jl

@@ -251,6 +251,17 @@ function contract_projectors(U, S, V, Q, Q_next)
     return P_left, P_right
 end
 
+function contract_projectors(U, S, V, fq::FourQuadrants)
+    isqS = sdiag_pow(S, -0.5)
+    # use * to respect fermionic case
+    p1 = (codomainind(fq), domainind(fq))
+    p2 = (codomainind(fq), (numout(fq) + 1,))
+    P_left = tensorcontract(fq.Q3, p1, false, fq.Q4 * V' * isqS, p2, false, p2)
+    p3 = ((1,), codomainind(fq) .+ 1)
+    P_right = tensorcontract(isqS * U' * fq.Q1, p3, false, fq.Q2, p1, false, p3)
+    return P_left, P_right
+end
+
 """
     half_infinite_environment(quadrant1, quadrant2)
     half_infinite_environment(C_1, C_2, E_1, E_2, E_3, E_4, A_1, A_2)

src/algorithms/ctmrg/projectors.jl

@@ -87,7 +87,7 @@ Construct the half-infinite projector algorithm based on the following keyword a
     - `:notrunc` : No singular values are truncated and the performed SVDs are exact
     - `:truncerr` : Additionally supply error threshold `η`; truncate to the maximal virtual dimension of `η`
     - `:truncdim` : Additionally supply truncation dimension `η`; truncate such that the 2-norm of the truncated values is smaller than `η`
-    - `:truncspace` : Additionally supply truncation space `η`; truncate according to the supplied vector space 
+    - `:truncspace` : Additionally supply truncation space `η`; truncate according to the supplied vector space
     - `:truncbelow` : Additionally supply singular value cutoff `η`; truncate such that every retained singular value is larger than `η`
 * `verbosity::Int=$(Defaults.projector_verbosity)` : Projector output verbosity which can be:
     0. Suppress output information
@@ -125,7 +125,7 @@ Construct the full-infinite projector algorithm based on the following keyword a
     - `:notrunc` : No singular values are truncated and the performed SVDs are exact
     - `:truncerr` : Additionally supply error threshold `η`; truncate to the maximal virtual dimension of `η`
     - `:truncdim` : Additionally supply truncation dimension `η`; truncate such that the 2-norm of the truncated values is smaller than `η`
-    - `:truncspace` : Additionally supply truncation space `η`; truncate according to the supplied vector space 
+    - `:truncspace` : Additionally supply truncation space `η`; truncate according to the supplied vector space
     - `:truncbelow` : Additionally supply singular value cutoff `η`; truncate such that every retained singular value is larger than `η`
 * `verbosity::Int=$(Defaults.projector_verbosity)` : Projector output verbosity which can be:
     0. Suppress output information
@@ -148,7 +148,7 @@ PROJECTOR_SYMBOLS[:fullinfinite] = FullInfiniteProjector
 Determine left and right projectors at the bond given determined by the enlarged corners
 and the given coordinate using the specified `alg`.
 """
-function compute_projector(enlarged_corners, coordinate, alg::HalfInfiniteProjector)
+function compute_projector(enlarged_corners, coordinate, last_space, alg::HalfInfiniteProjector)
     # SVD half-infinite environment
     halfinf = half_infinite_environment(enlarged_corners...)
     svd_alg = svd_algorithm(alg, coordinate)
@@ -166,7 +166,7 @@ function compute_projector(enlarged_corners, coordinate, alg::HalfInfiniteProjec
     P_left, P_right = contract_projectors(U, S, V, enlarged_corners...)
     return (P_left, P_right), (; U, S, V, info...)
 end
-function compute_projector(enlarged_corners, coordinate, alg::FullInfiniteProjector)
+function compute_projector(enlarged_corners, coordinate, last_space, alg::FullInfiniteProjector)
     halfinf_left = half_infinite_environment(enlarged_corners[1], enlarged_corners[2])
     halfinf_right = half_infinite_environment(enlarged_corners[3], enlarged_corners[4])
 
@@ -187,3 +187,121 @@ function compute_projector(enlarged_corners, coordinate, alg::FullInfiniteProjec
     P_left, P_right = contract_projectors(U, S, V, halfinf_left, halfinf_right)
     return (P_left, P_right), (; U, S, V, info...)
 end
+
+# ==========================================================================================
+# TBD: compute ncv vectors proj.ncv_ratio * old_nvectors?
+struct RandomizedProjector{S, T, R} <: ProjectorAlgorithm
+    svd_alg::S
+    trscheme::T
+    rng::R
+    oversampling::Int64
+    max_full::Int64
+    verbosity::Int
+end
+
+PROJECTOR_SYMBOLS[:randomized] = RandomizedProjector
+
+svd_algorithm(alg::RandomizedProjector) = alg.svd_alg
+
+#=
+# TBD is this needed?
+function RandomizedProjector(;
+        rng = Random.default_rng(),
+        svd_alg = (;),
+        trscheme = (;),
+        oversampling = 10,
+        max_full = 100,
+        verbosity = Defaults.projector_verbosity,
+    )
+
+    # parse SVD forward & rrule algorithm
+    svd_algorithm = _alg_or_nt(SVDAdjoint, svd_alg)
+
+    # parse truncation scheme
+    truncation_scheme = if trscheme isa TruncationScheme
+        trscheme
+    elseif trscheme isa NamedTuple
+        _TruncationScheme(; trscheme...)
+    else
+        throw(ArgumentError("unknown trscheme $trscheme"))
+    end
+
+    return RandomizedProjector(
+        svd_algorithm, truncation_scheme, rng, oversampling, max_full, verbosity
+    )
+end
+=#
+
+
+# needed as default interface in PEPSKit.ProjectorAlgorithm
+function RandomizedProjector(svd_algorithm, trscheme, verbosity)
+    @show "Hi RandomizedProjector"
+    @show which(
+        RandomizedProjector, typeof.(
+            (
+                svd_algorithm, trscheme, Random.default_rng(), 10, 100, verbosity,
+            )
+        )
+    )
+    return RandomizedProjector(
+        svd_algorithm, trscheme, Random.default_rng(), 10, 100, verbosity
+    )
+end
+
+function random_domain(alg::RandomizedProjector, full_space, last_space)
+    sector_dims = map(sectors(full_space)) do s
+        if dim(full_space, s) <= alg.max_full
+            n = dim(full_space, s)
+        else
+            n = dim(last_space, s) + alg.oversampling
+        end
+        return s => n
+    end
+    return Vect[sectortype(last_space)](sector_dims)
+end
+
+
+function randomized_range_finder(A::AbstractTensorMap, alg::RandomizedProjector, randomized_space)
+    Q = TensorMap{eltype(A)}(undef, domain(A) ← randomized_space)
+    foreach(blocks(Q)) do (s, b)
+        m, n = size(b)
+        if m <= alg.max_full
+            b .= LinearAlgebra.I(m)
+        else
+            Ω = randn(alg.rng, eltype(A), domain(A) ← Vect[sectortype(A)](s => n))
+            Y = A * Ω
+            Qs, _ = leftorth!(Y)
+            b .= block(Qs, s)
+        end
+    end
+    return Q
+end
+
+# impose full env, could also be defined for half_infinite_environment, little gain
+function compute_projector(fq, coordinate, last_space, alg::RandomizedProjector)
+    full_space = fuse(domain(fq))
+    randomized_space = random_domain(alg, full_space, last_space)
+    Q = randomized_range_finder(fq, alg, randomized_space)
+    B = Q' * fq
+
+    svd_alg = svd_algorithm(alg, coordinate)
+    U′, S, V, info = tsvd!(B, svd_alg; trunc = alg.trscheme)
+    U = Q * U′
+    foreach(blocks(S)) do (s, b)
+        if size(b, 1) == dim(randomized_space, s) && size(b, 1) < dim(full_space, s)
+            @warn("Sector is too small, kept all computed values: ", s)
+        end
+    end
+
+    # Check for degenerate singular values; still needed for exact blocks
+    Zygote.isderiving() && 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
+
+    @reset info.truncation_error = info.truncation_error / norm(S) # normalize truncation error
+    P_left, P_right = contract_projectors(U, S, V, fq)
+    return (P_left, P_right), (; U, S, V, info...)
+end

src/algorithms/ctmrg/simultaneous.jl

@@ -97,7 +97,8 @@ function simultaneous_projectors(
     trscheme = truncation_scheme(alg, env.edges[coordinate[1], coordinate′[2:3]...])
     alg′ = @set alg.trscheme = trscheme
     ec = (enlarged_corners[coordinate...], enlarged_corners[coordinate′...])
-    return compute_projector(ec, coordinate, alg′)
+    last_space = space(env.edges[coordinate[1], coordinate′[2:3]...], 1)
+    return compute_projector(ec, coordinate, last_space, alg′)
 end
 function simultaneous_projectors(
         coordinate, enlarged_corners::Array{E, 3}, env, alg::FullInfiniteProjector
@@ -109,13 +110,35 @@ function simultaneous_projectors(
     coordinate2 = _next_coordinate(coordinate, rowsize, colsize)
     coordinate3 = _next_coordinate(coordinate2, rowsize, colsize)
     coordinate4 = _next_coordinate(coordinate3, rowsize, colsize)
+    last_space = space(env.edges[coordinate[1], coordinate′[2:3]...], 1)
     ec = (
         enlarged_corners[coordinate4...],
         enlarged_corners[coordinate...],
         enlarged_corners[coordinate2...],
         enlarged_corners[coordinate3...],
     )
-    return compute_projector(ec, coordinate, alg′)
+    return compute_projector(ec, coordinate, last_space, alg′)
+end
+
+# TBD share code with FullInfiniteProjector?
+function simultaneous_projectors(
+        coordinate, enlarged_corners::Array{E, 3}, env, alg::RandomizedProjector
+    ) where {E}
+    coordinate′ = _next_coordinate(coordinate, size(env)[2:3]...)
+    trscheme = truncation_scheme(alg, env.edges[coordinate[1], coordinate′[2:3]...])
+    alg′ = @set alg.trscheme = trscheme
+    rowsize, colsize = size(enlarged_corners)[2:3]
+    coordinate2 = _next_coordinate(coordinate, rowsize, colsize)
+    coordinate3 = _next_coordinate(coordinate2, rowsize, colsize)
+    coordinate4 = _next_coordinate(coordinate3, rowsize, colsize)
+    fq = FourQuadrants(
+        enlarged_corners[coordinate4...],
+        enlarged_corners[coordinate...],
+        enlarged_corners[coordinate2...],
+        enlarged_corners[coordinate3...],
+    )
+    last_space = space(env.edges[coordinate[1], coordinate′[2:3]...], 1)
+    return compute_projector(fq, coordinate, last_space, alg′)
 end
 
 """

src/environments/ctmrg_environments.jl

@@ -477,3 +477,34 @@ function VI.inner(env₁::CTMRGEnv, env₂::CTMRGEnv)
     return inner((env₁.corners, env₁.edges), (env₂.corners, env₂.edges))
 end
 VI.norm(env::CTMRGEnv) = norm((env.corners, env.edges))
+
+# ==========================================================================================
+
+struct FourQuadrants{S, T, N, TM} <: AbstractTensorMap{S, T, N, N}
+    Q1::TM
+    Q2::TM
+    Q3::TM
+    Q4::TM
+
+    function FourQuadrants{T, S, N, TM}(
+            Q1, Q2, Q3, Q4
+        ) where {S, T, N, TM <: AbstractTensorMap{T, S, N, N}}
+        return new{T, S, N, TM}(Q1, Q2, Q3, Q4)
+    end
+end
+
+function FourQuadrants(Q1, Q2, Q3, Q4)
+    return FourQuadrants{eltype(Q1), spacetype(Q1), numin(Q1), typeof(Q1)}(Q1, Q2, Q3, Q4)
+end
+
+TensorKit.TensorMap(fq::FourQuadrants) = fq.Q1 * fq.Q2 * fq.Q3 * fq.Q4
+
+TensorKit.space(fq::FourQuadrants) = codomain(fq.Q1) ← domain(fq.Q4)
+
+# TBD use tensorcontract to handle fermions?
+(fq::FourQuadrants)(m) = fq.Q1 * (fq.Q2 * (fq.Q3 * (fq.Q4 * m)))
+
+Base.:*(m::AbstractTensorMap, fq::FourQuadrants) = m * fq.Q1 * fq.Q2 * fq.Q3 * fq.Q4
+Base.:*(fq::FourQuadrants, m::AbstractTensorMap) = fq.Q1 * (fq.Q2 * (fq.Q3 * (fq.Q4 * m)))
+
+Base.adjoint(fq::FourQuadrants) = FourQuadrants(fq.Q4', fq.Q3', fq.Q2', fq.Q1')