|
2 | 2 | # --------- |
3 | 3 | # Generic implementation: `findtruncated` followed by indexing |
4 | 4 | function truncate!(::typeof(svd_trunc!), (U, S, Vᴴ), strategy::TruncationStrategy) |
5 | | - ind = findtruncated_sorted(diagview(S), strategy) |
| 5 | + ind = findtruncated_svd(diagview(S), strategy) |
6 | 6 | return U[:, ind], Diagonal(diagview(S)[ind]), Vᴴ[ind, :] |
7 | 7 | end |
8 | 8 | function truncate!(::typeof(eig_trunc!), (D, V), strategy::TruncationStrategy) |
|
28 | 28 |
|
29 | 29 | # findtruncated |
30 | 30 | # ------------- |
| 31 | +# Generic fallback |
| 32 | +function findtruncated_svd(values::AbstractVector, strategy::TruncationStrategy) |
| 33 | + return findtruncated(values, strategy) |
| 34 | +end |
| 35 | + |
31 | 36 | # specific implementations for finding truncated values |
32 | 37 | findtruncated(values::AbstractVector, ::NoTruncation) = Colon() |
33 | 38 |
|
34 | | -function findtruncated(values::AbstractVector, strategy::TruncationKeepSorted) |
| 39 | +function findtruncated(values::AbstractVector, strategy::TruncationByOrder) |
35 | 40 | howmany = min(strategy.howmany, length(values)) |
36 | | - return partialsortperm(values, 1:howmany; by=strategy.by, rev=strategy.rev) |
| 41 | + return partialsortperm(values, 1:howmany; strategy.by, strategy.rev) |
37 | 42 | end |
38 | | -function findtruncated_sorted(values::AbstractVector, strategy::TruncationKeepSorted) |
| 43 | +function findtruncated_svd(values::AbstractVector, strategy::TruncationByOrder) |
| 44 | + strategy.by === abs || return findtruncated(values, strategy) |
39 | 45 | howmany = min(strategy.howmany, length(values)) |
40 | | - return 1:howmany |
| 46 | + return strategy.rev ? (1:howmany) : ((length(values) - howmany + 1):length(values)) |
41 | 47 | end |
42 | 48 |
|
43 | | -# TODO: consider if worth using that values are sorted when filter is `<` or `>`. |
44 | | -function findtruncated(values::AbstractVector, strategy::TruncationKeepFiltered) |
45 | | - ind = findall(strategy.filter, values) |
46 | | - return ind |
| 49 | +function findtruncated(values::AbstractVector, strategy::TruncationByFilter) |
| 50 | + return findall(strategy.filter, values) |
47 | 51 | end |
48 | 52 |
|
49 | | -function findtruncated(values::AbstractVector, strategy::TruncationKeepBelow) |
| 53 | +function findtruncated(values::AbstractVector, strategy::TruncationByValue) |
50 | 54 | atol = max(strategy.atol, strategy.rtol * norm(values, strategy.p)) |
51 | | - return findall(≤(atol) ∘ strategy.by, values) |
| 55 | + filter = (strategy.keep_below ? ≤(atol) : ≥(atol)) ∘ strategy.by |
| 56 | + return findtruncated(values, truncfilter(filter)) |
52 | 57 | end |
53 | | -function findtruncated_sorted(values::AbstractVector, strategy::TruncationKeepBelow) |
| 58 | +function findtruncated_svd(values::AbstractVector, strategy::TruncationByValue) |
| 59 | + strategy.by === abs || return findtruncated(values, strategy) |
54 | 60 | atol = max(strategy.atol, strategy.rtol * norm(values, strategy.p)) |
55 | | - i = searchsortedfirst(values, atol; by=strategy.by, rev=true) |
56 | | - return i:length(values) |
57 | | -end |
58 | | - |
59 | | -function findtruncated(values::AbstractVector, strategy::TruncationKeepAbove) |
60 | | - atol = max(strategy.atol, strategy.rtol * norm(values, strategy.p)) |
61 | | - return findall(≥(atol) ∘ strategy.by, values) |
62 | | -end |
63 | | -function findtruncated_sorted(values::AbstractVector, strategy::TruncationKeepAbove) |
64 | | - atol = max(strategy.atol, strategy.rtol * norm(values, strategy.p)) |
65 | | - i = searchsortedlast(values, atol; by=strategy.by, rev=true) |
66 | | - return 1:i |
67 | | -end |
68 | | - |
69 | | -function findtruncated(values::AbstractVector, strategy::TruncationIntersection) |
70 | | - inds = map(Base.Fix1(findtruncated, values), strategy.components) |
71 | | - return intersect(inds...) |
72 | | -end |
73 | | -function findtruncated_sorted(values::AbstractVector, strategy::TruncationIntersection) |
74 | | - inds = map(Base.Fix1(findtruncated_sorted, values), strategy.components) |
75 | | - return intersect(inds...) |
| 61 | + if strategy.keep_below |
| 62 | + i = searchsortedfirst(values, atol; by=abs, rev=true) |
| 63 | + return i:length(values) |
| 64 | + else |
| 65 | + i = searchsortedlast(values, atol; by=abs, rev=true) |
| 66 | + return 1:i |
| 67 | + end |
76 | 68 | end |
77 | 69 |
|
78 | | -function findtruncated(values::AbstractVector, strategy::TruncationError) |
| 70 | +function findtruncated(values::AbstractVector, strategy::TruncationByError) |
79 | 71 | I = sortperm(values; by=abs, rev=true) |
80 | | - I′ = _truncerr_impl(values, I, strategy) |
| 72 | + I′ = _truncerr_impl(values, I; strategy.atol, strategy.rtol, strategy.p) |
81 | 73 | return I[I′] |
82 | 74 | end |
83 | | -function findtruncated_sorted(values::AbstractVector, strategy::TruncationError) |
| 75 | +function findtruncated_svd(values::AbstractVector, strategy::TruncationByError) |
84 | 76 | I = eachindex(values) |
85 | | - I′ = _truncerr_impl(values, I, strategy) |
| 77 | + I′ = _truncerr_impl(values, I; strategy.atol, strategy.rtol, strategy.p) |
86 | 78 | return I[I′] |
87 | 79 | end |
88 | | -function _truncerr_impl(values::AbstractVector, I, strategy::TruncationError) |
89 | | - Nᵖ = sum(Base.Fix2(^, strategy.p) ∘ abs, values) |
90 | | - ϵᵖ = max(strategy.atol^strategy.p, strategy.rtol^strategy.p * Nᵖ) |
| 80 | +function _truncerr_impl(values::AbstractVector, I; atol::Real=0, rtol::Real=0, p::Real=2) |
| 81 | + by = Base.Fix2(^, p) ∘ abs |
| 82 | + Nᵖ = sum(by, values) |
| 83 | + ϵᵖ = max(atol^p, rtol^p * Nᵖ) |
| 84 | + |
| 85 | + # fast path to avoid checking all values |
91 | 86 | ϵᵖ ≥ Nᵖ && return Base.OneTo(0) |
92 | 87 |
|
93 | 88 | truncerrᵖ = zero(real(eltype(values))) |
94 | 89 | rank = length(values) |
95 | 90 | for i in reverse(I) |
96 | | - truncerrᵖ += abs(values[i])^strategy.p |
97 | | - if truncerrᵖ ≥ ϵᵖ |
98 | | - break |
99 | | - else |
100 | | - rank -= 1 |
101 | | - end |
| 91 | + truncerrᵖ += by(values[i]) |
| 92 | + truncerrᵖ ≥ ϵᵖ && break |
| 93 | + rank -= 1 |
102 | 94 | end |
| 95 | + |
103 | 96 | return Base.OneTo(rank) |
104 | 97 | end |
105 | 98 |
|
106 | | -# Generic fallback |
107 | | -function findtruncated_sorted(values::AbstractVector, strategy::TruncationStrategy) |
108 | | - return findtruncated(values, strategy) |
| 99 | +function findtruncated(values::AbstractVector, strategy::TruncationIntersection) |
| 100 | + return mapreduce(Base.Fix1(findtruncated, values), _ind_intersect, strategy.components; |
| 101 | + init=trues(length(values))) |
109 | 102 | end |
| 103 | +function findtruncated_svd(values::AbstractVector, strategy::TruncationIntersection) |
| 104 | + return mapreduce(Base.Fix1(findtruncated_svd, values), _ind_intersect, |
| 105 | + strategy.components; init=trues(length(values))) |
| 106 | +end |
| 107 | + |
| 108 | +# when one of the ind selections is a bitvector, have to handle differently |
| 109 | +function _ind_intersect(A::AbstractVector{Bool}, B::AbstractVector) |
| 110 | + result = falses(length(A)) |
| 111 | + result[B] .= @view A[B] |
| 112 | + return result |
| 113 | +end |
| 114 | +_ind_intersect(A::AbstractVector, B::AbstractVector{Bool}) = _ind_intersect(B, A) |
| 115 | +_ind_intersect(A::AbstractVector{Bool}, B::AbstractVector{Bool}) = A .& B |
| 116 | +_ind_intersect(A, B) = intersect(A, B) |
0 commit comments