Better handling of TruncatedAlgorithm in select_algorithm

Author: mtfishman

src/interface/eig.jl

@@ -100,6 +100,12 @@ for f in (:eig_full!, :eig_vals!)
 end
 
 function select_algorithm(::typeof(eig_trunc!), A, alg; trunc=nothing, kwargs...)
-    alg_eig = select_algorithm(eig_full!, A, alg; kwargs...)
-    return TruncatedAlgorithm(alg_eig, select_truncation(trunc))
+    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

@@ -101,6 +101,12 @@ for f in (:eigh_full!, :eigh_vals!)
 end
 
 function select_algorithm(::typeof(eigh_trunc!), A, alg; trunc=nothing, kwargs...)
-    alg_eigh = select_algorithm(eigh_full!, A, alg; kwargs...)
-    return TruncatedAlgorithm(alg_eigh, select_truncation(trunc))
+    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/svd.jl

@@ -105,6 +105,12 @@ for f in (:svd_full!, :svd_compact!, :svd_vals!)
 end
 
 function select_algorithm(::typeof(svd_trunc!), A, alg; trunc=nothing, kwargs...)
-    alg_svd = select_algorithm(svd_compact!, A, alg; kwargs...)
-    return TruncatedAlgorithm(alg_svd, select_truncation(trunc))
+    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

@@ -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

@@ -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 = qr_compact(randn(rng, T, m, m))[1]
+    D = Diagonal([0.9, 0.3, 0.1, 0.01])
+    A = V * D * 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

test/eigh.jl

@@ -62,3 +62,19 @@ end
         @test V2 * (V2' * V1) ≈ V1
     end
 end
+
+@testset "eigh_trunc! specify truncation algorithm T = $T" for T in
+                                                               (Float32, Float64,
+                                                                ComplexF32,
+                                                                ComplexF64)
+    rng = StableRNG(123)
+    m = 4
+    V = qr_compact(randn(rng, T, m, m))[1]
+    D = Diagonal([0.9, 0.3, 0.1, 0.01])
+    A = V * D * V'
+    A = (A + A') / 2
+    alg = TruncatedAlgorithm(LAPACK_QRIteration(), truncrank(2))
+    D2, V2 = @constinferred eigh_trunc(A; alg)
+    @test diagview(D2) ≈ diagview(D)[1:2] rtol = sqrt(eps(real(T)))
+    @test_throws ArgumentError eigh_trunc(A; alg, trunc=(; maxrank=2))
+end

test/svd.jl

@@ -152,3 +152,18 @@ end
         end
     end
 end
+
+@testset "svd_trunc! specify truncation algorithm T = $T" for T in
+                                                              (Float32, Float64, ComplexF32,
+                                                               ComplexF64)
+    rng = StableRNG(123)
+    m = 4
+    U = qr_compact(randn(rng, T, m, m))[1]
+    S = Diagonal([0.9, 0.3, 0.1, 0.01])
+    Vᴴ = qr_compact(randn(rng, T, m, m))[1]
+    A = U * S * Vᴴ
+    alg = TruncatedAlgorithm(LAPACK_DivideAndConquer(), TruncationKeepAbove(0.2, 0.0))
+    U2, S2, V2ᴴ = @constinferred svd_trunc(A; alg)
+    @test diagview(S2) ≈ diagview(S)[1:2] rtol = sqrt(eps(real(T)))
+    @test_throws ArgumentError svd_trunc(A; alg, trunc=(; maxrank=2))
+end