Implement schur_full/schur_vals for Diagonal inputs

docs/src/changelog.md

@@ -37,6 +37,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec
 
 ### Fixed
 
+- `schur_full` and `schur_vals` now support `Diagonal` inputs through `DiagonalAlgorithm` ([#276](https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/issues/276)).
 - LQ decompositions no longer gauge fix `Q` when `positive = false` and `L` is not computed.
 
 ### Performance

src/implementations/schur.jl

@@ -24,6 +24,28 @@ function check_input(::typeof(schur_vals!), A::AbstractMatrix, vals, ::AbstractA
     return nothing
 end
 
+function check_input(::typeof(schur_full!), A::AbstractMatrix, TZv, ::DiagonalAlgorithm)
+    m = LinearAlgebra.checksquare(A)
+    isdiag(A) || throw(DimensionMismatch("diagonal input matrix expected"))
+    T, Z, vals = TZv
+    @assert T isa AbstractMatrix && Z isa AbstractMatrix && vals isa AbstractVector
+    @check_size(T, (m, m))
+    @check_scalar(T, A)
+    @check_size(Z, (m, m))
+    @check_scalar(Z, A)
+    @check_size(vals, (m,))
+    @check_scalar(vals, A)
+    return nothing
+end
+function check_input(::typeof(schur_vals!), A::AbstractMatrix, vals, ::DiagonalAlgorithm)
+    m = LinearAlgebra.checksquare(A)
+    isdiag(A) || throw(DimensionMismatch("diagonal input matrix expected"))
+    @assert vals isa AbstractVector
+    @check_size(vals, (m,))
+    @check_scalar(vals, A)
+    return nothing
+end
+
 # Outputs
 # -------
 function initialize_output(::typeof(schur_full!), A::AbstractMatrix, ::AbstractAlgorithm)
@@ -38,6 +60,17 @@ function initialize_output(::typeof(schur_vals!), A::AbstractMatrix, ::AbstractA
     return vals
 end
 
+# a diagonal matrix is already in Schur form, so the eigenvalues need not be complex
+function initialize_output(::typeof(schur_full!), A::AbstractMatrix, ::DiagonalAlgorithm)
+    n = size(A, 1) # square check will happen later
+    return (A, similar(A), similar(A, eltype(A), n))
+end
+function initialize_output(::typeof(schur_vals!), A::AbstractMatrix, ::DiagonalAlgorithm)
+    n = size(A, 1) # square check will happen later
+    return similar(A, eltype(A), n)
+end
+initialize_output(::typeof(schur_vals!), A::Diagonal, ::DiagonalAlgorithm) = diagview(A)
+
 # DefaultAlgorithm intercepts
 # ---------------------------
 for f! in (:schur_full!, :schur_vals!)
@@ -100,6 +133,26 @@ function schur_vals!(A::AbstractMatrix, vals, alg::QRIteration)
     return vals
 end
 
+# Diagonal logic
+# --------------
+# `A` is already in Schur form, so `T = A` and `Z = I`, without any reordering
+function schur_full!(A::AbstractMatrix, TZv, alg::DiagonalAlgorithm)
+    check_input(schur_full!, A, TZv, alg)
+    T, Z, vals = TZv
+    if !has_equal_storage(A, T)
+        zero!(T)
+        diagview(T) .= diagview(A)
+    end
+    one!(Z)
+    vals .= diagview(A)
+    return TZv
+end
+function schur_vals!(A::AbstractMatrix, vals, alg::DiagonalAlgorithm)
+    check_input(schur_vals!, A, vals, alg)
+    has_equal_storage(A, vals) || copy!(vals, diagview(A))
+    return vals
+end
+
 # Deprecations
 # ------------
 for (lapack_algtype, expert_val) in ((:LAPACK_Simple, false), (:LAPACK_Expert, true))

test/decompositions/schur.jl

@@ -3,6 +3,7 @@ using Test
 using TestExtras
 using StableRNGs
 using LinearAlgebra: I, Diagonal
+using CUDA, AMDGPU
 
 if @isdefined(fast_tests) && fast_tests
     BLASFloats = (Float64, ComplexF64)
@@ -21,22 +22,25 @@ m = 54
 for T in (BLASFloats..., GenericFloats...)
     TestSuite.seed_rng!(123)
     if T ∈ BLASFloats
-        #=if CUDA.functional()
-            TestSuite.test_schur(CuMatrix{T}, (m, m); test_blocksize = false)
-            TestSuite.test_schur(Diagonal{T, CuVector{T}}, m; test_blocksize = false)
+        if CUDA.functional()
+            # dense GPU schur is not yet supported: there is no `gees!` for CUSOLVER
+            TestSuite.test_schur(Diagonal{T, CuVector{T}}, m)
+            TestSuite.test_schur_algs(Diagonal{T, CuVector{T}}, m, (DiagonalAlgorithm(),))
         end
+        #= not yet supported
         if AMDGPU.functional()
             TestSuite.test_schur(ROCMatrix{T}, (m, m); test_blocksize = false)
             TestSuite.test_schur(Diagonal{T, ROCVector{T}}, m; test_blocksize = false)
-        end=# # not yet supported
+        end=#
     end
     if !is_buildkite
         TestSuite.test_schur(T, (m, m))
         if T ∈ BLASFloats
             LAPACK_SCHUR_ALGS = (QRIteration(), QRIteration(expert = true))
             TestSuite.test_schur_algs(T, (m, m), LAPACK_SCHUR_ALGS)
         end
-        #AT = Diagonal{T, Vector{T}}
-        #TestSuite.test_schur(AT, m) # not supported yet
+        AT = Diagonal{T, Vector{T}}
+        TestSuite.test_schur(AT, m)
+        TestSuite.test_schur_algs(AT, m, (DiagonalAlgorithm(),))
     end
 end

test/testsuite/decompositions/schur.jl

@@ -33,6 +33,8 @@ function test_schur_full(
         @test eltype(vals) == Tc
         @test isisometric(Z)
         @test A * Z ≈ Z * TA
+        # a diagonal matrix is already in Schur form and is not reordered
+        A isa Diagonal && @test TA ≈ A
 
         TA2, Z2, vals2 = @testinferred schur_full!(Ac, (TA, Z, vals))
         @test TA2 === TA
@@ -78,14 +80,17 @@ function test_schur_vals(
         Ac = deepcopy(A)
         Tc = isa(A, Diagonal) ? eltype(T) : complex(eltype(T))
 
+        # a diagonal matrix is not reordered, unlike `eig_vals`
+        vals₀ = A isa Diagonal ? diagview(A) : eig_vals(A)
+
         valsc = @testinferred schur_vals(A)
         @test eltype(valsc) == Tc
-        @test valsc ≈ eig_vals(A)
+        @test valsc ≈ vals₀
 
         valsc = similar(A, Tc, size(A, 1))
         valsc = @testinferred schur_vals!(Ac, valsc)
         @test eltype(valsc) == Tc
-        @test valsc ≈ eig_vals(A)
+        @test valsc ≈ vals₀
     end
 end
 
@@ -100,13 +105,15 @@ function test_schur_vals_algs(
         Ac = deepcopy(A)
         Tc = isa(A, Diagonal) ? eltype(T) : complex(eltype(T))
 
+        vals₀ = A isa Diagonal ? diagview(A) : eig_vals(A)
+
         valsc = @testinferred schur_vals(A; alg)
         @test eltype(valsc) == Tc
-        @test valsc ≈ eig_vals(A)
+        @test valsc ≈ vals₀
 
         valsc = similar(A, Tc, size(A, 1))
         valsc = @testinferred schur_vals!(Ac, valsc; alg)
         @test eltype(valsc) == Tc
-        @test valsc ≈ eig_vals(A)
+        @test valsc ≈ vals₀
     end
 end