Blockdiagonal factorizations refactor (#166)

Author: lkdvos

Project.toml

@@ -1,7 +1,7 @@
 name = "BlockSparseArrays"
 uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.9.4"
+version = "0.10.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

docs/Project.toml

@@ -6,6 +6,6 @@ Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 
 [compat]
 BlockArrays = "1"
-BlockSparseArrays = "0.9"
+BlockSparseArrays = "0.10"
 Documenter = "1"
 Literate = "2"

examples/Project.toml

@@ -5,5 +5,5 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 BlockArrays = "1"
-BlockSparseArrays = "0.9"
+BlockSparseArrays = "0.10"
 Test = "1"

src/BlockSparseArrays.jl

@@ -45,7 +45,7 @@ include("blocksparsearray/blockdiagonalarray.jl")
 include("BlockArraysSparseArraysBaseExt/BlockArraysSparseArraysBaseExt.jl")
 
 # factorizations
-include("factorizations/tensorproducts.jl")
+include("factorizations/utility.jl")
 include("factorizations/svd.jl")
 include("factorizations/truncation.jl")
 include("factorizations/qr.jl")

src/factorizations/eig.jl

@@ -1,85 +1,82 @@
 using BlockArrays: blocksizes
 using DiagonalArrays: diagonal
 using LinearAlgebra: LinearAlgebra, Diagonal
-using MatrixAlgebraKit:
-  MatrixAlgebraKit,
-  TruncationStrategy,
-  check_input,
-  default_eig_algorithm,
-  default_eigh_algorithm,
-  diagview,
-  eig_full!,
-  eig_trunc!,
-  eig_vals!,
-  eigh_full!,
-  eigh_trunc!,
-  eigh_vals!,
-  findtruncated
+using MatrixAlgebraKit: MatrixAlgebraKit, diagview
+using MatrixAlgebraKit: default_eig_algorithm, eig_full!, eig_vals!
+using MatrixAlgebraKit: default_eigh_algorithm, eigh_full!, eigh_vals!
 
 for f in [:default_eig_algorithm, :default_eigh_algorithm]
   @eval begin
     function MatrixAlgebraKit.$f(::Type{<:AbstractBlockSparseMatrix}; kwargs...)
-      return BlockPermutedDiagonalAlgorithm() do block
+      return BlockDiagonalAlgorithm() do block
         return $f(block; kwargs...)
       end
     end
   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.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
+      ::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, (D, V), alg::BlockDiagonalAlgorithm
     )
-      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, (D, V), alg)
+
+      # do decomposition on each block
+      for bI in blockdiagindices(A)
+        if isstored(A, bI)
+          block = @view!(A[bI])
+          block_alg = block_algorithm(alg, block)
+          bD, bV = $f(block, block_alg)
+          D[bI] = bD
+          V[bI] = bV
+        else
+          # TODO: this should be `V[bI] = LinearAlgebra.I`
+          copyto!(@view!(V[bI]), LinearAlgebra.I)
+        end
       end
       return (D, V)
     end
@@ -100,17 +97,29 @@ end
 for f in [:eig_vals!, :eigh_vals!]
   @eval begin
     function MatrixAlgebraKit.initialize_output(
-      ::typeof($f), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm
+      ::typeof($f), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm
     )
       T = output_type($f, blocktype(A))
       return similar(A, BlockType(T), axes(A, 1))
     end
+    function MatrixAlgebraKit.check_input(
+      ::typeof($f), A::AbstractBlockSparseMatrix, D, ::BlockDiagonalAlgorithm
+    )
+      @assert isa(D, AbstractBlockSparseVector)
+      @assert eltype(D) === $(f == :eig_vals! ? complex : real)(eltype(A))
+      @assert axes(A, 1) == axes(A, 2)
+      @assert (axes(A, 1),) == axes(D)
+      @assert isblockdiagonal(A)
+      return nothing
+    end
+
     function MatrixAlgebraKit.$f(
-      A::AbstractBlockSparseMatrix, D, alg::BlockPermutedDiagonalAlgorithm
+      A::AbstractBlockSparseMatrix, D, alg::BlockDiagonalAlgorithm
     )
+      MatrixAlgebraKit.check_input($f, A, D, alg)
       for I in eachblockstoredindex(A)
         block = @view!(A[I])
-        D[I] = $f(block, block_algorithm(alg, block))
+        D[Tuple(I)[1]] = $f(block, block_algorithm(alg, block))
       end
       return D
     end