Add support for block sparse QR decomposition

Project.toml

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

src/BlockSparseArrays.jl

@@ -46,5 +46,6 @@ include("BlockArraysSparseArraysBaseExt/BlockArraysSparseArraysBaseExt.jl")
 # factorizations
 include("factorizations/svd.jl")
 include("factorizations/truncation.jl")
+include("factorizations/qr.jl")
 
 end

src/factorizations/qr.jl

@@ -0,0 +1,221 @@
+using MatrixAlgebraKit: MatrixAlgebraKit, qr_compact!, qr_full!
+
+# TODO: this is a hardcoded for now to get around this function not being defined in the
+# type domain
+function default_blocksparse_qr_algorithm(A::AbstractMatrix; kwargs...)
+  blocktype(A) <: StridedMatrix{<:LinearAlgebra.BLAS.BlasFloat} ||
+    error("unsupported type: $(blocktype(A))")
+  alg = MatrixAlgebraKit.LAPACK_HouseholderQR(; kwargs...)
+  return BlockPermutedDiagonalAlgorithm(alg)
+end
+function MatrixAlgebraKit.default_algorithm(
+  ::typeof(qr_compact!), A::AbstractBlockSparseMatrix; kwargs...
+)
+  return default_blocksparse_qr_algorithm(A; kwargs...)
+end
+function MatrixAlgebraKit.default_algorithm(
+  ::typeof(qr_full!), A::AbstractBlockSparseMatrix; kwargs...
+)
+  return default_blocksparse_qr_algorithm(A; kwargs...)
+end
+
+function similar_output(
+  ::typeof(qr_compact!), A, R_axis, alg::MatrixAlgebraKit.AbstractAlgorithm
+)
+  Q = similar(A, axes(A, 1), R_axis)
+  R = similar(A, R_axis, axes(A, 2))
+  return Q, R
+end
+
+function similar_output(
+  ::typeof(qr_full!), A, R_axis, alg::MatrixAlgebraKit.AbstractAlgorithm
+)
+  Q = similar(A, axes(A, 1), R_axis)
+  R = similar(A, R_axis, axes(A, 2))
+  return Q, R
+end
+
+function MatrixAlgebraKit.initialize_output(
+  ::typeof(qr_compact!), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm
+)
+  bm, bn = blocksize(A)
+  bmn = min(bm, bn)
+
+  brows = eachblockaxis(axes(A, 1))
+  bcols = eachblockaxis(axes(A, 2))
+  r_axes = similar(brows, bmn)
+
+  # fill in values for blocks that are present
+  bIs = collect(eachblockstoredindex(A))
+  browIs = Int.(first.(Tuple.(bIs)))
+  bcolIs = Int.(last.(Tuple.(bIs)))
+  for bI in eachblockstoredindex(A)
+    row, col = Int.(Tuple(bI))
+    len = minimum(length, (brows[row], bcols[col]))
+    r_axes[col] = brows[row][Base.OneTo(len)]
+  end
+
+  # fill in values for blocks that aren't present, pairing them in order of occurence
+  # this is a convention, which at least gives the expected results for blockdiagonal
+  emptyrows = setdiff(1:bm, browIs)
+  emptycols = setdiff(1:bn, bcolIs)
+  for (row, col) in zip(emptyrows, emptycols)
+    len = minimum(length, (brows[row], bcols[col]))
+    r_axes[col] = brows[row][Base.OneTo(len)]
+  end
+
+  r_axis = mortar_axis(r_axes)
+  Q, R = similar_output(qr_compact!, A, r_axis, alg)
+
+  # allocate output
+  for bI in eachblockstoredindex(A)
+    brow, bcol = Tuple(bI)
+    Q[brow, bcol], R[bcol, bcol] = MatrixAlgebraKit.initialize_output(
+      qr_compact!, @view!(A[bI]), alg.alg
+    )
+  end
+
+  # allocate output for blocks that aren't present -- do we also fill identities here?
+  for (row, col) in zip(emptyrows, emptycols)
+    @view!(Q[Block(row, col)])
+  end
+
+  return Q, R
+end
+
+function MatrixAlgebraKit.initialize_output(
+  ::typeof(qr_full!), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm
+)
+  bm, bn = blocksize(A)
+
+  brows = eachblockaxis(axes(A, 1))
+  r_axes = copy(brows)
+
+  # fill in values for blocks that are present
+  bIs = collect(eachblockstoredindex(A))
+  browIs = Int.(first.(Tuple.(bIs)))
+  bcolIs = Int.(last.(Tuple.(bIs)))
+  for bI in eachblockstoredindex(A)
+    row, col = Int.(Tuple(bI))
+    r_axes[col] = brows[row]
+  end
+
+  # fill in values for blocks that aren't present, pairing them in order of occurence
+  # this is a convention, which at least gives the expected results for blockdiagonal
+  emptyrows = setdiff(1:bm, browIs)
+  emptycols = setdiff(1:bn, bcolIs)
+  for (row, col) in zip(emptyrows, emptycols)
+    r_axes[col] = brows[row]
+  end
+  for (i, k) in enumerate((length(emptycols) + 1):length(emptyrows))
+    r_axes[bn + i] = brows[emptyrows[k]]
+  end
+
+  r_axis = mortar_axis(r_axes)
+  Q, R = similar_output(qr_full!, A, r_axis, alg)
+
+  # allocate output
+  for bI in eachblockstoredindex(A)
+    brow, bcol = Tuple(bI)
+    Q[brow, bcol], R[bcol, bcol] = MatrixAlgebraKit.initialize_output(
+      qr_full!, @view!(A[bI]), alg.alg
+    )
+  end
+
+  # allocate output for blocks that aren't present -- do we also fill identities here?
+  for (row, col) in zip(emptyrows, emptycols)
+    @view!(Q[Block(row, col)])
+  end
+  # also handle extra rows/cols
+  for (i, k) in enumerate((length(emptycols) + 1):length(emptyrows))
+    @view!(Q[Block(emptyrows[k], bn + i)])
+  end
+
+  return Q, R
+end
+
+function MatrixAlgebraKit.check_input(
+  ::typeof(qr_compact!), A::AbstractBlockSparseMatrix, QR
+)
+  Q, R = QR
+  @assert isa(Q, AbstractBlockSparseMatrix) && isa(R, AbstractBlockSparseMatrix)
+  @assert eltype(A) == eltype(Q) == eltype(R)
+  @assert axes(A, 1) == axes(Q, 1) && axes(A, 2) == axes(R, 2)
+  @assert axes(Q, 2) == axes(R, 1)
+
+  return nothing
+end
+
+function MatrixAlgebraKit.check_input(::typeof(qr_full!), A::AbstractBlockSparseMatrix, QR)
+  Q, R = QR
+  @assert isa(Q, AbstractBlockSparseMatrix) && isa(R, AbstractBlockSparseMatrix)
+  @assert eltype(A) == eltype(Q) == eltype(R)
+  @assert axes(A, 1) == axes(Q, 1) && axes(A, 2) == axes(R, 2)
+  @assert axes(Q, 2) == axes(R, 1)
+
+  return nothing
+end
+
+function MatrixAlgebraKit.qr_compact!(
+  A::AbstractBlockSparseMatrix, QR, alg::BlockPermutedDiagonalAlgorithm
+)
+  MatrixAlgebraKit.check_input(qr_compact!, A, QR)
+  Q, R = QR
+
+  # do decomposition on each block
+  for bI in eachblockstoredindex(A)
+    brow, bcol = Tuple(bI)
+    qr = (@view!(Q[brow, bcol]), @view!(R[bcol, bcol]))
+    qr′ = qr_compact!(@view!(A[bI]), qr, alg.alg)
+    @assert qr === qr′ "qr_compact! might not be in-place"
+  end
+
+  # fill in identities for blocks that aren't present
+  bIs = collect(eachblockstoredindex(A))
+  browIs = Int.(first.(Tuple.(bIs)))
+  bcolIs = Int.(last.(Tuple.(bIs)))
+  emptyrows = setdiff(1:blocksize(A, 1), browIs)
+  emptycols = setdiff(1:blocksize(A, 2), bcolIs)
+  # needs copyto! instead because size(::LinearAlgebra.I) doesn't work
+  # Q[Block(row, col)] = LinearAlgebra.I
+  for (row, col) in zip(emptyrows, emptycols)
+    copyto!(@view!(Q[Block(row, col)]), LinearAlgebra.I)
+  end
+
+  return QR
+end
+
+function MatrixAlgebraKit.qr_full!(
+  A::AbstractBlockSparseMatrix, QR, alg::BlockPermutedDiagonalAlgorithm
+)
+  MatrixAlgebraKit.check_input(qr_full!, A, QR)
+  Q, R = QR
+
+  # do decomposition on each block
+  for bI in eachblockstoredindex(A)
+    brow, bcol = Tuple(bI)
+    qr = (@view!(Q[brow, bcol]), @view!(R[bcol, bcol]))
+    qr′ = qr_full!(@view!(A[bI]), qr, alg.alg)
+    @assert qr === qr′ "qr_full! might not be in-place"
+  end
+
+  # fill in identities for blocks that aren't present
+  bIs = collect(eachblockstoredindex(A))
+  browIs = Int.(first.(Tuple.(bIs)))
+  bcolIs = Int.(last.(Tuple.(bIs)))
+  emptyrows = setdiff(1:blocksize(A, 1), browIs)
+  emptycols = setdiff(1:blocksize(A, 2), bcolIs)
+  # needs copyto! instead because size(::LinearAlgebra.I) doesn't work
+  # Q[Block(row, col)] = LinearAlgebra.I
+  for (row, col) in zip(emptyrows, emptycols)
+    copyto!(@view!(Q[Block(row, col)]), LinearAlgebra.I)
+  end
+
+  # also handle extra rows/cols
+  bn = blocksize(A, 2)
+  for (i, k) in enumerate((length(emptycols) + 1):length(emptyrows))
+    copyto!(@view!(Q[Block(emptyrows[k], bn + i)]), LinearAlgebra.I)
+  end
+
+  return QR
+end

