Implement truncated eigenvalues

Author: lkdvos

ext/TensorKitChainRulesCoreExt/factorizations.jl

@@ -1,32 +1,38 @@
 # Factorizations rules
 # --------------------
-function ChainRulesCore.rrule(::typeof(TensorKit.tsvd!), t::AbstractTensorMap;
-                              trunc::TruncationStrategy=TensorKit.notrunc(),
-                              kwargs...)
-    # TODO: I think we can use tsvd! here without issues because we don't actually require
-    # the data of `t` anymore.
-    USVᴴ = tsvd(t; trunc=TensorKit.notrunc(), kwargs...)
-
-    if trunc != TensorKit.notrunc() && !isempty(blocksectors(t))
-        USVᴴ′ = MatrixAlgebraKit.truncate!(svd_trunc!, USVᴴ, trunc)
-    else
-        USVᴴ′ = USVᴴ
-    end
+for f in (:tsvd, :eig, :eigh)
+    f! = Symbol(f, :!)
+    f_trunc! = f == :tsvd ? :svd_trunc! : Symbol(f, :_trunc!)
+    f_pullback = Symbol(f, :_pullback)
+    f_pullback! = f == :tsvd ? :svd_compact_pullback! : Symbol(f, :_full_pullback!)
+    @eval function ChainRulesCore.rrule(::typeof(TensorKit.$f!), t::AbstractTensorMap;
+                                        trunc::TruncationStrategy=TensorKit.notrunc(),
+                                        kwargs...)
+        # TODO: I think we can use f! here without issues because we don't actually require
+        # the data of `t` anymore.
+        F = $f(t; trunc=TensorKit.notrunc(), kwargs...)
+
+        if trunc != TensorKit.notrunc() && !isempty(blocksectors(t))
+            F′ = MatrixAlgebraKit.truncate!($f_trunc!, F, trunc)
+        else
+            F′ = F
+        end
 
-    function tsvd!_pullback(ΔUSVᴴ′)
-        ΔUSVᴴ = unthunk.(ΔUSVᴴ′)
-        Δt = zerovector(t)
-        foreachblock(Δt) do c, (b,)
-            USVᴴc = block.(USVᴴ, Ref(c))
-            ΔUSVᴴc = block.(ΔUSVᴴ, Ref(c))
-            svd_compact_pullback!(b, USVᴴc, ΔUSVᴴc)
-            return nothing
+        function $f_pullback(ΔF′)
+            ΔF = unthunk.(ΔF′)
+            Δt = zerovector(t)
+            foreachblock(Δt) do c, (b,)
+                Fc = block.(F, Ref(c))
+                ΔFc = block.(ΔF, Ref(c))
+                $f_pullback!(b, Fc, ΔFc)
+                return nothing
+            end
+            return NoTangent(), Δt
         end
-        return NoTangent(), Δt
-    end
-    tsvd!_pullback(::NTuple{3,ZeroTangent}) = NoTangent(), ZeroTangent()
+        $f_pullback(::Tuple{ZeroTangent,Vararg{ZeroTangent}}) = NoTangent(), ZeroTangent()
 
-    return USVᴴ′, tsvd!_pullback
+        return F′, $f_pullback
+    end
 end
 
 function ChainRulesCore.rrule(::typeof(LinearAlgebra.svdvals!), t::AbstractTensorMap)
