Refactor eig

Author: lkdvos

src/factorizations/eig.jl

@@ -4,7 +4,6 @@ using LinearAlgebra: LinearAlgebra, Diagonal
 using MatrixAlgebraKit:
   MatrixAlgebraKit,
   TruncationStrategy,
-  check_input,
   default_eig_algorithm,
   default_eigh_algorithm,
   diagview,
@@ -26,60 +25,100 @@ for f in [:default_eig_algorithm, :default_eigh_algorithm]
   end
 end
 
+function output_type(::typeof(eig_full!), A::Type{<:AbstractMatrix{T}}) where {T}
+  DV = Base.promote_op(eig_full!, A)
+  return if isconcretetype(DV)
+    DV
+  else
+    Tuple{AbstractMatrix{complex(T)},AbstractMatrix{complex(T)}}
+  end
+end
+function output_type(::typeof(eigh_full!), A::Type{<:AbstractMatrix{T}}) where {T}
+  DV = Base.promote_op(eigh_full!, A)
+  return isconcretetype(DV) ? DV : Tuple{AbstractMatrix{real(T)},AbstractMatrix{T}}
+end
+
+function MatrixAlgebraKit.initialize_output(
+  ::Union{typeof(eig_full!),typeof(eigh_full!)},
+  ::AbstractBlockSparseMatrix,
+  ::BlockPermutedDiagonalAlgorithm,
+)
+  return nothing
+end
+
+function MatrixAlgebraKit.check_input(
+  ::Union{typeof(eig_full!),typeof(eigh_full!)},
+  A::AbstractBlockSparseMatrix,
+  DV,
+  ::BlockPermutedDiagonalAlgorithm,
+)
+  @assert isblockpermuteddiagonal(A)
+  return nothing
+end
 function MatrixAlgebraKit.check_input(
-  ::typeof(eig_full!), A::AbstractBlockSparseMatrix, (D, V)
+  ::typeof(eig_full!), A::AbstractBlockSparseMatrix, (D, V), ::BlockDiagonalAlgorithm
 )
   @assert isa(D, AbstractBlockSparseMatrix) && isa(V, AbstractBlockSparseMatrix)
   @assert eltype(V) === eltype(D) === complex(eltype(A))
   @assert axes(A, 1) == axes(A, 2)
   @assert axes(A) == axes(D) == axes(V)
+  @assert isblockdiagonal(A)
   return nothing
 end
 function MatrixAlgebraKit.check_input(
-  ::typeof(eigh_full!), A::AbstractBlockSparseMatrix, (D, V)
+  ::typeof(eigh_full!), A::AbstractBlockSparseMatrix, (D, V), ::BlockDiagonalAlgorithm
 )
   @assert isa(D, AbstractBlockSparseMatrix) && isa(V, AbstractBlockSparseMatrix)
   @assert eltype(V) === eltype(A)
   @assert eltype(D) === real(eltype(A))
   @assert axes(A, 1) == axes(A, 2)
   @assert axes(A) == axes(D) == axes(V)
+  @assert isblockdiagonal(A)
   return nothing
 end
 
-function output_type(f::typeof(eig_full!), A::Type{<:AbstractMatrix{T}}) where {T}
-  DV = Base.promote_op(f, A)
-  !isconcretetype(DV) && return Tuple{AbstractMatrix{complex(T)},AbstractMatrix{complex(T)}}
-  return DV
-end
-function output_type(f::typeof(eigh_full!), A::Type{<:AbstractMatrix{T}}) where {T}
-  DV = Base.promote_op(f, A)
-  !isconcretetype(DV) && return Tuple{AbstractMatrix{real(T)},AbstractMatrix{T}}
-  return DV
-end
-
 for f in [:eig_full!, :eigh_full!]
   @eval begin
     function MatrixAlgebraKit.initialize_output(
       ::typeof($f), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm
+    )
+      return nothing
+    end
+    function MatrixAlgebraKit.initialize_output(
+      ::typeof($f), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm
     )
       Td, Tv = fieldtypes(output_type($f, blocktype(A)))
       D = similar(A, BlockType(Td))
       V = similar(A, BlockType(Tv))
       return (D, V)
     end
     function MatrixAlgebraKit.$f(
-      A::AbstractBlockSparseMatrix, (D, V), alg::BlockPermutedDiagonalAlgorithm
+      A::AbstractBlockSparseMatrix, DV, alg::BlockPermutedDiagonalAlgorithm
     )
-      check_input($f, A, (D, V))
-      for I in eachstoredblockdiagindex(A)
-        block = @view!(A[I])
-        block_alg = block_algorithm(alg, block)
-        D[I], V[I] = $f(block, block_alg)
-      end
-      for I in eachunstoredblockdiagindex(A)
-        # TODO: Support setting `LinearAlgebra.I` directly, and/or
-        # using `FillArrays.Eye`.
-        V[I] = LinearAlgebra.I(size(@view(V[I]), 1))
+      MatrixAlgebraKit.check_input($f, A, DV, alg)
+      Ad, transform_rows, transform_cols = blockdiagonalize(A)
+      Dd, Vd = $f(Ad, BlockDiagonalAlgorithm(alg))
+      D = transform_rows(Dd)
+      V = transform_cols(Vd)
+      return D, V
+    end
+    function MatrixAlgebraKit.$f(
+      A::AbstractBlockSparseMatrix, (D, V), alg::BlockDiagonalAlgorithm
+    )
+      MatrixAlgebraKit.check_input($f, A, (D, V), alg)
+
+      # do decomposition on each block
+      for I in 1:min(blocksize(A)...)
+        bI = Block(I, I)
+        if isstored(blocks(A), CartesianIndex(I, I)) # TODO: isblockstored
+          block = @view!(A[bI])
+          block_alg = block_algorithm(alg, block)
+          bD, bV = $f(block, block_alg)
+          D[bI] = bD
+          V[bI] = bV
+        else
+          copyto!(@view!(V[bI]), LinearAlgebra.I)
+        end
       end
       return (D, V)
     end