Add a new truncation strategy to protect degenerate clusters

docs/src/user_interface/truncations.md

@@ -145,6 +145,19 @@ combined_trunc = truncrank(10) & trunctol(; atol = 1e-6);
 combined_trunc = trunctol(; atol = 1e-6) | truncrank(3);
 ```
 
+Truncating with a fixed number of values can split a group of (near-)degenerate values, such as 
+when these degeneracies are due to a symmetry. Wrapping a strategy with [`truncmultiplets`](@ref)
+ensures that such groups are either kept or discarded as a whole cluster:
+
+```@docs; canonical=false
+truncmultiplets
+```
+
+```julia
+# Keep at most 10 values, without splitting degenerate groups
+degenerate_trunc = truncmultiplets(truncrank(10); rtol = 1e-10);
+```
+
 ## Truncation Error
 
 When using truncated decompositions such as [`svd_trunc`](@ref), [`eig_trunc`](@ref), or [`eigh_trunc`](@ref), an additional truncation error value is returned.

src/MatrixAlgebraKit.jl

@@ -52,7 +52,7 @@ export CUSOLVER_Simple, CUSOLVER_HouseholderQR, CUSOLVER_QRIteration, CUSOLVER_S
 export ROCSOLVER_HouseholderQR, ROCSOLVER_QRIteration, ROCSOLVER_Jacobi,
     ROCSOLVER_DivideAndConquer, ROCSOLVER_Bisection
 
-export notrunc, truncrank, trunctol, truncerror, truncfilter
+export notrunc, truncrank, trunctol, truncerror, truncfilter, truncmultiplets
 
 @static if VERSION >= v"1.11.0-DEV.469"
     eval(
@@ -64,7 +64,8 @@ export notrunc, truncrank, trunctol, truncerror, truncfilter
     eval(
         Expr(
             :public, :TruncationByOrder, :TruncationByFilter, :TruncationByValue,
-            :TruncationByError, :TruncationIntersection, :TruncationUnion, :truncate
+            :TruncationByError, :TruncationIntersection, :TruncationUnion, :TruncationMultiplets,
+            :truncate
         )
     )
     eval(

src/implementations/truncation.jl

@@ -163,6 +163,40 @@ _ind_union(A::Base.OneTo, B::Base.OneTo) = Base.OneTo(max(length(A), length(B)))
 _ind_union(A::AbstractUnitRange, B::AbstractUnitRange) = union(A, B)
 _ind_union(A, B) = union(A, B)
 
+function findtruncated(values::AbstractVector, strategy::TruncationMultiplets)
+    keep = _ind_to_mask(values, findtruncated(values, strategy.trunc))
+    perm = sortperm(values; by = abs, rev = true)
+    return findall(_keep_multiplets!(keep, values, perm, strategy))
+end
+function findtruncated_svd(values::AbstractVector, strategy::TruncationMultiplets)
+    # values are already sorted by decreasing magnitude
+    keep = _ind_to_mask(values, findtruncated_svd(values, strategy.trunc))
+    return findall(_keep_multiplets!(keep, values, eachindex(values), strategy))
+end
+
+_ind_to_mask(values::AbstractVector, ::Colon) = trues(length(values))
+_ind_to_mask(values::AbstractVector, ind::AbstractVector{Bool}) = BitVector(ind)
+function _ind_to_mask(values::AbstractVector, ind)
+    keep = falses(length(values))
+    keep[ind] .= true
+    return keep
+end
+
+# walk the values in sorted order `perm`, and make sure each group of consecutive degenerate
+# values is either kept or discarded as a whole cluster
+function _keep_multiplets!(keep::AbstractVector{Bool}, values::AbstractVector, perm, strategy::TruncationMultiplets)
+    σs = abs.(collect(view(values, perm)))
+    start = firstindex(σs)
+    for k in eachindex(σs)
+        if k == lastindex(σs) || σs[k] - σs[k + 1] > max(strategy.atol, strategy.rtol * σs[k])
+            group = view(perm, start:k)
+            allequal(view(keep, group)) || (keep[group] .= strategy.mode == :up)
+            start = k + 1
+        end
+    end
+    return keep
+end
+
 # Truncation error
 # ----------------
 truncation_error(values::AbstractVector, ind) = truncation_error!(copy(values), ind)

src/interface/truncation.jl

@@ -270,6 +270,45 @@ Base.:|(::NoTruncation, ::NoTruncation) = notrunc()
 Base.:|(::NoTruncation, ::TruncationUnion) = notrunc()
 Base.:|(::TruncationUnion, ::NoTruncation) = notrunc()
 
+"""
+    TruncationMultiplets(trunc::TruncationStrategy, atol::Real, rtol::Real, mode::Symbol)
+
+Truncation strategy that applies `trunc` and then adjusts the selection such that groups of
+(near-)degenerate values are never split.
+
+See also [`truncmultiplets`](@ref).
+"""
+struct TruncationMultiplets{S <: TruncationStrategy, T <: Real} <: TruncationStrategy
+    trunc::S
+    atol::T
+    rtol::T
+    mode::Symbol
+end
+function TruncationMultiplets(trunc::TruncationStrategy, atol::Real, rtol::Real, mode::Symbol)
+    return TruncationMultiplets(trunc, promote(atol, rtol)..., mode)
+end
+
+"""
+    truncmultiplets(strategy::TruncationStrategy; atol::Real=0, rtol::Real=1e-10, mode::Symbol=:down)
+
+Truncation strategy that applies `strategy`, and then makes sure that groups of (near-)degenerate
+values are either kept or discarded as a whole.
+
+Values are sorted by decreasing magnitude, and two consecutive values `abs(vᵢ) ≥ abs(vᵢ₊₁)` are
+considered degenerate whenever `abs(vᵢ) - abs(vᵢ₊₁) ≤ max(atol, rtol * abs(vᵢ))`. Groups of
+degenerate values that are only partially kept by `strategy` are discarded entirely if
+`mode = :down`, or kept entirely if `mode = :up`.
+
+This is useful when the degeneracies are due to a symmetry, where splitting a
+degenerate group would break the symmetry due to rounding errors.
+In particular, `truncmultiplets(truncrank(howmany))` keeps at most `howmany` values for
+`mode = :down`, and at least `min(howmany, length(values))` values for `mode = :up`.
+"""
+function truncmultiplets(strategy::TruncationStrategy; atol::Real = 0, rtol::Real = 1.0e-10, mode::Symbol = :down)
+    mode in (:down, :up) || throw(ArgumentError(lazy"mode must be `:down` or `:up`, got `$mode`"))
+    return TruncationMultiplets(strategy, atol, rtol, mode)
+end
+
 @doc """
     truncation_error(values, ind)
 Compute the truncation error as the 2-norm of the values that are not kept by `ind`.

test/common/truncate.jl

@@ -2,7 +2,9 @@ using MatrixAlgebraKit
 using Test
 using TestExtras
 using MatrixAlgebraKit: NoTruncation, TruncationIntersection, TruncationUnion,
-    TruncationByOrder, TruncationByValue, TruncationStrategy, findtruncated, findtruncated_svd
+    TruncationByOrder, TruncationByValue, TruncationMultiplets, TruncationStrategy,
+    findtruncated, findtruncated_svd
+using LinearAlgebra: Diagonal, norm
 
 @testset "truncate" begin
     trunc = @constinferred TruncationStrategy()
@@ -94,3 +96,78 @@ using MatrixAlgebraKit: NoTruncation, TruncationIntersection, TruncationUnion,
     @test union2 isa TruncationUnion
     @test length(union2.components) == 3
 end
+
+@testset "truncmultiplets" begin
+    trunc = @constinferred truncmultiplets(truncrank(2))
+    @test trunc isa TruncationMultiplets
+    @test trunc.trunc == truncrank(2)
+    @test trunc.atol == 0 && trunc.rtol == 1.0e-10
+    @test trunc.mode === :down
+    @test trunc == truncmultiplets(truncrank(2))
+    @test trunc != truncmultiplets(truncrank(2); mode = :up)
+    @test_throws ArgumentError truncmultiplets(truncrank(2); mode = :sideways)
+
+    # a degenerate group of three values around 0.5, split by truncrank(2) and truncrank(3)
+    ϵ = 1.0e-14
+    vals_sorted = [1.0, 0.5 + ϵ, 0.5, -0.5 + ϵ, 0.1]
+    vals_unsorted = vals_sorted[[3, 5, 1, 4, 2]]
+    for howmany in (2, 3)
+        down = truncmultiplets(truncrank(howmany))
+        up = truncmultiplets(truncrank(howmany); mode = :up)
+        @test @constinferred(findtruncated_svd(vals_sorted, down)) == [1]
+        @test @constinferred(findtruncated_svd(vals_sorted, up)) == 1:4
+        @test sort(vals_unsorted[@constinferred(findtruncated(vals_unsorted, down))]) == [1.0]
+        @test sort(vals_unsorted[@constinferred(findtruncated(vals_unsorted, up))]) ==
+            sort(vals_sorted[1:4])
+    end
+
+    # nothing changes if the cut does not split a degenerate cluster
+    for howmany in (0, 1, 4, 5, 6)
+        for mode in (:down, :up)
+            strategy = truncmultiplets(truncrank(howmany); mode)
+            @test findtruncated_svd(vals_sorted, strategy) ==
+                findtruncated_svd(vals_sorted, truncrank(howmany))
+            @test sort(findtruncated(vals_unsorted, strategy)) ==
+                sort(findtruncated(vals_unsorted, truncrank(howmany)))
+        end
+    end
+    @test findtruncated_svd(vals_sorted, truncmultiplets(notrunc())) == 1:5
+    @test isempty(findtruncated_svd(Float64[], truncmultiplets(truncrank(2))))
+
+    # without tolerance, only exactly degenerate values are grouped
+    @test findtruncated_svd(vals_sorted, truncmultiplets(truncrank(2); rtol = 0)) == 1:2
+    @test findtruncated_svd([1.0, 0.5, 0.5, 0.1], truncmultiplets(truncrank(2); rtol = 0)) == [1]
+
+    vals_chain = [1.0, 0.9, 0.8, 0.7, 0.1]
+    @test findtruncated_svd(vals_chain, truncmultiplets(truncrank(2); atol = 0.15)) == Int[]
+    @test findtruncated_svd(vals_chain, truncmultiplets(truncrank(2); atol = 0.15, mode = :up)) == 1:4
+    @test findtruncated_svd(vals_chain, truncmultiplets(truncrank(2); rtol = 0.12)) == Int[]
+    @test findtruncated_svd(vals_chain, truncmultiplets(truncrank(2); rtol = 0.12, mode = :up)) == 1:3
+
+    strategy = truncmultiplets(truncrank(3) & trunctol(; atol = 0.2))
+    @test findtruncated_svd(vals_sorted, strategy) == [1]
+    strategy = truncmultiplets(trunctol(; atol = 0.5) | truncrank(1); mode = :up)
+    @test findtruncated_svd(vals_sorted, strategy) == 1:4
+    strategy = truncmultiplets(truncfilter(x -> x < 0.5 + ϵ / 2); mode = :up)
+    @test sort(findtruncated(vals_unsorted, strategy)) == sort(findtruncated(vals_unsorted, truncfilter(x -> x < 0.9)))
+    @test truncmultiplets(truncrank(2)) & trunctol(; atol = 0.2) isa TruncationIntersection
+
+    # exactly degenerate singular values
+    for T in (Float64, ComplexF64)
+        n = 6
+        Q₁ = qr_compact(randn(T, n, n))[1]
+        Q₂ = qr_compact(randn(T, n, n))[1]
+        Sref = [3.0, 2.0, 2.0, 2.0, 1.0, 0.5]
+        A = Q₁ * Diagonal(Sref) * Q₂'
+        for howmany in 2:3
+            U, S, Vᴴ, ϵ = svd_trunc(A; trunc = truncmultiplets(truncrank(howmany); rtol = 1.0e-8))
+            @test diagview(S) ≈ Sref[1:1]
+            @test ϵ ≈ norm(Sref[2:end])
+            @test U * S * Vᴴ ≈ Sref[1] * Q₁[:, 1] * Q₂[:, 1]'
+            U, S, Vᴴ, ϵ = svd_trunc(A; trunc = truncmultiplets(truncrank(howmany); rtol = 1.0e-8, mode = :up))
+            @test diagview(S) ≈ Sref[1:4]
+            @test ϵ ≈ norm(Sref[5:end])
+            @test U * S * Vᴴ ≈ Q₁[:, 1:4] * Diagonal(Sref[1:4]) * Q₂[:, 1:4]'
+        end
+    end
+end

test/testsuite/decompositions/svd.jl

@@ -686,6 +686,21 @@ function test_svd_trunc_algs(
             @test ϵ2 ≈ norm(diagview(S)[3:4]) atol = atol
             @test_throws ArgumentError svd_trunc(A; alg = truncalg, trunc = (; maxrank = 2))
             @test_throws ArgumentError svd_trunc_no_error(A; alg = truncalg, trunc = (; maxrank = 2))
+
+            # truncmultiplets
+            Sdiag = similar(A, real(eltype(T)), m4)
+            copyto!(Sdiag, real(eltype(T))[0.9, 0.3, 0.3, 0.01])
+            S = Diagonal(Sdiag)
+            Vᴴ = instantiate_unitary(T, A, m4)
+            A = U * S * Vᴴ
+            truncalg = TruncatedAlgorithm(alg, truncmultiplets(truncrank(2); atol = 0.2))
+            U2, S2, V2ᴴ, ϵ2 = @testinferred svd_trunc(A; alg = truncalg)
+            @test collect(diagview(S2)) ≈ collect(diagview(S)[1:1])
+            @test ϵ2 ≈ norm(diagview(S)[2:4]) atol = atol
+            truncalg = TruncatedAlgorithm(alg, truncmultiplets(truncrank(3); atol = 0.2))
+            U2, S2, V2ᴴ, ϵ2 = @testinferred svd_trunc(A; alg = truncalg)
+            @test collect(diagview(S2)) ≈ collect(diagview(S)[1:3])
+            @test ϵ2 ≈ norm(diagview(S)[4:4]) atol = atol
         end
     end
 end