Merge branch 'main' into mf/fillarraysext

Author: mtfishman

src/algorithms.jl

@@ -77,9 +77,6 @@ passed as the third positional argument in the form of a `NamedTuple`.
 """ select_algorithm
 
 function select_algorithm(f::F, A, alg::Alg=nothing; kwargs...) where {F,Alg}
-    return select_algorithm(f, typeof(A), alg; kwargs...)
-end
-function select_algorithm(f::F, ::Type{A}, alg::Alg=nothing; kwargs...) where {F,A,Alg}
     if isnothing(alg)
         return default_algorithm(f, A; kwargs...)
     elseif alg isa Symbol
@@ -193,10 +190,24 @@ macro functiondef(f)
         end
 
         # define fallbacks for algorithm selection
-        @inline function select_algorithm(::typeof($f), ::Type{A}, alg::Alg;
-                                          kwargs...) where {Alg,A}
+        @inline function select_algorithm(::typeof($f), A, alg::Alg; kwargs...) where {Alg}
             return select_algorithm($f!, A, alg; kwargs...)
         end
+        # define default algorithm fallbacks for out-of-place functions
+        # in terms of the corresponding in-place function
+        @inline function default_algorithm(::typeof($f), A; kwargs...)
+            return default_algorithm($f!, A; kwargs...)
+        end
+        # define default algorithm fallbacks for out-of-place functions
+        # in terms of the corresponding in-place function for types,
+        # in principle this is covered by the definition above but
+        # it is necessary to avoid ambiguity errors with the generic definitions:
+        # ```julia
+        # default_algorithm(f::F, A; kwargs...) where {F} = default_algorithm(f, typeof(A); kwargs...)
+        # function default_algorithm(f::F, ::Type{T}; kwargs...) where {F,T}
+        #     throw(MethodError(default_algorithm, (f, T)))
+        # end
+        # ```
         @inline function default_algorithm(::typeof($f), ::Type{A}; kwargs...) where {A}
             return default_algorithm($f!, A; kwargs...)
         end

src/implementations/truncation.jl

@@ -67,22 +67,30 @@ struct TruncationKeepFiltered{F} <: TruncationStrategy
     filter::F
 end
 
-struct TruncationKeepAbove{T<:Real} <: TruncationStrategy
+struct TruncationKeepAbove{T<:Real,F} <: TruncationStrategy
     atol::T
     rtol::T
     p::Int
+    by::F
+end
+function TruncationKeepAbove(; atol::Real, rtol::Real, p::Int=2, by=abs)
+    return TruncationKeepAbove(atol, rtol, p, by)
 end
-function TruncationKeepAbove(atol::Real, rtol::Real, p::Int=2)
-    return TruncationKeepAbove(promote(atol, rtol)..., p)
+function TruncationKeepAbove(atol::Real, rtol::Real, p::Int=2, by=abs)
+    return TruncationKeepAbove(promote(atol, rtol)..., p, by)
 end
 
-struct TruncationKeepBelow{T<:Real} <: TruncationStrategy
+struct TruncationKeepBelow{T<:Real,F} <: TruncationStrategy
     atol::T
     rtol::T
     p::Int
+    by::F
+end
+function TruncationKeepBelow(; atol::Real, rtol::Real, p::Int=2, by=abs)
+    return TruncationKeepBelow(atol, rtol, p, by)
 end
-function TruncationKeepBelow(atol::Real, rtol::Real, p::Int=2)
-    return TruncationKeepBelow(promote(atol, rtol)..., p)
+function TruncationKeepBelow(atol::Real, rtol::Real, p::Int=2, by=abs)
+    return TruncationKeepBelow(promote(atol, rtol)..., p, by)
 end
 
 # TODO: better names for these functions of the above types
@@ -94,18 +102,18 @@ Truncation strategy to keep the first `howmany` values when sorted according to
 truncrank(howmany::Int; by=abs, rev=true) = TruncationKeepSorted(howmany, by, rev)
 
 """
-    trunctol(atol::Real)
+    trunctol(atol::Real; by=abs)
 
-Truncation strategy to discard the values that are smaller than `atol` in absolute value.
+Truncation strategy to discard the values that are smaller than `atol` according to `by`.
 """
-trunctol(atol) = TruncationKeepFiltered(≥(atol) ∘ abs)
+trunctol(atol; by=abs) = TruncationKeepFiltered(≥(atol) ∘ by)
 
 """
-    truncabove(atol::Real)
+    truncabove(atol::Real; by=abs)
 
-Truncation strategy to discard the values that are larger than `atol` in absolute value.
+Truncation strategy to discard the values that are larger than `atol` according to `by`.
 """
-truncabove(atol) = TruncationKeepFiltered(≤(atol) ∘ abs)
+truncabove(atol; by=abs) = TruncationKeepFiltered(≤(atol) ∘ by)
 
 """
     TruncationIntersection(trunc::TruncationStrategy, truncs::TruncationStrategy...)
