diff --git a/Project.toml b/Project.toml index fbb1ce2f..7f3b7dfc 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.2" +version = "0.6.3" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/BlockSparseArrays.jl b/src/BlockSparseArrays.jl index b7658a5a..40b55492 100644 --- a/src/BlockSparseArrays.jl +++ b/src/BlockSparseArrays.jl @@ -47,5 +47,6 @@ include("BlockArraysSparseArraysBaseExt/BlockArraysSparseArraysBaseExt.jl") include("factorizations/svd.jl") include("factorizations/truncation.jl") include("factorizations/qr.jl") +include("factorizations/lq.jl") end diff --git a/src/factorizations/lq.jl b/src/factorizations/lq.jl new file mode 100644 index 00000000..4a07cfa6 --- /dev/null +++ b/src/factorizations/lq.jl @@ -0,0 +1,221 @@ +using MatrixAlgebraKit: MatrixAlgebraKit, lq_compact!, lq_full! + +# TODO: this is a hardcoded for now to get around this function not being defined in the +# type domain +function default_blocksparse_lq_algorithm(A::AbstractMatrix; kwargs...) + blocktype(A) <: StridedMatrix{<:LinearAlgebra.BLAS.BlasFloat} || + error("unsupported type: $(blocktype(A))") + alg = MatrixAlgebraKit.LAPACK_HouseholderLQ(; kwargs...) + return BlockPermutedDiagonalAlgorithm(alg) +end +function MatrixAlgebraKit.default_algorithm( + ::typeof(lq_compact!), A::AbstractBlockSparseMatrix; kwargs... +) + return default_blocksparse_lq_algorithm(A; kwargs...) +end +function MatrixAlgebraKit.default_algorithm( + ::typeof(lq_full!), A::AbstractBlockSparseMatrix; kwargs... +) + return default_blocksparse_lq_algorithm(A; kwargs...) +end + +function similar_output( + ::typeof(lq_compact!), A, L_axis, alg::MatrixAlgebraKit.AbstractAlgorithm +) + L = similar(A, axes(A, 1), L_axis) + Q = similar(A, L_axis, axes(A, 2)) + return L, Q +end + +function similar_output( + ::typeof(lq_full!), A, L_axis, alg::MatrixAlgebraKit.AbstractAlgorithm +) + L = similar(A, axes(A, 1), L_axis) + Q = similar(A, L_axis, axes(A, 2)) + return L, Q +end + +function MatrixAlgebraKit.initialize_output( + ::typeof(lq_compact!), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm +) + bm, bn = blocksize(A) + bmn = min(bm, bn) + + brows = eachblockaxis(axes(A, 1)) + bcols = eachblockaxis(axes(A, 2)) + l_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])) + l_axes[row] = bcols[col][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])) + l_axes[row] = bcols[col][Base.OneTo(len)] + end + + l_axis = mortar_axis(l_axes) + L, Q = similar_output(lq_compact!, A, l_axis, alg) + + # allocate output + for bI in eachblockstoredindex(A) + brow, bcol = Tuple(bI) + L[brow, brow], Q[brow, bcol] = MatrixAlgebraKit.initialize_output( + lq_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 L, Q +end + +function MatrixAlgebraKit.initialize_output( + ::typeof(lq_full!), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm +) + bm, bn = blocksize(A) + + bcols = eachblockaxis(axes(A, 2)) + l_axes = copy(bcols) + + # 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)) + l_axes[row] = bcols[col] + 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) + l_axes[row] = bcols[col] + end + for (i, k) in enumerate((length(emptycols) + 1):length(emptyrows)) + l_axes[bn + i] = bcols[emptycols[k]] + end + + l_axis = mortar_axis(l_axes) + L, Q = similar_output(lq_full!, A, l_axis, alg) + + # allocate output + for bI in eachblockstoredindex(A) + brow, bcol = Tuple(bI) + L[brow, brow], Q[brow, bcol] = MatrixAlgebraKit.initialize_output( + lq_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(emptyrows) + 1):length(emptycols)) + @view!(Q[Block(bm + i, emptycols[k])]) + end + + return L, Q +end + +function MatrixAlgebraKit.check_input( + ::typeof(lq_compact!), A::AbstractBlockSparseMatrix, LQ +) + L, Q = LQ + @assert isa(L, AbstractBlockSparseMatrix) && isa(Q, AbstractBlockSparseMatrix) + @assert eltype(A) == eltype(L) == eltype(Q) + @assert axes(A, 1) == axes(L, 1) && axes(A, 2) == axes(Q, 2) + @assert axes(L, 2) == axes(Q, 1) + + return nothing +end + +function MatrixAlgebraKit.check_input(::typeof(lq_full!), A::AbstractBlockSparseMatrix, LQ) + L, Q = LQ + @assert isa(L, AbstractBlockSparseMatrix) && isa(Q, AbstractBlockSparseMatrix) + @assert eltype(A) == eltype(L) == eltype(Q) + @assert axes(A, 1) == axes(L, 1) && axes(A, 2) == axes(Q, 2) + @assert axes(L, 2) == axes(Q, 1) + + return nothing +end + +function MatrixAlgebraKit.lq_compact!( + A::AbstractBlockSparseMatrix, LQ, alg::BlockPermutedDiagonalAlgorithm +) + MatrixAlgebraKit.check_input(lq_compact!, A, LQ) + L, Q = LQ + + # do decomposition on each block + for bI in eachblockstoredindex(A) + brow, bcol = Tuple(bI) + lq = (@view!(L[brow, brow]), @view!(Q[brow, bcol])) + lq′ = lq_compact!(@view!(A[bI]), lq, alg.alg) + @assert lq === lq′ "lq_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 LQ +end + +function MatrixAlgebraKit.lq_full!( + A::AbstractBlockSparseMatrix, LQ, alg::BlockPermutedDiagonalAlgorithm +) + MatrixAlgebraKit.check_input(lq_full!, A, LQ) + L, Q = LQ + + # do decomposition on each block + for bI in eachblockstoredindex(A) + brow, bcol = Tuple(bI) + lq = (@view!(L[brow, brow]), @view!(Q[brow, bcol])) + lq′ = lq_full!(@view!(A[bI]), lq, alg.alg) + @assert lq === lq′ "lq_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 + bm = blocksize(A, 1) + for (i, k) in enumerate((length(emptyrows) + 1):length(emptycols)) + copyto!(@view!(Q[Block(bm + i, emptycols[k])]), LinearAlgebra.I) + end + + return LQ +end diff --git a/test/test_factorizations.jl b/test/test_factorizations.jl index 02b42f86..e528ca69 100644 --- a/test/test_factorizations.jl +++ b/test/test_factorizations.jl @@ -1,7 +1,15 @@ using BlockArrays: Block, BlockedMatrix, BlockedVector, blocks, mortar using BlockSparseArrays: BlockSparseArray, BlockDiagonal, eachblockstoredindex using MatrixAlgebraKit: - qr_compact, qr_full, svd_compact, svd_full, svd_trunc, truncrank, trunctol + lq_compact, + lq_full, + qr_compact, + qr_full, + svd_compact, + svd_full, + svd_trunc, + truncrank, + trunctol using LinearAlgebra: LinearAlgebra using Random: Random using Test: @inferred, @testset, @test @@ -157,7 +165,7 @@ 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] + 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) @@ -181,3 +189,29 @@ end @test A ≈ Q * R end end + +@testset "lq_compact" 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) + L, Q = lq_compact(A) + @test Matrix(Q * Q') ≈ LinearAlgebra.I + @test A ≈ L * Q + end +end + +@testset "lq_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) + L, Q = lq_full(A) + L′, Q′ = lq_full(Matrix(A)) + @test size(L) == size(L′) + @test size(Q) == size(Q′) + @test Matrix(Q * Q') ≈ LinearAlgebra.I + @test Matrix(Q'Q) ≈ LinearAlgebra.I + @test A ≈ L * Q + end +end