make eigvals(::Diagonal) accept sortby (#1477)

Solves the `Diagonal` part of #1471 and #1450.

I thought it would be better to start with `Diagonal`, as it was the
only part that required a bit of thinking. I think it also makes sense
in itself, since `eigen` accepted `sortby`, just `eigvals` and `eigvecs`
did not. Even though the docstring said none did.

If this is welcome I'll also do `SymTridiagonal`, and then we can
document that eigenvalues are sorted by default.

---------

Co-authored-by: Daniel Karrasch <[email protected]>

Author: araujoms

src/diagonal.jl

@@ -1011,16 +1011,16 @@ end
 _ortho_eltype(T) = Base.promote_op(/, T, T)
 _ortho_eltype(T::Type{<:Number}) = typeof(one(T)/one(T))
 
-# TODO Docstrings for eigvals, eigvecs, eigen all mention permute, scale, sortby as keyword args
-# but not all of them below provide them. Do we need to fix that?
 #Eigensystem
-eigvals(D::Diagonal{<:Number}; permute::Bool=true, scale::Bool=true) = copy(D.diag)
-eigvals(D::Diagonal; permute::Bool=true, scale::Bool=true) =
-    reduce(vcat, eigvals(x) for x in D.diag) #For block matrices, etc.
-function eigvecs(D::Diagonal{T}) where {T<:AbstractMatrix}
-    diag_vecs = [ eigvecs(x) for x in D.diag ]
+eigvals(D::Diagonal{<:Number}; permute::Bool=true, scale::Bool=true, sortby::Union{Function,Nothing}=nothing) = sorteig!(copy(D.diag), sortby)
+eigvals(D::Diagonal; permute::Bool=true, scale::Bool=true, sortby::Union{Function,Nothing}=nothing) =
+    sorteig!(reduce(vcat, eigvals(x; sortby=nothing) for x in D.diag), sortby) #For block matrices, etc.
+function _eigen(D::Diagonal{T}) where {T<:AbstractMatrix}
+    facts = [eigen(x; sortby=nothing) for x in D.diag]
+    λ = reduce(vcat, f.values for f in facts)
+    diag_vecs = [f.vectors for f in facts]
     matT = promote_type(map(typeof, diag_vecs)...)
-    ncols_diag = [ size(x, 2) for x in D.diag ]
+    ncols_diag = [size(x, 2) for x in D.diag]
     nrows = size(D, 1)
     vecs = Matrix{Vector{eltype(matT)}}(undef, nrows, sum(ncols_diag))
     for j in axes(D, 2), i in axes(D, 1)
@@ -1035,14 +1035,14 @@ function eigvecs(D::Diagonal{T}) where {T<:AbstractMatrix}
             end
         end
     end
-    return vecs
+    return λ, vecs
 end
 function eigen(D::Diagonal; permute::Bool=true, scale::Bool=true, sortby::Union{Function,Nothing}=nothing)
     if any(!isfinite, D.diag)
         throw(ArgumentError("matrix contains Infs or NaNs"))
     end
     Td = _ortho_eltype(eltype(D))
-    λ = eigvals(D)
+    λ = eigvals(D; sortby=nothing)
     if !isnothing(sortby)
         p = sortperm(λ; alg=QuickSort, by=sortby)
         λ = λ[p]
@@ -1059,8 +1059,7 @@ function eigen(D::Diagonal{<:AbstractMatrix}; permute::Bool=true, scale::Bool=tr
     if any(any(!isfinite, x) for x in D.diag)
         throw(ArgumentError("matrix contains Infs or NaNs"))
     end
-    λ = eigvals(D)
-    evecs = eigvecs(D)
+    λ, evecs = _eigen(D)
     if !isnothing(sortby)
         p = sortperm(λ; alg=QuickSort, by=sortby)
         λ = λ[p]

test/diagonal.jl

@@ -413,7 +413,7 @@ LinearAlgebra.istril(N::NotDiagonal) = istril(N.a)
     @test factorize(D) == D
 
     @testset "Eigensystem" begin
-        eigD = eigen(D)
+        eigD = eigen(D, sortby=nothing)
         @test Diagonal(eigD.values) == D
         @test eigD.vectors == Matrix(I, size(D))
         eigsortD = eigen(D, sortby=LinearAlgebra.eigsortby)
@@ -564,7 +564,7 @@ end
 @testset "svdvals and eigvals (#11120/#11247)" begin
     D = Diagonal(Matrix{Float64}[randn(3,3), randn(2,2)])
     @test sort([svdvals(D)...;], rev = true) ≈ svdvals([D.diag[1] zeros(3,2); zeros(2,3) D.diag[2]])
-    @test sort([eigvals(D)...;], by=LinearAlgebra.eigsortby) ≈ eigvals([D.diag[1] zeros(3,2); zeros(2,3) D.diag[2]])
+    @test eigvals(D, sortby=LinearAlgebra.eigsortby) ≈ eigvals([D.diag[1] zeros(3,2); zeros(2,3) D.diag[2]])
 end
 
 @testset "eigvals should return a copy of the diagonal" begin
@@ -853,7 +853,7 @@ end
 @testset "Eigensystem for block diagonal (issue #30681)" begin
     I2 = Matrix(I, 2,2)
     D = Diagonal([2.0*I2, 3.0*I2])
-    eigD = eigen(D)
+    eigD = eigen(D; sortby=LinearAlgebra.eigsortby)
     evals = [ 2.0, 2.0, 3.0, 3.0 ]
     evecs = [ [[ 1.0, 0.0 ]]  [[ 0.0, 1.0 ]]  [[ 0.0, 0.0 ]]  [[ 0.0, 0.0 ]];
               [[ 0.0, 0.0 ]]  [[ 0.0, 0.0 ]]  [[ 1.0, 0.0 ]]  [[ 0.0, 1.0 ]] ]
@@ -863,7 +863,7 @@ end
 
     I3 = Matrix(I, 3,3)
     D = Diagonal([[0.0 -1.0; 1.0 0.0], 2.0*I3])
-    eigD = eigen(D)
+    eigD = eigen(D; sortby=LinearAlgebra.eigsortby)
     evals = [ -1.0im, 1.0im, 2.0, 2.0, 2.0 ]
     evecs = [ [[ 1/sqrt(2)+0im, 1/sqrt(2)*im ]]  [[ 1/sqrt(2)+0im, -1/sqrt(2)*im ]]  [[ 0.0, 0.0 ]]       [[ 0.0, 0.0 ]]      [[ 0.0, 0.0]];
               [[ 0.0, 0.0, 0.0 ]]                [[ 0.0, 0.0, 0.0 ]]                 [[ 1.0, 0.0, 0.0 ]]  [[ 0.0, 1.0, 0.0 ]] [[ 0.0, 0.0, 1.0]] ]
@@ -1066,8 +1066,9 @@ end
 @testset "eigenvalue sorting" begin
     D = Diagonal([0.4, 0.2, -1.3])
     @test eigvals(D) == eigen(D).values == [0.4, 0.2, -1.3] # not sorted by default
+    @test eigvals(D; sortby=LinearAlgebra.eigsortby) == [-1.3, 0.2, 0.4] # sortby keyword supported for eigvals(::Diagonal)
     @test eigvals(Matrix(D)) == eigen(Matrix(D)).values == [-1.3, 0.2, 0.4] # sorted even if diagonal special case is detected
-    E = eigen(D, sortby=abs) # sortby keyword supported for eigen(::Diagonal)
+    E = eigen(D; sortby=abs) # sortby keyword supported for eigen(::Diagonal)
     @test E.values == [0.2, 0.4, -1.3]
     @test E.vectors == [0 1 0; 1 0 0; 0 0 1]
 end

test/unitful.jl

@@ -86,7 +86,7 @@ end
         @test U isa AbstractMatrix{<:Union{Real,Complex}}
         @test V isa AbstractMatrix{<:Union{Real,Complex}}
         @test s isa AbstractVector{<:Furlong{1}}
-        E = eigen(Du)
+        E = eigen(Du; sortby=nothing)
         vals, vecs = E
         @test Matrix(E) == Du
         @test vals isa AbstractVector{<:Furlong{1}}