@@ -43,44 +49,6 @@ function ChainRulesCore.rrule(::typeof(LinearAlgebra.svdvals!), t::AbstractTenso
     return s, svdvals_pullback
 end
 
-function ChainRulesCore.rrule(::typeof(TensorKit.eig!), t::AbstractTensorMap; kwargs...)
-    DV = eig(t; kwargs...)
-
-    function eig!_pullback(ΔDV′)
-        ΔDV = unthunk.(ΔDV′)
-        Δt = zerovector(t)
-        foreachblock(Δt) do c, (b,)
-            DVc = block.(DV, Ref(c))
-            ΔDVc = block.(ΔDV, Ref(c))
-            eig_full_pullback!(b, DVc, ΔDVc)
-            return nothing
-        end
-        return NoTangent(), Δt
-    end
-    eig!_pullback(::NTuple{2,ZeroTangent}) = NoTangent(), ZeroTangent()
-
-    return DV, eig!_pullback
-end
-
-function ChainRulesCore.rrule(::typeof(TensorKit.eigh!), t::AbstractTensorMap; kwargs...)
-    DV = eigh(t; kwargs...)
-
-    function eigh!_pullback(ΔDV′)
-        ΔDV = unthunk.(ΔDV′)
-        Δt = zerovector(t)
-        foreachblock(Δt) do c, (b,)
-            DVc = block.(DV, Ref(c))
-            ΔDVc = block.(ΔDV, Ref(c))
-            eigh_full_pullback!(b, DVc, ΔDVc)
-            return nothing
-        end
-        return NoTangent(), Δt
-    end
-    eigh!_pullback(::NTuple{2,ZeroTangent}) = NoTangent(), ZeroTangent()
-
-    return DV, eigh!_pullback
-end
-
 function ChainRulesCore.rrule(::typeof(LinearAlgebra.eigvals!), t::AbstractTensorMap;
                               sortby=nothing, kwargs...)
     @assert sortby === nothing "only `sortby=nothing` is supported"

src/tensors/factorizations/implementations.jl

@@ -149,9 +149,27 @@ rightpolar!(t::AbstractTensorMap; kwargs...) = right_polar!(t; kwargs...)
 
 # Eigenvalue decomposition
 # ------------------------
-eigh!(t::AbstractTensorMap) = eigh_full!(t)
-eig!(t::AbstractTensorMap) = eig_full!(t)
-eigen!(t::AbstractTensorMap) = ishermitian(t) ? eigh!(t) : eig!(t)
+function eigh!(t::AbstractTensorMap; trunc=notrunc(), kwargs...)
+    InnerProductStyle(t) === EuclideanInnerProduct() || throw_invalid_innerproduct(:eigh!)
+    if trunc == notrunc()
+        return eigh_full!(t; kwargs...)
+    else
+        return eigh_trunc!(t; trunc, kwargs...)
+    end
+end
+
+function eig!(t::AbstractTensorMap; trunc=notrunc(), kwargs...)
+    InnerProductStyle(t) === EuclideanInnerProduct() || throw_invalid_innerproduct(:eig!)
+    if trunc == notrunc()
+        return eig_full!(t; kwargs...)
+    else
+        return eig_trunc!(t; trunc, kwargs...)
+    end
+end
+
+function eigen!(t::AbstractTensorMap; kwargs...)
+    return ishermitian(t) ? eigh!(t; kwargs...) : eig!(t; kwargs...)
+end
 
 # Singular value decomposition
 # ----------------------------

src/tensors/factorizations/matrixalgebrakit.jl

@@ -193,6 +193,26 @@ function initialize_output(::typeof(eig_full!), t::AbstractTensorMap, ::Abstract
     return D, V
 end
 
+function initialize_output(::typeof(eigh_trunc!), t::AbstractTensorMap,
+                           alg::TruncatedAlgorithm)
+    return initialize_output(eigh_full!, t, alg.alg)
+end
+
+function initialize_output(::typeof(eig_trunc!), t::AbstractTensorMap,
+                           alg::TruncatedAlgorithm)
+    return initialize_output(eig_full!, t, alg.alg)
+end
+
+function eigh_trunc!(t::AbstractTensorMap, DV, alg::TruncatedAlgorithm)
+    DV′ = eigh_full!(t, DV, alg.alg)
+    return truncate!(eigh_trunc!, DV′, alg.trunc)
+end
+
+function eig_trunc!(t::AbstractTensorMap, DV, alg::TruncatedAlgorithm)
+    DV′ = eig_full!(t, DV, alg.alg)
+    return truncate!(eig_trunc!, DV′, alg.trunc)
+end
+
 # QR decomposition
 # ----------------
 const _T_QR = Tuple{<:AbstractTensorMap,<:AbstractTensorMap}

src/tensors/factorizations/truncation.jl

@@ -65,6 +65,47 @@ function truncate!(::typeof(left_null!),
     return Ũ
 end
 
+function truncate!(::typeof(eigh_trunc!), (D, V)::_T_DV, strategy::TruncationStrategy)
+    ind = findtruncated(diagview(D), strategy)
+    V_truncated = spacetype(D)(c => length(I) for (c, I) in ind)
+
+    D̃ = DiagonalTensorMap{scalartype(D)}(undef, V_truncated)
+    for (c, b) in blocks(D̃)
+        I = get(ind, c, nothing)
+        @assert !isnothing(I)
+        copy!(b.diag, @view(block(D, c).diag[I]))
+    end
+
+    Ṽ = similar(V, V_truncated ← domain(V))
+    for (c, b) in blocks(Ṽ)
+        I = get(ind, c, nothing)
+        @assert !isnothing(I)
+        copy!(b, @view(block(V, c)[I, :]))
+    end
+
+    return D̃, Ṽ
+end
+function truncate!(::typeof(eig_trunc!), (D, V)::_T_DV, strategy::TruncationStrategy)
+    ind = findtruncated(diagview(D), strategy)
+    V_truncated = spacetype(D)(c => length(I) for (c, I) in ind)
+
+    D̃ = DiagonalTensorMap{scalartype(D)}(undef, V_truncated)
+    for (c, b) in blocks(D̃)
+        I = get(ind, c, nothing)
+        @assert !isnothing(I)
+        copy!(b.diag, @view(block(D, c).diag[I]))
+    end
+
+    Ṽ = similar(V, V_truncated ← domain(V))
+    for (c, b) in blocks(Ṽ)
+        I = get(ind, c, nothing)
+        @assert !isnothing(I)
+        copy!(b, @view(block(V, c)[I, :]))
+    end
+
+    return D̃, Ṽ
+end
+
 # Find truncation
 # ---------------
 # auxiliary functions
@@ -88,18 +129,28 @@ function _findnexttruncvalue(S, truncdim::SectorDict{I,Int}) where {I<:Sector}
     return σmin, keys(truncdim)[imin]
 end
 
-# sorted implementations
+# implementations
 function findtruncated_sorted(S::SectorDict, strategy::TruncationKeepAbove)
     atol = rtol_to_atol(S, strategy.p, strategy.atol, strategy.rtol)
     findtrunc = Base.Fix2(findtruncated_sorted, truncbelow(atol))
     return SectorDict(c => findtrunc(d) for (c, d) in Sd)
 end
+function findtruncated(S::SectorDict, strategy::TruncationKeepAbove)
+    atol = rtol_to_atol(S, strategy.p, strategy.atol, strategy.rtol)
+    findtrunc = Base.Fix2(findtruncated, truncbelow(atol))
+    return SectorDict(c => findtrunc(d) for (c, d) in Sd)
+end
 
 function findtruncated_sorted(S::SectorDict, strategy::TruncationKeepBelow)
     atol = rtol_to_atol(S, strategy.p, strategy.atol, strategy.rtol)
     findtrunc = Base.Fix2(findtruncated_sorted, truncabove(atol))
     return SectorDict(c => findtrunc(d) for (c, d) in Sd)
 end
+function findtruncated(S::SectorDict, strategy::TruncationKeepBelow)
+    atol = rtol_to_atol(S, strategy.p, strategy.atol, strategy.rtol)
+    findtrunc = Base.Fix2(findtruncated, truncabove(atol))
+    return SectorDict(c => findtrunc(d) for (c, d) in Sd)
+end
 
 function findtruncated_sorted(Sd::SectorDict, strategy::TruncationError)
     I = keytype(Sd)
@@ -153,9 +204,17 @@ end
 function findtruncated_sorted(Sd::SectorDict, strategy::TruncationKeepFiltered)
     return SectorDict(c => findtruncated_sorted(d, strategy) for (c, d) in Sd)
 end
+function findtruncated(Sd::SectorDict, strategy::TruncationKeepFiltered)
+    return SectorDict(c => findtruncated(d, strategy) for (c, d) in Sd)
+end
 
 function findtruncated_sorted(Sd::SectorDict, strategy::TruncationIntersection)
     inds = map(Base.Fix1(findtruncated_sorted, Sd), strategy)
     return SectorDict(c => intersect(map(Base.Fix2(getindex, c), inds)...)
                       for c in intersect(map(keys, inds)...))
 end
+function findtruncated(Sd::SectorDict, strategy::TruncationIntersection)
+    inds = map(Base.Fix1(findtruncated, Sd), strategy)
+    return SectorDict(c => intersect(map(Base.Fix2(getindex, c), inds)...)
+                      for c in intersect(map(keys, inds)...))
+end