Add Diagonal eigh implementation and tests

Author: lkdvos

src/implementations/eigh.jl

@@ -8,6 +8,8 @@ function copy_input(::typeof(eigh_vals), A::AbstractMatrix)
 end
 copy_input(::typeof(eigh_trunc), A) = copy_input(eigh_full, A)
 
+copy_input(::typeof(eigh_full), A::Diagonal) = copy(A)
+
 function check_input(::typeof(eigh_full!), A::AbstractMatrix, DV, ::AbstractAlgorithm)
     m, n = size(A)
     m == n || throw(DimensionMismatch("square input matrix expected"))
@@ -21,6 +23,27 @@ function check_input(::typeof(eigh_full!), A::AbstractMatrix, DV, ::AbstractAlgo
 end
 function check_input(::typeof(eigh_vals!), A::AbstractMatrix, D, ::AbstractAlgorithm)
     m, n = size(A)
+    m == n || throw(DimensionMismatch("square input matrix expected"))
+    @assert D isa AbstractVector
+    @check_size(D, (n,))
+    @check_scalar(D, A, real)
+    return nothing
+end
+
+function check_input(::typeof(eigh_full!), A::AbstractMatrix, DV, ::DiagonalAlgorithm)
+    m, n = size(A)
+    @assert m == n && isdiag(A)
+    D, V = DV
+    @assert D isa Diagonal && V isa Diagonal
+    @check_size(D, (m, m))
+    @check_scalar(D, A, real)
+    @check_size(V, (m, m))
+    @check_scalar(V, A)
+    return nothing
+end
+function check_input(::typeof(eigh_vals!), A::AbstractMatrix, D, ::DiagonalAlgorithm)
+    m, n = size(A)
+    @assert m == n && isdiag(A)
     @assert D isa AbstractVector
     @check_size(D, (n,))
     @check_scalar(D, A, real)
@@ -45,6 +68,13 @@ function initialize_output(::typeof(eigh_trunc!), A::AbstractMatrix,
     return initialize_output(eigh_full!, A, alg.alg)
 end
 
+function initialize_output(::typeof(eigh_full!), A::Diagonal, ::DiagonalAlgorithm)
+    return eltype(A) <: Real ? A : similar(A, real(eltype(A))), similar(A)
+end
+function initialize_output(::typeof(eigh_vals!), A::Diagonal, ::DiagonalAlgorithm)
+    return eltype(A) <: Real ? diagview(A) : similar(A, real(eltype(A)), size(A, 1))
+end
+
 # Implementation
 # --------------
 function eigh_full!(A::AbstractMatrix, DV, alg::LAPACK_EighAlgorithm)
@@ -85,6 +115,25 @@ function eigh_trunc!(A::AbstractMatrix, DV, alg::TruncatedAlgorithm)
     return truncate!(eigh_trunc!, (D, V), alg.trunc)
 end
 
+# Diagonal logic
+# --------------
+function eigh_full!(A::Diagonal, DV, alg::DiagonalAlgorithm)
+    check_input(eigh_full!, A, DV, alg)
+    D, V = DV
+    D === A || (diagview(D) .= real.(diagview(A)))
+    one!(V)
+    return D, V
+end
+
+function eigh_vals!(A::Diagonal, D, alg::DiagonalAlgorithm)
+    check_input(eigh_vals!, A, D, alg)
+    Ad = diagview(A)
+    D === Ad || (D .= real.(Ad))
+    return D
+end
+
+# GPU logic
+# ---------
 _gpu_heevj!(A::AbstractMatrix, Dd::AbstractVector, V::AbstractMatrix; kwargs...) = throw(MethodError(_gpu_heevj!, (A, Dd, V)))
 _gpu_heevd!(A::AbstractMatrix, Dd::AbstractVector, V::AbstractMatrix; kwargs...) = throw(MethodError(_gpu_heevd!, (A, Dd, V)))
 _gpu_heev!(A::AbstractMatrix, Dd::AbstractVector, V::AbstractMatrix; kwargs...) = throw(MethodError(_gpu_heev!, (A, Dd, V)))

src/interface/eigh.jl

@@ -93,6 +93,9 @@ end
 function default_eigh_algorithm(::Type{T}; kwargs...) where {T<:YALAPACK.BlasMat}
     return LAPACK_MultipleRelativelyRobustRepresentations(; kwargs...)
 end
+function default_eigh_algorithm(::Type{T}; kwargs...) where {T<:Diagonal}
+    return DiagonalAlgorithm(; kwargs...)
+end
 
 for f in (:eigh_full!, :eigh_vals!)
     @eval function default_algorithm(::typeof($f), ::Type{A}; kwargs...) where {A}

test/eigh.jl

@@ -5,7 +5,9 @@ using StableRNGs
 using LinearAlgebra: LinearAlgebra, Diagonal, I
 using MatrixAlgebraKit: TruncatedAlgorithm, diagview
 
-@testset "eigh_full! for T = $T" for T in (Float32, Float64, ComplexF32, ComplexF64)
+const BLASFloats = (Float32, Float64, ComplexF32, ComplexF64)
+
+@testset "eigh_full! for T = $T" for T in BLASFloats
     rng = StableRNG(123)
     m = 54
     for alg in (LAPACK_MultipleRelativelyRobustRepresentations(),
@@ -29,7 +31,7 @@ using MatrixAlgebraKit: TruncatedAlgorithm, diagview
     end
 end
 
-@testset "eigh_trunc! for T = $T" for T in (Float32, Float64, ComplexF32, ComplexF64)
+@testset "eigh_trunc! for T = $T" for T in BLASFloats
     rng = StableRNG(123)
     m = 54
     for alg in (LAPACK_QRIteration(),
@@ -62,10 +64,7 @@ end
     end
 end
 
-@testset "eigh_trunc! specify truncation algorithm T = $T" for T in
-                                                               (Float32, Float64,
-                                                                ComplexF32,
-                                                                ComplexF64)
+@testset "eigh_trunc! specify truncation algorithm T = $T" for T in BLASFloats
     rng = StableRNG(123)
     m = 4
     V = qr_compact(randn(rng, T, m, m))[1]
@@ -77,3 +76,25 @@ end
     @test diagview(D2) ≈ diagview(D)[1:2] rtol = sqrt(eps(real(T)))
     @test_throws ArgumentError eigh_trunc(A; alg, trunc=(; maxrank=2))
 end
+
+@testset "eigh for Diagonal{$T}" for T in BLASFloats
+    rng = StableRNG(123)
+    m = 54
+    Ad = randn(rng, T, m)
+    Ad .+= conj.(Ad)
+    A = Diagonal(Ad)
+
+    D, V = @constinferred eigh_full(A)
+    @test D isa Diagonal{real(T)} && size(D) == size(A)
+    @test V isa Diagonal{T} && size(V) == size(A)
+    @test A * V ≈ V * D
+
+    D2 = @constinferred eigh_vals(A)
+    @test D2 isa AbstractVector{real(T)} && length(D2) == m
+    @test diagview(D) ≈ D2
+
+    A2 = Diagonal(T[0.9, 0.3, 0.1, 0.01])
+    alg = TruncatedAlgorithm(DiagonalAlgorithm(), truncrank(2))
+    D2, V2 = @constinferred eigh_trunc(A2; alg)
+    @test diagview(D2) ≈ diagview(A2)[1:2]
+end