Define block sparse eig_vals

Author: mtfishman

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -45,6 +45,9 @@ end
 function eachstoredblockdiagindex(a::AbstractArray)
   return eachblockstoredindex(a) ∩ blockdiagindices(a)
 end
+function eachunstoredblockdiagindex(a::AbstractArray)
+  return setdiff(blockdiagindices(a), eachblockstoredindex(a))
+end
 
 # Like `BlockArrays.eachblock` but only iterating
 # over stored blocks.

src/factorizations/eig.jl

@@ -1,24 +1,13 @@
+using BlockArrays: blocksizes
+using LinearAlgebra: LinearAlgebra
 using MatrixAlgebraKit:
-  MatrixAlgebraKit, default_eig_algorithm, default_eigh_algorithm, eig_full!, eigh_full!
-
-function initialize_blocksparse_eig_output(
-  f, A::AbstractMatrix, alg::BlockPermutedDiagonalAlgorithm
-)
-  Td, Tv = fieldtypes(Base.promote_op(f, blocktype(A), typeof(alg.alg)))
-  D = similar(A, BlockType(Td))
-  V = similar(A, BlockType(Tv))
-  return (D, V)
-end
-
-function blocksparse_eig_full!(
-  f, A::AbstractMatrix, (D, V), alg::BlockPermutedDiagonalAlgorithm
-)
-  for I in blockdiagindices(A)
-    d, v = f(@view!(A[I]), alg.alg)
-    D[I], V[I] = d, v
-  end
-  return (D, V)
-end
+  MatrixAlgebraKit,
+  default_eig_algorithm,
+  default_eigh_algorithm,
+  eig_full!,
+  eig_vals!,
+  eigh_full!,
+  eigh_vals!
 
 for f in [:default_eig_algorithm, :default_eigh_algorithm]
   @eval begin
@@ -34,12 +23,41 @@ for f in [:eig_full!, :eigh_full!]
     function MatrixAlgebraKit.initialize_output(
       ::typeof($f), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm
     )
-      return initialize_blocksparse_eig_output($f, A, alg)
+      Td, Tv = fieldtypes(Base.promote_op($f, blocktype(A), typeof(alg.alg)))
+      D = similar(A, BlockType(Td))
+      V = similar(A, BlockType(Tv))
+      return (D, V)
     end
     function MatrixAlgebraKit.$f(
       A::AbstractBlockSparseMatrix, (D, V), alg::BlockPermutedDiagonalAlgorithm
     )
-      return blocksparse_eig_full!($f, A, (D, V), alg)
+      for I in eachstoredblockdiagindex(A)
+        D[I], V[I] = $f(@view(A[I]), alg.alg)
+      end
+      for I in eachunstoredblockdiagindex(A)
+        # TODO: Support setting `LinearAlgebra.I` directly, and/or
+        # using `FillArrays.Eye`.
+        V[I] = LinearAlgebra.I(first(blocksizes(A)[Int.(Tuple(I))...]))
+      end
+      return (D, V)
+    end
+  end
+end
+
+for f in [:eig_vals!, :eigh_vals!]
+  @eval begin
+    function MatrixAlgebraKit.initialize_output(
+      ::typeof($f), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm
+    )
+      return similar(A, axes(A, 1))
+    end
+    function MatrixAlgebraKit.$f(
+      A::AbstractBlockSparseMatrix, D, alg::BlockPermutedDiagonalAlgorithm
+    )
+      for I in eachblockstoredindex(A)
+        D[I] = $f(@view!(A[I]), alg.alg)
+      end
+      return D
     end
   end
 end