@@ -177,17 +185,18 @@ Generic interface for finding truncated values of the spectrum of a decompositio
 based on the `strategy`. The output should be a collection of indices specifying
 which values to keep. `MatrixAlgebraKit.findtruncated` is used inside of the default
 implementation of [`truncate!`](@ref) to perform the truncation. It does not assume that the
-values are sorted. For a version that assumes the values are reverse sorted by
-absolute value (which is the standard case for SVD) see
-[`MatrixAlgebraKit.findtruncated_sorted`](@ref).
+values are sorted. For a version that assumes the values are reverse sorted (which is the
+standard case for SVD) see [`MatrixAlgebraKit.findtruncated_sorted`](@ref).
 """ findtruncated
 
 @doc """
     MatrixAlgebraKit.findtruncated_sorted(values::AbstractVector, strategy::TruncationStrategy)
 
-Like [`MatrixAlgebraKit.findtruncated`](@ref) but assumes that the values are sorted in reverse order by
-absolute value. However, note that this assumption is not checked, so passing values that are not sorted
-in that way can silently give unexpected results. This is used in the default implementation of
+Like [`MatrixAlgebraKit.findtruncated`](@ref) but assumes that the values are sorted in reverse order.
+They are assumed to be sorted in a way that is consistent with the truncation strategy,
+which generally means they are sorted by absolute value but some truncation strategies allow
+customizing that. However, note that this assumption is not checked, so passing values that are not sorted
+in the correct way can silently give unexpected results. This is used in the default implementation of
 [`svd_trunc!`](@ref).
 """ findtruncated_sorted
 
@@ -212,21 +221,21 @@ end
 
 function findtruncated(values::AbstractVector, strategy::TruncationKeepBelow)
     atol = max(strategy.atol, strategy.rtol * norm(values, strategy.p))
-    return findall(≤(atol), values)
+    return findall(≤(atol) ∘ strategy.by, values)
 end
 function findtruncated_sorted(values::AbstractVector, strategy::TruncationKeepBelow)
     atol = max(strategy.atol, strategy.rtol * norm(values, strategy.p))
-    i = searchsortedfirst(values, atol; by=abs, rev=true)
+    i = searchsortedfirst(values, atol; by=strategy.by, rev=true)
     return i:length(values)
 end
 
 function findtruncated(values::AbstractVector, strategy::TruncationKeepAbove)
     atol = max(strategy.atol, strategy.rtol * norm(values, strategy.p))
-    return findall(≥(atol), values)
+    return findall(≥(atol) ∘ strategy.by, values)
 end
 function findtruncated_sorted(values::AbstractVector, strategy::TruncationKeepAbove)
     atol = max(strategy.atol, strategy.rtol * norm(values, strategy.p))
-    i = searchsortedlast(values, atol; by=abs, rev=true)
+    i = searchsortedlast(values, atol; by=strategy.by, rev=true)
     return 1:i
 end
 

src/interface/eig.jl

@@ -99,8 +99,13 @@ for f in (:eig_full!, :eig_vals!)
     end
 end
 
-function select_algorithm(::typeof(eig_trunc!), ::Type{A}, alg; trunc=nothing,
-                          kwargs...) where {A<:YALAPACK.BlasMat}
-    alg_eig = select_algorithm(eig_full!, A, alg; kwargs...)
-    return TruncatedAlgorithm(alg_eig, select_truncation(trunc))
+function select_algorithm(::typeof(eig_trunc!), A, alg; trunc=nothing, kwargs...)
+    if alg isa TruncatedAlgorithm
+        isnothing(trunc) ||
+            throw(ArgumentError("`trunc` can't be specified when `alg` is a `TruncatedAlgorithm`"))
+        return alg
+    else
+        alg_eig = select_algorithm(eig_full!, A, alg; kwargs...)
+        return TruncatedAlgorithm(alg_eig, select_truncation(trunc))
+    end
 end

src/interface/eigh.jl

@@ -100,8 +100,13 @@ for f in (:eigh_full!, :eigh_vals!)
     end
 end
 
-function select_algorithm(::typeof(eigh_trunc!), ::Type{A}, alg; trunc=nothing,
-                          kwargs...) where {A<:YALAPACK.BlasMat}
-    alg_eigh = select_algorithm(eigh_full!, A, alg; kwargs...)
-    return TruncatedAlgorithm(alg_eigh, select_truncation(trunc))
+function select_algorithm(::typeof(eigh_trunc!), A, alg; trunc=nothing, kwargs...)
+    if alg isa TruncatedAlgorithm
+        isnothing(trunc) ||
+            throw(ArgumentError("`trunc` can't be specified when `alg` is a `TruncatedAlgorithm`"))
+        return alg
+    else
+        alg_eig = select_algorithm(eigh_full!, A, alg; kwargs...)
+        return TruncatedAlgorithm(alg_eig, select_truncation(trunc))
+    end
 end

src/interface/lq.jl

@@ -77,9 +77,7 @@ function default_lq_algorithm(::Type{T}; kwargs...) where {T<:YALAPACK.BlasMat}
 end
 
 for f in (:lq_full!, :lq_compact!, :lq_null!)
-    @eval begin
-        function default_algorithm(::typeof($f), ::Type{A}; kwargs...) where {A}
-            return default_lq_algorithm(A; kwargs...)
-        end
+    @eval function default_algorithm(::typeof($f), ::Type{A}; kwargs...) where {A}
+        return default_lq_algorithm(A; kwargs...)
     end
 end

src/interface/polar.jl

@@ -61,10 +61,7 @@ end
 # Algorithm selection
 # -------------------
 default_polar_algorithm(A; kwargs...) = default_polar_algorithm(typeof(A); kwargs...)
-function default_polar_algorithm(T::Type; kwargs...)
-    throw(MethodError(default_polar_algorithm, (T,)))
-end
-function default_polar_algorithm(::Type{T}; kwargs...) where {T<:YALAPACK.BlasMat}
+function default_polar_algorithm(::Type{T}; kwargs...) where {T}
     return PolarViaSVD(default_algorithm(svd_compact!, T; kwargs...))
 end
 

src/interface/qr.jl

@@ -77,10 +77,7 @@ function default_qr_algorithm(::Type{T}; kwargs...) where {T<:YALAPACK.BlasMat}
 end
 
 for f in (:qr_full!, :qr_compact!, :qr_null!)
-    @eval begin
-        function default_algorithm(::typeof($f), ::Type{A};
-                                   kwargs...) where {A<:YALAPACK.BlasMat}
-            return default_qr_algorithm(A; kwargs...)
-        end
+    @eval function default_algorithm(::typeof($f), ::Type{A}; kwargs...) where {A}
+        return default_qr_algorithm(A; kwargs...)
     end
 end

src/interface/svd.jl

@@ -104,8 +104,13 @@ for f in (:svd_full!, :svd_compact!, :svd_vals!)
     end
 end
 
-function select_algorithm(::typeof(svd_trunc!), ::Type{A}, alg; trunc=nothing,
-                          kwargs...) where {A<:YALAPACK.BlasMat}
-    alg_svd = select_algorithm(svd_compact!, A, alg; kwargs...)
-    return TruncatedAlgorithm(alg_svd, select_truncation(trunc))
+function select_algorithm(::typeof(svd_trunc!), A, alg; trunc=nothing, kwargs...)
+    if alg isa TruncatedAlgorithm
+        isnothing(trunc) ||
+            throw(ArgumentError("`trunc` can't be specified when `alg` is a `TruncatedAlgorithm`"))
+        return alg
+    else
+        alg_svd = select_algorithm(svd_compact!, A, alg; kwargs...)
+        return TruncatedAlgorithm(alg_svd, select_truncation(trunc))
+    end
 end

test/algorithms.jl

@@ -2,7 +2,7 @@ using MatrixAlgebraKit
 using Test
 using TestExtras
 using MatrixAlgebraKit: LAPACK_SVDAlgorithm, NoTruncation, PolarViaSVD, TruncatedAlgorithm,
-                        default_algorithm, select_algorithm
+                        TruncationKeepBelow, default_algorithm, select_algorithm
 
 @testset "default_algorithm" begin
     A = randn(3, 3)
@@ -50,6 +50,12 @@ end
                                  NoTruncation())
     end
 
+    alg = TruncatedAlgorithm(LAPACK_Simple(), TruncationKeepBelow(0.1, 0.0))
+    for f in (eig_trunc!, eigh_trunc!, svd_trunc!)
+        @test @constinferred(select_algorithm(eig_trunc!, A, alg)) === alg
+        @test_throws ArgumentError select_algorithm(eig_trunc!, A, alg; trunc=(; maxrank=2))
+    end
+
     @test @constinferred(select_algorithm(svd_compact!, A)) === LAPACK_DivideAndConquer()
     @test @constinferred(select_algorithm(svd_compact!, A, nothing)) ===
           LAPACK_DivideAndConquer()

test/eig.jl

@@ -3,7 +3,7 @@ using Test
 using TestExtras
 using StableRNGs
 using LinearAlgebra: Diagonal
-using MatrixAlgebraKit: diagview
+using MatrixAlgebraKit: TruncatedAlgorithm, diagview
 
 @testset "eig_full! for T = $T" for T in (Float32, Float64, ComplexF32, ComplexF64)
     rng = StableRNG(123)
@@ -57,3 +57,17 @@ end
         @test V2 * ((V2' * V2) \ (V2' * V1)) ≈ V1
     end
 end
+
+@testset "eig_trunc! specify truncation algorithm T = $T" for T in
+                                                              (Float32, Float64, ComplexF32,
+                                                               ComplexF64)
+    rng = StableRNG(123)
+    m = 4
+    V = randn(rng, T, m, m)
+    D = Diagonal([0.9, 0.3, 0.1, 0.01])
+    A = V * D * inv(V)
+    alg = TruncatedAlgorithm(LAPACK_Simple(), truncrank(2))
+    D2, V2 = @constinferred eig_trunc(A; alg)
+    @test diagview(D2) ≈ diagview(D)[1:2] rtol = sqrt(eps(real(T)))
+    @test_throws ArgumentError eig_trunc(A; alg, trunc=(; maxrank=2))
+end