use truncation_error everywhere

Author: lkdvos

ext/MatrixAlgebraKitChainRulesCoreExt.jl

@@ -2,7 +2,7 @@ module MatrixAlgebraKitChainRulesCoreExt
 
 using MatrixAlgebraKit
 using MatrixAlgebraKit: copy_input, initialize_output, zero!, diagview,
-    TruncatedAlgorithm, findtruncated, findtruncated_svd, compute_truncerr!
+    TruncatedAlgorithm, findtruncated, findtruncated_svd, truncation_error
 using ChainRulesCore
 using LinearAlgebra
 
@@ -113,7 +113,7 @@ for eig in (:eig, :eigh)
             Ac = copy_input($eig_f, A)
             DV = $(eig_f!)(Ac, DV, alg.alg)
             DV′, ind = MatrixAlgebraKit.truncate($eig_t!, DV, alg.trunc)
-            ϵ = compute_truncerr!(copy(diagview(DV[1])), ind)
+            ϵ = truncation_error(diagview(DV[1]), ind)
             return (DV′..., ϵ), $(_make_eig_t_pb)(A, DV, ind)
         end
         function $(_make_eig_t_pb)(A, DV, ind)
@@ -157,7 +157,7 @@ function ChainRulesCore.rrule(::typeof(svd_trunc!), A, USVᴴ, alg::TruncatedAlg
     Ac = copy_input(svd_compact, A)
     USVᴴ = svd_compact!(Ac, USVᴴ, alg.alg)
     USVᴴ′, ind = MatrixAlgebraKit.truncate(svd_trunc!, USVᴴ, alg.trunc)
-    ϵ = compute_truncerr!(copy(diagview(USVᴴ[2])), ind)
+    ϵ = truncation_error(diagview(USVᴴ[2]), ind)
     return (USVᴴ′..., ϵ), _make_svd_trunc_pullback(A, USVᴴ, ind)
 end
 function _make_svd_trunc_pullback(A, USVᴴ, ind)

src/MatrixAlgebraKit.jl

@@ -54,7 +54,7 @@ export notrunc, truncrank, trunctol, truncerror, truncfilter
     eval(
         Expr(
             :public, :TruncationByOrder, :TruncationByFilter, :TruncationByValue,
-            :TruncationByError, :TruncationIntersection, :truncate
+            :TruncationByError, :TruncationIntersection, :truncate, :truncation_error
         )
     )
     eval(

src/implementations/eig.jl

@@ -109,7 +109,8 @@ end
 function eig_trunc!(A, DV, alg::TruncatedAlgorithm)
     D, V = eig_full!(A, DV, alg.alg)
     DVtrunc, ind = truncate(eig_trunc!, (D, V), alg.trunc)
-    return DVtrunc..., compute_truncerr!(diagview(D), ind)
+    ϵ = truncation_error(diagview(D), ind)
+    return DVtrunc..., ϵ
 end
 
 # Diagonal logic

src/implementations/eigh.jl

@@ -112,7 +112,8 @@ end
 function eigh_trunc!(A, DV, alg::TruncatedAlgorithm)
     D, V = eigh_full!(A, DV, alg.alg)
     DVtrunc, ind = truncate(eigh_trunc!, (D, V), alg.trunc)
-    return DVtrunc..., compute_truncerr!(diagview(D), ind)
+    ϵ = truncation_error(diagview(D), ind)
+    return DVtrunc..., ϵ
 end
 
 # Diagonal logic

src/implementations/svd.jl

@@ -239,7 +239,8 @@ end
 function svd_trunc!(A, USVᴴ, alg::TruncatedAlgorithm)
     U, S, Vᴴ = svd_compact!(A, USVᴴ, alg.alg)
     USVᴴtrunc, ind = truncate(svd_trunc!, (U, S, Vᴴ), alg.trunc)
-    return USVᴴtrunc..., compute_truncerr!(diagview(S), ind)
+    ϵ = truncation_error(diagview(S), ind)
+    return USVᴴtrunc..., ϵ
 end
 
 # Diagonal logic
@@ -385,7 +386,7 @@ function svd_trunc!(A::AbstractMatrix, USVᴴ, alg::TruncatedAlgorithm{<:GPU_Ran
     # TODO: make sure that truncation is based on maxrank, otherwise this might be wrong
     USVᴴtrunc, ind = truncate(svd_trunc!, (U, S, Vᴴ), alg.trunc)
     Strunc = diagview(USVᴴtrunc[2])
-    # normal `compute_truncerr!` does not work here since `S` is not the full singular value spectrum
+    # normal `truncation_error` does not work here since `S` is not the full singular value spectrum
     ϵ = sqrt(norm(A)^2 - norm(Strunc)^2) # is there a more accurate way to do this?
     return USVᴴtrunc..., ϵ
 end

src/implementations/truncation.jl

@@ -117,30 +117,26 @@ _ind_intersect(A::AbstractVector, B::AbstractVector{Bool}) = _ind_intersect(B, A
 _ind_intersect(A::AbstractVector{Bool}, B::AbstractVector{Bool}) = A .& B
 _ind_intersect(A, B) = intersect(A, B)
 
-# Compute truncation error as 2-norm of discarded values
-# by destroying original values
-function compute_truncerr!(values::AbstractVector, ind)
-    values[ind] .= zero(eltype(values))
-    return norm(values)
-end
+# Truncation error
+# ----------------
+# generic fallback: no allocations but inaccurate
+truncation_error(values::AbstractVector, ind) = sqrt(norm(values)^2 - norm(view(values, ind))^2)
 
-function compute_truncerr(values::AbstractVector{<:Number}, ind::AbstractUnitRange)
+function truncation_error(values::AbstractVector{<:Number}, ind::AbstractUnitRange)
     init = abs2(zero(eltype(values)))
     return sqrt(
         sum(abs2, view(values, firstindex(values):(first(ind) - 1)); init) +
             sum(abs2, view(values, (last(ind) + 1):lastindex(values)); init)
     )
 end
-
-function compute_truncerr(values::AbstractVector{<:Number}, ind::AbstractVector{Bool})
+function truncation_error(values::AbstractVector{<:Number}, ind::AbstractVector{Bool})
     init = abs2(zero(eltype(values)))
     @inbounds for i in eachindex(values, ind)
         init += abs2(values[i] * ~(ind[i]))
     end
     return sqrt(init)
 end
-
-function compute_truncerr(values::AbstractVector{<:Number}, ind::AbstractVector{<:Integer})
+function truncation_error(values::AbstractVector{<:Number}, ind::AbstractVector{<:Integer})
     sort!(ind)
     allind = eachindex(IndexLinear(), values)
     next_i, next_j = iterate(allind), iterate(ind)
@@ -161,6 +157,3 @@ function compute_truncerr(values::AbstractVector{<:Number}, ind::AbstractVector{
 
     return sqrt(init)
 end
-
-# generic fallback: no allocations but inaccurate
-compute_truncerr(values::AbstractVector, ind) = sqrt(norm(values)^2 - norm(view(values, ind))^2)

src/interface/truncation.jl

@@ -181,3 +181,9 @@ Base.:&(::NoTruncation, ::NoTruncation) = notrunc()
 # disambiguate
 Base.:&(::NoTruncation, trunc::TruncationIntersection) = trunc
 Base.:&(trunc::TruncationIntersection, ::NoTruncation) = trunc
+
+@doc """
+    truncation_error(values, ind)
+
+Compute the truncation error as the 2-norm of the values that are not kept by `ind`.
+""" truncation_error