test/test_factorizations.jl

@@ -1,6 +1,7 @@
 using BlockArrays: Block, BlockedMatrix, BlockedVector, blocks, mortar
 using BlockSparseArrays: BlockSparseArray, BlockDiagonal, eachblockstoredindex
-using MatrixAlgebraKit: svd_compact, svd_full, svd_trunc, truncrank, trunctol
+using MatrixAlgebraKit:
+  qr_compact, qr_full, svd_compact, svd_full, svd_trunc, truncrank, trunctol
 using LinearAlgebra: LinearAlgebra
 using Random: Random
 using Test: @inferred, @testset, @test
@@ -154,3 +155,29 @@ end
     @test (V1ᴴ * V1ᴴ' ≈ LinearAlgebra.I)
   end
 end
+
+@testset "qr_compact" for T in (Float32, Float64, ComplexF32, ComplexF64)
+  for i in [1, 2], j in [1, 2], k in [1, 2], l in [1, 2]
+    A = BlockSparseArray{T}(undef, ([i, j], [k, l]))
+    A[Block(1, 1)] = randn(T, i, k)
+    A[Block(2, 2)] = randn(T, j, l)
+    Q, R = qr_compact(A)
+    @test Matrix(Q'Q) ≈ LinearAlgebra.I
+    @test A ≈ Q * R
+  end
+end
+
+@testset "qr_full" for T in (Float32, Float64, ComplexF32, ComplexF64)
+  for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3]
+    A = BlockSparseArray{T}(undef, ([i, j], [k, l]))
+    A[Block(1, 1)] = randn(T, i, k)
+    A[Block(2, 2)] = randn(T, j, l)
+    Q, R = qr_full(A)
+    Q′, R′ = qr_full(Matrix(A))
+    @test size(Q) == size(Q′)
+    @test size(R) == size(R′)
+    @test Matrix(Q'Q) ≈ LinearAlgebra.I
+    @test Matrix(Q * Q') ≈ LinearAlgebra.I
+    @test A ≈ Q * R
+  end
+end