diff --git a/Project.toml b/Project.toml index e1edff8e..fbb1ce2f 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "BlockSparseArrays" uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4" authors = ["ITensor developers and contributors"] -version = "0.6.1" +version = "0.6.2" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/BlockSparseArrays.jl b/src/BlockSparseArrays.jl index 893bcf0c..b7658a5a 100644 --- a/src/BlockSparseArrays.jl +++ b/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 diff --git a/src/factorizations/qr.jl b/src/factorizations/qr.jl new file mode 100644 index 00000000..c28da925 --- /dev/null +++ b/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 diff --git a/test/test_factorizations.jl b/test/test_factorizations.jl index 5152237b..02b42f86 100644 --- a/test/test_factorizations.jl +++ b/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