diff --git a/Project.toml b/Project.toml index 0ca7f6c..1905cdc 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.9.1" +version = "0.10.0" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/docs/Project.toml b/docs/Project.toml index edbaa0b..6f5417f 100644 --- a/docs/Project.toml +++ b/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" diff --git a/examples/Project.toml b/examples/Project.toml index 7e02e0d..14c5df3 100644 --- a/examples/Project.toml +++ b/examples/Project.toml @@ -5,5 +5,5 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] BlockArrays = "1" -BlockSparseArrays = "0.9" +BlockSparseArrays = "0.10" Test = "1" diff --git a/src/BlockSparseArrays.jl b/src/BlockSparseArrays.jl index 4229343..8949bad 100644 --- a/src/BlockSparseArrays.jl +++ b/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") diff --git a/src/factorizations/eig.jl b/src/factorizations/eig.jl index 502ac10..ce1288b 100644 --- a/src/factorizations/eig.jl +++ b/src/factorizations/eig.jl @@ -1,20 +1,9 @@ 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 @@ -26,41 +15,67 @@ for f in [:default_eig_algorithm, :default_eigh_algorithm] 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.initialize_output( + ::Union{typeof(eig_full!),typeof(eigh_full!)}, + ::AbstractBlockSparseMatrix, + ::BlockPermutedDiagonalAlgorithm, +) + return nothing +end + +function MatrixAlgebraKit.check_input( + ::Union{typeof(eig_full!),typeof(eigh_full!)}, + A::AbstractBlockSparseMatrix, + DV, + ::BlockPermutedDiagonalAlgorithm, +) + @assert isblockpermuteddiagonal(A) + return nothing +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 + ) + return nothing + end + function MatrixAlgebraKit.initialize_output( + ::typeof($f), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm ) Td, Tv = fieldtypes(output_type($f, blocktype(A))) D = similar(A, BlockType(Td)) @@ -68,18 +83,32 @@ for f in [:eig_full!, :eigh_full!] return (D, V) end function MatrixAlgebraKit.$f( - A::AbstractBlockSparseMatrix, (D, V), alg::BlockPermutedDiagonalAlgorithm + A::AbstractBlockSparseMatrix, DV, alg::BlockPermutedDiagonalAlgorithm ) - 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, DV, alg) + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + Dd, Vd = $f(Ad, BlockDiagonalAlgorithm(alg)) + D = transform_rows(Dd, invrowperm) + V = transform_cols(Vd, invcolperm) + return D, V + end + function MatrixAlgebraKit.$f( + A::AbstractBlockSparseMatrix, (D, V), alg::BlockDiagonalAlgorithm + ) + MatrixAlgebraKit.check_input($f, A, (D, V), alg) + + # do decomposition on each block + for I in 1:min(blocksize(A)...) + bI = Block(I, I) + if isstored(blocks(A), CartesianIndex(I, I)) # TODO: isblockstored + block = @view!(A[bI]) + block_alg = block_algorithm(alg, block) + bD, bV = $f(block, block_alg) + D[bI] = bD + V[bI] = bV + else + copyto!(@view!(V[bI]), LinearAlgebra.I) + end end return (D, V) end @@ -101,16 +130,47 @@ for f in [:eig_vals!, :eigh_vals!] @eval begin function MatrixAlgebraKit.initialize_output( ::typeof($f), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm + ) + return nothing + end + function MatrixAlgebraKit.initialize_output( + ::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, ::BlockPermutedDiagonalAlgorithm + ) + @assert isblockpermuteddiagonal(A) + return nothing + 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 ) + MatrixAlgebraKit.check_input($f, A, D, alg) + Ad, (invrowperm, _) = blockdiagonalize(A) + Dd = $f(Ad, BlockDiagonalAlgorithm(alg)) + return transform_rows(Dd, invrowperm) + end + function MatrixAlgebraKit.$f( + 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 diff --git a/src/factorizations/lq.jl b/src/factorizations/lq.jl index 5dc1301..1f5bed4 100644 --- a/src/factorizations/lq.jl +++ b/src/factorizations/lq.jl @@ -8,211 +8,140 @@ function MatrixAlgebraKit.default_lq_algorithm( end 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 +function output_type( + f::Union{typeof(lq_compact!),typeof(lq_full!)}, A::Type{<:AbstractMatrix{T}} +) where {T} + LQ = Base.promote_op(f, A) + return isconcretetype(LQ) ? LQ : Tuple{AbstractMatrix{T},AbstractMatrix{T}} end -function similar_output( - ::typeof(lq_full!), A, L_axis, alg::MatrixAlgebraKit.AbstractAlgorithm +function MatrixAlgebraKit.initialize_output( + ::typeof(lq_compact!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm ) - L = similar(A, axes(A, 1), L_axis) - Q = similar(A, L_axis, axes(A, 2)) - return L, Q + return nothing end - function MatrixAlgebraKit.initialize_output( - ::typeof(lq_compact!), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm + ::typeof(lq_compact!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm ) - 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 - + # using the property that zip stops as soon as one of the iterators is exhausted + l_axes = map(splat(infimum), zip(brows, bcols)) 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) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - L[brow, brow], Q[brow, bcol] = MatrixAlgebraKit.initialize_output( - lq_compact!, block, block_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 + BL, BQ = fieldtypes(output_type(lq_compact!, blocktype(A))) + L = similar(A, BlockType(BL), (axes(A, 1), l_axis)) + Q = similar(A, BlockType(BQ), (l_axis, axes(A, 2))) return L, Q end function MatrixAlgebraKit.initialize_output( - ::typeof(lq_full!), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm + ::typeof(lq_full!), ::AbstractBlockSparseMatrix, ::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) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - L[brow, brow], Q[brow, bcol] = MatrixAlgebraKit.initialize_output( - lq_full!, block, block_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 nothing +end +function MatrixAlgebraKit.initialize_output( + ::typeof(lq_full!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm +) + BL, BQ = fieldtypes(output_type(lq_full!, blocktype(A))) + L = similar(A, BlockType(BL), (axes(A, 1), axes(A, 2))) + Q = similar(A, BlockType(BQ), (axes(A, 2), axes(A, 2))) return L, Q end function MatrixAlgebraKit.check_input( - ::typeof(lq_compact!), A::AbstractBlockSparseMatrix, LQ + ::typeof(lq_compact!), A::AbstractBlockSparseMatrix, LQ, ::BlockPermutedDiagonalAlgorithm +) + @assert isblockpermuteddiagonal(A) + return nothing +end +function MatrixAlgebraKit.check_input( + ::typeof(lq_compact!), A::AbstractBlockSparseMatrix, (L, Q), ::BlockDiagonalAlgorithm ) - 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) - + @assert isblockdiagonal(A) return nothing end - -function MatrixAlgebraKit.check_input(::typeof(lq_full!), A::AbstractBlockSparseMatrix, LQ) - L, Q = LQ +function MatrixAlgebraKit.check_input( + ::typeof(lq_full!), A::AbstractBlockSparseMatrix, LQ, ::BlockPermutedDiagonalAlgorithm +) + @assert isblockpermuteddiagonal(A) + return nothing +end +function MatrixAlgebraKit.check_input( + ::typeof(lq_full!), A::AbstractBlockSparseMatrix, (L, Q), ::BlockDiagonalAlgorithm +) @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) - + @assert isblockdiagonal(A) return nothing end function MatrixAlgebraKit.lq_compact!( A::AbstractBlockSparseMatrix, LQ, alg::BlockPermutedDiagonalAlgorithm ) - MatrixAlgebraKit.check_input(lq_compact!, A, LQ) - L, Q = LQ + MatrixAlgebraKit.check_input(lq_compact!, A, LQ, alg) + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + Ld, Qd = lq_compact!(Ad, BlockDiagonalAlgorithm(alg)) + L = transform_rows(Ld, invrowperm) + Q = transform_cols(Qd, invcolperm) + return L, Q +end +function MatrixAlgebraKit.lq_compact!( + A::AbstractBlockSparseMatrix, (L, Q), alg::BlockDiagonalAlgorithm +) + MatrixAlgebraKit.check_input(lq_compact!, A, (L, Q), alg) # do decomposition on each block - for bI in eachblockstoredindex(A) - brow, bcol = Tuple(bI) - lq = (@view!(L[brow, brow]), @view!(Q[brow, bcol])) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - lq′ = lq_compact!(block, lq, block_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) + for I in 1:min(blocksize(A)...) + bI = Block(I, I) + if isstored(blocks(A), CartesianIndex(I, I)) # TODO: isblockstored + block = @view!(A[bI]) + block_alg = block_algorithm(alg, block) + bL, bQ = lq_compact!(block, block_alg) + L[bI] = bL + Q[bI] = bQ + else + copyto!(@view!(Q[bI]), LinearAlgebra.I) + end end - return LQ + return L, Q 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])) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - lq′ = lq_full!(block, lq, block_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 + MatrixAlgebraKit.check_input(lq_full!, A, LQ, alg) + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + Ld, Qd = lq_full!(Ad, BlockDiagonalAlgorithm(alg)) + L = transform_rows(Ld, invrowperm) + Q = transform_cols(Qd, invcolperm) + return L, Q +end +function MatrixAlgebraKit.lq_full!( + A::AbstractBlockSparseMatrix, (L, Q), alg::BlockDiagonalAlgorithm +) + MatrixAlgebraKit.check_input(lq_full!, A, (L, Q), alg) - # 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) + for I in 1:min(blocksize(A)...) + bI = Block(I, I) + if isstored(blocks(A), CartesianIndex(I, I)) # TODO: isblockstored + block = @view!(A[bI]) + block_alg = block_algorithm(alg, block) + bL, bQ = lq_full!(block, block_alg) + L[bI] = bL + Q[bI] = bQ + else + copyto!(@view!(Q[bI]), LinearAlgebra.I) + end end - return LQ + return L, Q end diff --git a/src/factorizations/qr.jl b/src/factorizations/qr.jl index 10c40d8..a564c1e 100644 --- a/src/factorizations/qr.jl +++ b/src/factorizations/qr.jl @@ -1,5 +1,4 @@ -using MatrixAlgebraKit: - MatrixAlgebraKit, default_qr_algorithm, lq_compact!, lq_full!, qr_compact!, qr_full! +using MatrixAlgebraKit: MatrixAlgebraKit, default_qr_algorithm, qr_compact!, qr_full! function MatrixAlgebraKit.default_qr_algorithm( ::Type{<:AbstractBlockSparseMatrix}; kwargs... @@ -9,211 +8,143 @@ function MatrixAlgebraKit.default_qr_algorithm( end 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 +function output_type( + f::Union{typeof(qr_compact!),typeof(qr_full!)}, A::Type{<:AbstractMatrix{T}} +) where {T} + QR = Base.promote_op(f, A) + return isconcretetype(QR) ? QR : Tuple{AbstractMatrix{T},AbstractMatrix{T}} end -function similar_output( - ::typeof(qr_full!), A, R_axis, alg::MatrixAlgebraKit.AbstractAlgorithm +function MatrixAlgebraKit.initialize_output( + ::typeof(qr_compact!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm ) - Q = similar(A, axes(A, 1), R_axis) - R = similar(A, R_axis, axes(A, 2)) - return Q, R + return nothing end - function MatrixAlgebraKit.initialize_output( - ::typeof(qr_compact!), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm + ::typeof(qr_compact!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm ) - 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 - + # using the property that zip stops as soon as one of the iterators is exhausted + r_axes = map(splat(infimum), zip(brows, bcols)) 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) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - Q[brow, bcol], R[bcol, bcol] = MatrixAlgebraKit.initialize_output( - qr_compact!, block, block_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 + BQ, BR = fieldtypes(output_type(qr_compact!, blocktype(A))) + Q = similar(A, BlockType(BQ), (axes(A, 1), r_axis)) + R = similar(A, BlockType(BR), (r_axis, axes(A, 2))) return Q, R end function MatrixAlgebraKit.initialize_output( - ::typeof(qr_full!), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm + ::typeof(qr_full!), ::AbstractBlockSparseMatrix, ::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) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - Q[brow, bcol], R[bcol, bcol] = MatrixAlgebraKit.initialize_output( - qr_full!, block, block_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 nothing +end +function MatrixAlgebraKit.initialize_output( + ::typeof(qr_full!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm +) + BQ, BR = fieldtypes(output_type(qr_full!, blocktype(A))) + Q = similar(A, BlockType(BQ), (axes(A, 1), axes(A, 1))) + R = similar(A, BlockType(BR), (axes(A, 1), axes(A, 2))) return Q, R end function MatrixAlgebraKit.check_input( - ::typeof(qr_compact!), A::AbstractBlockSparseMatrix, QR + ::typeof(qr_compact!), A::AbstractBlockSparseMatrix, QR, ::BlockPermutedDiagonalAlgorithm +) + @assert isblockpermuteddiagonal(A) + return nothing +end +function MatrixAlgebraKit.check_input( + ::typeof(qr_compact!), A::AbstractBlockSparseMatrix, (Q, R), ::BlockDiagonalAlgorithm ) - 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) - + @assert isblockdiagonal(A) return nothing end -function MatrixAlgebraKit.check_input(::typeof(qr_full!), A::AbstractBlockSparseMatrix, QR) - Q, R = QR +function MatrixAlgebraKit.check_input( + ::typeof(qr_full!), A::AbstractBlockSparseMatrix, QR, ::BlockPermutedDiagonalAlgorithm +) + @assert isblockpermuteddiagonal(A) + return nothing +end +function MatrixAlgebraKit.check_input( + ::typeof(qr_full!), A::AbstractBlockSparseMatrix, (Q, R), ::BlockDiagonalAlgorithm +) @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) - + @assert isblockdiagonal(A) return nothing end function MatrixAlgebraKit.qr_compact!( A::AbstractBlockSparseMatrix, QR, alg::BlockPermutedDiagonalAlgorithm ) - MatrixAlgebraKit.check_input(qr_compact!, A, QR) - Q, R = QR + check_input(qr_compact!, A, QR, alg) + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + Qd, Rd = qr_compact!(Ad, BlockDiagonalAlgorithm(alg)) + Q = transform_rows(Qd, invrowperm) + R = transform_cols(Rd, invcolperm) + return Q, R +end - # do decomposition on each block - for bI in eachblockstoredindex(A) - brow, bcol = Tuple(bI) - qr = (@view!(Q[brow, bcol]), @view!(R[bcol, bcol])) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - qr′ = qr_compact!(block, qr, block_alg) - @assert qr === qr′ "qr_compact! might not be in-place" - end +function MatrixAlgebraKit.qr_compact!( + A::AbstractBlockSparseMatrix, (Q, R), alg::BlockDiagonalAlgorithm +) + MatrixAlgebraKit.check_input(qr_compact!, A, (Q, R), alg) - # 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) + # do decomposition on each block + for I in 1:min(blocksize(A)...) + bI = Block(I, I) + if isstored(blocks(A), CartesianIndex(I, I)) # TODO: isblockstored + block = @view!(A[bI]) + block_alg = block_algorithm(alg, block) + bQ, bR = qr_compact!(block, block_alg) + Q[bI] = bQ + R[bI] = bR + else + copyto!(@view!(Q[bI]), LinearAlgebra.I) + end end - return QR + return Q, R 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])) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - qr′ = qr_full!(block, qr, block_alg) - @assert qr === qr′ "qr_full! might not be in-place" - end + check_input(qr_full!, A, QR, alg) + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + Qd, Rd = qr_full!(Ad, BlockDiagonalAlgorithm(alg)) + Q = transform_rows(Qd, invrowperm) + R = transform_cols(Rd, invcolperm) + return Q, R +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 +function MatrixAlgebraKit.qr_full!( + A::AbstractBlockSparseMatrix, (Q, R), alg::BlockDiagonalAlgorithm +) + MatrixAlgebraKit.check_input(qr_full!, A, (Q, R), alg) - # 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) + for I in 1:min(blocksize(A)...) + bI = Block(I, I) + if isstored(blocks(A), CartesianIndex(I, I)) # TODO: isblockstored + block = @view!(A[bI]) + block_alg = block_algorithm(alg, block) + bQ, bR = qr_full!(block, block_alg) + Q[bI] = bQ + R[bI] = bR + else + copyto!(@view!(Q[bI]), LinearAlgebra.I) + end end - return QR + return Q, R end diff --git a/src/factorizations/svd.jl b/src/factorizations/svd.jl index 8c616d4..c2482f7 100644 --- a/src/factorizations/svd.jl +++ b/src/factorizations/svd.jl @@ -1,25 +1,8 @@ using DiagonalArrays: diagonaltype using MatrixAlgebraKit: - MatrixAlgebraKit, check_input, default_svd_algorithm, svd_compact!, svd_full! + MatrixAlgebraKit, check_input, default_svd_algorithm, svd_compact!, svd_full!, svd_vals! using TypeParameterAccessors: realtype -""" - BlockPermutedDiagonalAlgorithm(A::MatrixAlgebraKit.AbstractAlgorithm) - -A wrapper for `MatrixAlgebraKit.AbstractAlgorithm` that implements the wrapped algorithm on -a block-by-block basis, which is possible if the input matrix is a block-diagonal matrix or -a block permuted block-diagonal matrix. -""" -struct BlockPermutedDiagonalAlgorithm{F} <: MatrixAlgebraKit.AbstractAlgorithm - falg::F -end -function block_algorithm(alg::BlockPermutedDiagonalAlgorithm, a::AbstractMatrix) - return block_algorithm(alg, typeof(a)) -end -function block_algorithm(alg::BlockPermutedDiagonalAlgorithm, A::Type{<:AbstractMatrix}) - return alg.falg(A) -end - function MatrixAlgebraKit.default_svd_algorithm( ::Type{<:AbstractBlockSparseMatrix}; kwargs... ) @@ -28,147 +11,88 @@ function MatrixAlgebraKit.default_svd_algorithm( end end -function output_type(::typeof(svd_compact!), A::Type{<:AbstractMatrix{T}}) where {T} - USVᴴ = Base.promote_op(svd_compact!, A) - !isconcretetype(USVᴴ) && - return Tuple{AbstractMatrix{T},AbstractMatrix{realtype(T)},AbstractMatrix{T}} - return USVᴴ +function output_type( + f::Union{typeof(svd_compact!),typeof(svd_full!)}, A::Type{<:AbstractMatrix{T}} +) where {T} + USVᴴ = Base.promote_op(f, A) + return if isconcretetype(USVᴴ) + USVᴴ + else + Tuple{AbstractMatrix{T},AbstractMatrix{realtype(T)},AbstractMatrix{T}} + end +end +function output_type(::typeof(svd_vals!), A::Type{<:AbstractMatrix{T}}) where {T} + S = Base.promote_op(svd_vals!, A) + return isconcretetype(S) ? S : AbstractVector{real(T)} end -function similar_output( - ::typeof(svd_compact!), A, S_axes, alg::MatrixAlgebraKit.AbstractAlgorithm +function MatrixAlgebraKit.initialize_output( + ::typeof(svd_compact!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm +) + return nothing +end +function MatrixAlgebraKit.initialize_output( + ::typeof(svd_compact!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm ) + brows = eachblockaxis(axes(A, 1)) + bcols = eachblockaxis(axes(A, 2)) + # using the property that zip stops as soon as one of the iterators is exhausted + s_axes = map(splat(infimum), zip(brows, bcols)) + s_axis = mortar_axis(s_axes) + S_axes = (s_axis, s_axis) + BU, BS, BVᴴ = fieldtypes(output_type(svd_compact!, blocktype(A))) U = similar(A, BlockType(BU), (axes(A, 1), S_axes[1])) S = similar(A, BlockType(BS), S_axes) Vᴴ = similar(A, BlockType(BVᴴ), (S_axes[2], axes(A, 2))) + return U, S, Vᴴ end function MatrixAlgebraKit.initialize_output( - ::typeof(svd_compact!), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm + ::typeof(svd_full!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm ) - bm, bn = blocksize(A) - bmn = min(bm, bn) - - brows = eachblockaxis(axes(A, 1)) - bcols = eachblockaxis(axes(A, 2)) - u_axes = similar(brows, bmn) - v_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)) - u_axes[col] = infimum(brows[row], bcols[col]) - v_axes[col] = infimum(bcols[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) - u_axes[col] = infimum(brows[row], bcols[col]) - v_axes[col] = infimum(bcols[col], brows[row]) - end - - u_axis = mortar_axis(u_axes) - v_axis = mortar_axis(v_axes) - S_axes = (u_axis, v_axis) - U, S, Vt = similar_output(svd_compact!, A, S_axes, alg) - - # allocate output - for bI in eachblockstoredindex(A) - brow, bcol = Tuple(bI) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - U[brow, bcol], S[bcol, bcol], Vt[bcol, bcol] = MatrixAlgebraKit.initialize_output( - svd_compact!, block, block_alg - ) - end - - # allocate output for blocks that aren't present -- do we also fill identities here? - for (row, col) in zip(emptyrows, emptycols) - @view!(U[Block(row, col)]) - @view!(Vt[Block(col, col)]) - end + return nothing +end +function MatrixAlgebraKit.initialize_output( + ::typeof(svd_full!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm +) + BU, BS, BVᴴ = fieldtypes(output_type(svd_full!, blocktype(A))) + U = similar(A, BlockType(BU), (axes(A, 1), axes(A, 1))) + S = similar(A, BlockType(BS), axes(A)) + Vᴴ = similar(A, BlockType(BVᴴ), (axes(A, 2), axes(A, 2))) - return U, S, Vt + return U, S, Vᴴ end -function similar_output( - ::typeof(svd_full!), A, S_axes, alg::MatrixAlgebraKit.AbstractAlgorithm +function MatrixAlgebraKit.initialize_output( + ::typeof(svd_vals!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm ) - U = similar(A, axes(A, 1), S_axes[1]) - T = real(eltype(A)) - S = similar(A, T, S_axes) - Vt = similar(A, S_axes[2], axes(A, 2)) - return U, S, Vt + return nothing end - function MatrixAlgebraKit.initialize_output( - ::typeof(svd_full!), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm + ::typeof(svd_vals!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm ) - bm, bn = blocksize(A) - brows = eachblockaxis(axes(A, 1)) - u_axes = similar(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)) - u_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) - u_axes[col] = brows[row] - end - for (i, k) in enumerate((length(emptycols) + 1):length(emptyrows)) - u_axes[bn + i] = brows[emptyrows[k]] - end - - u_axis = mortar_axis(u_axes) - S_axes = (u_axis, axes(A, 2)) - U, S, Vt = similar_output(svd_full!, A, S_axes, alg) - - # allocate output - for bI in eachblockstoredindex(A) - brow, bcol = Tuple(bI) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - U[brow, bcol], S[bcol, bcol], Vt[bcol, bcol] = MatrixAlgebraKit.initialize_output( - svd_full!, block, block_alg - ) - end - - # allocate output for blocks that aren't present -- do we also fill identities here? - for (row, col) in zip(emptyrows, emptycols) - @view!(U[Block(row, col)]) - @view!(Vt[Block(col, col)]) - end - # also handle extra rows/cols - for i in (length(emptyrows) + 1):length(emptycols) - @view!(Vt[Block(emptycols[i], emptycols[i])]) - end - for (i, k) in enumerate((length(emptycols) + 1):length(emptyrows)) - @view!(U[Block(emptyrows[k], bn + i)]) - end + bcols = eachblockaxis(axes(A, 2)) + # using the property that zip stops as soon as one of the iterators is exhausted + s_axes = map(splat(infimum), zip(brows, bcols)) + s_axis = mortar_axis(s_axes) - return U, S, Vt + BS = output_type(svd_vals!, blocktype(A)) + return similar(A, BlockType(BS), S_axes) end function MatrixAlgebraKit.check_input( - ::typeof(svd_compact!), A::AbstractBlockSparseMatrix, (U, S, Vᴴ) + ::typeof(svd_compact!), + A::AbstractBlockSparseMatrix, + USVᴴ, + ::BlockPermutedDiagonalAlgorithm, +) + @assert isblockpermuteddiagonal(A) +end +function MatrixAlgebraKit.check_input( + ::typeof(svd_compact!), A::AbstractBlockSparseMatrix, (U, S, Vᴴ), ::BlockDiagonalAlgorithm ) @assert isa(U, AbstractBlockSparseMatrix) && isa(S, AbstractBlockSparseMatrix) && @@ -177,11 +101,18 @@ function MatrixAlgebraKit.check_input( @assert real(eltype(A)) == eltype(S) @assert axes(A, 1) == axes(U, 1) && axes(A, 2) == axes(Vᴴ, 2) @assert axes(S, 1) == axes(S, 2) + @assert isblockdiagonal(A) return nothing end function MatrixAlgebraKit.check_input( - ::typeof(svd_full!), A::AbstractBlockSparseMatrix, (U, S, Vᴴ) + ::typeof(svd_full!), A::AbstractBlockSparseMatrix, USVᴴ, ::BlockPermutedDiagonalAlgorithm +) + @assert isblockpermuteddiagonal(A) + return nothing +end +function MatrixAlgebraKit.check_input( + ::typeof(svd_full!), A::AbstractBlockSparseMatrix, (U, S, Vᴴ), ::BlockDiagonalAlgorithm ) @assert isa(U, AbstractBlockSparseMatrix) && isa(S, AbstractBlockSparseMatrix) && @@ -190,78 +121,119 @@ function MatrixAlgebraKit.check_input( @assert real(eltype(A)) == eltype(S) @assert axes(A, 1) == axes(U, 1) && axes(A, 2) == axes(Vᴴ, 1) == axes(Vᴴ, 2) @assert axes(S, 2) == axes(A, 2) + @assert isblockdiagonal(A) + return nothing +end + +function MatrixAlgebraKit.check_input( + ::typeof(svd_vals!), A::AbstractBlockSparseMatrix, S, ::BlockPermutedDiagonalAlgorithm +) + @assert isblockpermuteddiagonal(A) + return nothing +end +function MatrixAlgebraKit.check_input( + ::typeof(svd_vals!), A::AbstractBlockSparseMatrix, S, ::BlockDiagonalAlgorithm +) + @assert isa(S, AbstractBlockSparseVector) + @assert real(eltype(A)) == eltype(S) + @assert isblockdiagonal(A) return nothing end function MatrixAlgebraKit.svd_compact!( - A::AbstractBlockSparseMatrix, (U, S, Vᴴ), alg::BlockPermutedDiagonalAlgorithm + A::AbstractBlockSparseMatrix, USVᴴ, alg::BlockPermutedDiagonalAlgorithm ) - check_input(svd_compact!, A, (U, S, Vᴴ)) + check_input(svd_compact!, A, USVᴴ, alg) - # do decomposition on each block - for bI in eachblockstoredindex(A) - brow, bcol = Tuple(bI) - usvᴴ = (@view!(U[brow, bcol]), @view!(S[bcol, bcol]), @view!(Vᴴ[bcol, bcol])) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - usvᴴ′ = svd_compact!(block, usvᴴ, block_alg) - @assert usvᴴ === usvᴴ′ "svd_compact! might not be in-place" - end + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + Ud, S, Vᴴd = svd_compact!(Ad, BlockDiagonalAlgorithm(alg)) + U = transform_rows(Ud, invrowperm) + Vᴴ = transform_cols(Vᴴd, invcolperm) + + return U, S, Vᴴ +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 - # U[Block(row, col)] = LinearAlgebra.I - # Vᴴ[Block(col, col)] = LinearAlgebra.I - for (row, col) in zip(emptyrows, emptycols) - copyto!(@view!(U[Block(row, col)]), LinearAlgebra.I) - copyto!(@view!(Vᴴ[Block(col, col)]), LinearAlgebra.I) +function MatrixAlgebraKit.svd_compact!( + A::AbstractBlockSparseMatrix, (U, S, Vᴴ), alg::BlockDiagonalAlgorithm +) + check_input(svd_compact!, A, (U, S, Vᴴ), alg) + + for I in 1:min(blocksize(A)...) + bI = Block(I, I) + if isstored(blocks(A), CartesianIndex(I, I)) # TODO: isblockstored + block = @view!(A[bI]) + block_alg = block_algorithm(alg, block) + bU, bS, bVᴴ = svd_compact!(block, block_alg) + U[bI] = bU + S[bI] = bS + Vᴴ[bI] = bVᴴ + else + copyto!(@view!(U[bI]), LinearAlgebra.I) + copyto!(@view!(Vᴴ[bI]), LinearAlgebra.I) + end end - return (U, S, Vᴴ) + return U, S, Vᴴ end function MatrixAlgebraKit.svd_full!( - A::AbstractBlockSparseMatrix, (U, S, Vᴴ), alg::BlockPermutedDiagonalAlgorithm + A::AbstractBlockSparseMatrix, USVᴴ, alg::BlockPermutedDiagonalAlgorithm ) - check_input(svd_full!, A, (U, S, Vᴴ)) + check_input(svd_full!, A, USVᴴ, alg) - # do decomposition on each block - for bI in eachblockstoredindex(A) - brow, bcol = Tuple(bI) - usvᴴ = (@view!(U[brow, bcol]), @view!(S[bcol, bcol]), @view!(Vᴴ[bcol, bcol])) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - usvᴴ′ = svd_full!(block, usvᴴ, block_alg) - @assert usvᴴ === usvᴴ′ "svd_full! might not be in-place" - end + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + Ud, S, Vᴴd = svd_full!(Ad, BlockDiagonalAlgorithm(alg)) + U = transform_rows(Ud, invrowperm) + Vᴴ = transform_cols(Vᴴd, invcolperm) + + return U, S, Vᴴ +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 - # U[Block(row, col)] = LinearAlgebra.I - # Vt[Block(col, col)] = LinearAlgebra.I - for (row, col) in zip(emptyrows, emptycols) - copyto!(@view!(U[Block(row, col)]), LinearAlgebra.I) - copyto!(@view!(Vᴴ[Block(col, col)]), LinearAlgebra.I) +function MatrixAlgebraKit.svd_full!( + A::AbstractBlockSparseMatrix, (U, S, Vᴴ), alg::BlockDiagonalAlgorithm +) + check_input(svd_full!, A, (U, S, Vᴴ), alg) + + for I in 1:min(blocksize(A)...) + bI = Block(I, I) + if isstored(blocks(A), CartesianIndex(I, I)) # TODO: isblockstored + block = @view!(A[bI]) + block_alg = block_algorithm(alg, block) + bU, bS, bVᴴ = svd_full!(block, block_alg) + U[bI] = bU + S[bI] = bS + Vᴴ[bI] = bVᴴ + else + copyto!(@view!(U[bI]), LinearAlgebra.I) + copyto!(@view!(Vᴴ[bI]), LinearAlgebra.I) + end end - # also handle extra rows/cols - for i in (length(emptyrows) + 1):length(emptycols) - copyto!(@view!(Vᴴ[Block(emptycols[i], emptycols[i])]), LinearAlgebra.I) + # Complete the unitaries for rectangular inputs + for I in (blocksize(A, 2) + 1):blocksize(A, 1) + copyto!(@view!(U[Block(I, I)]), LinearAlgebra.I) end - bn = blocksize(A, 2) - for (i, k) in enumerate((length(emptycols) + 1):length(emptyrows)) - copyto!(@view!(U[Block(emptyrows[k], bn + i)]), LinearAlgebra.I) + for I in (blocksize(A, 1) + 1):blocksize(A, 2) + copyto!(@view!(Vᴴ[Block(I, I)]), LinearAlgebra.I) end - return (U, S, Vᴴ) + return U, S, Vᴴ +end + +function MatrixAlgebraKit.svd_vals!( + A::AbstractBlockSparseMatrix, S, alg::BlockPermutedDiagonalAlgorithm +) + MatrixAlgebraKit.check_input(svd_vals!, A, S, alg) + Ad, _ = blockdiagonalize(A) + return svd_vals!(Ad, BlockDiagonalAlgorithm(alg)) +end +function MatrixAlgebraKit.svd_vals!( + A::AbstractBlockSparseMatrix, S, alg::BlockDiagonalAlgorithm +) + MatrixAlgebraKit.check_input(svd_vals!, A, S, alg) + for I in eachblockstoredindex(A) + block = @view!(A[I]) + S[Tuple(I)[1]] = svd_vals!(block, block_algorithm(alg, block)) + end + return S end diff --git a/src/factorizations/tensorproducts.jl b/src/factorizations/tensorproducts.jl deleted file mode 100644 index bbf02ae..0000000 --- a/src/factorizations/tensorproducts.jl +++ /dev/null @@ -1,19 +0,0 @@ -function infimum(r1::AbstractUnitRange, r2::AbstractUnitRange) - (isone(first(r1)) && isone(first(r2))) || - throw(ArgumentError("infimum only defined for ranges starting at 1")) - if length(r1) ≤ length(r2) - return r1 - else - return r1[r2] - end -end - -function supremum(r1::AbstractUnitRange, r2::AbstractUnitRange) - (isone(first(r1)) && isone(first(r2))) || - throw(ArgumentError("supremum only defined for ranges starting at 1")) - if length(r1) ≥ length(r2) - return r1 - else - return r2 - end -end diff --git a/src/factorizations/truncation.jl b/src/factorizations/truncation.jl index d173f88..a28995b 100644 --- a/src/factorizations/truncation.jl +++ b/src/factorizations/truncation.jl @@ -1,4 +1,6 @@ using MatrixAlgebraKit: + MatrixAlgebraKit, + TruncatedAlgorithm, TruncationStrategy, diagview, eig_trunc!, @@ -8,34 +10,50 @@ using MatrixAlgebraKit: truncate! """ - BlockPermutedDiagonalTruncationStrategy(strategy::TruncationStrategy) + BlockDiagonalTruncationStrategy(strategy::TruncationStrategy) A wrapper for `TruncationStrategy` that implements the wrapped strategy on a block-by-block -basis, which is possible if the input matrix is a block-diagonal matrix or a block permuted -block-diagonal matrix. +basis, which is possible if the input matrix is a block-diagonal matrix. """ -struct BlockPermutedDiagonalTruncationStrategy{T<:TruncationStrategy} <: TruncationStrategy +struct BlockDiagonalTruncationStrategy{T<:TruncationStrategy} <: TruncationStrategy strategy::T end -function MatrixAlgebraKit.truncate!( - ::typeof(svd_trunc!), - (U, S, Vᴴ)::NTuple{3,AbstractBlockSparseMatrix}, - strategy::TruncationStrategy, +function BlockDiagonalTruncationStrategy(alg::BlockPermutedDiagonalAlgorithm) + return BlockDiagonalTruncationStrategy(alg.strategy) +end + +function MatrixAlgebraKit.svd_trunc!( + A::AbstractBlockSparseMatrix, + out, + alg::TruncatedAlgorithm{<:BlockPermutedDiagonalAlgorithm}, ) - # TODO assert blockdiagonal - return truncate!( - svd_trunc!, (U, S, Vᴴ), BlockPermutedDiagonalTruncationStrategy(strategy) - ) + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + blockalg = BlockDiagonalAlgorithm(alg.alg) + blockstrategy = BlockDiagonalTruncationStrategy(alg.trunc) + Ud, S, Vᴴd = svd_trunc!(Ad, TruncatedAlgorithm(blockalg, blockstrategy)) + + U = transform_rows(Ud, invrowperm) + Vᴴ = transform_cols(Vᴴd, invcolperm) + + return U, S, Vᴴ end + for f in [:eig_trunc!, :eigh_trunc!] @eval begin - function MatrixAlgebraKit.truncate!( - ::typeof($f), - (D, V)::NTuple{2,AbstractBlockSparseMatrix}, - strategy::TruncationStrategy, + function MatrixAlgebraKit.$f( + A::AbstractBlockSparseMatrix, + out, + alg::TruncatedAlgorithm{<:BlockPermutedDiagonalAlgorithm}, ) - return truncate!($f, (D, V), BlockPermutedDiagonalTruncationStrategy(strategy)) + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + blockalg = BlockDiagonalAlgorithm(alg.alg) + blockstrategy = BlockDiagonalTruncationStrategy(alg.trunc) + Dd, Vd = $f(Ad, TruncatedAlgorithm(blockalg, blockstrategy)) + + D = transform_rows(Dd, invrowperm) + V = transform_cols(Vd, invcolperm) + return D, V end end end @@ -43,7 +61,7 @@ end # cannot use regular slicing here: I want to slice without altering blockstructure # solution: use boolean indexing and slice the mask, effectively cheaply inverting the map function MatrixAlgebraKit.findtruncated( - values::AbstractVector, strategy::BlockPermutedDiagonalTruncationStrategy + values::AbstractVector, strategy::BlockDiagonalTruncationStrategy ) ind = findtruncated(Vector(values), strategy.strategy) indexmask = falses(length(values)) @@ -66,7 +84,7 @@ end function MatrixAlgebraKit.truncate!( ::typeof(svd_trunc!), (U, S, Vᴴ)::NTuple{3,AbstractBlockSparseMatrix}, - strategy::BlockPermutedDiagonalTruncationStrategy, + strategy::BlockDiagonalTruncationStrategy, ) I = findtruncated(diag(S), strategy) return (U[:, I], S[I, I], Vᴴ[I, :]) @@ -76,7 +94,7 @@ for f in [:eig_trunc!, :eigh_trunc!] function MatrixAlgebraKit.truncate!( ::typeof($f), (D, V)::NTuple{2,AbstractBlockSparseMatrix}, - strategy::BlockPermutedDiagonalTruncationStrategy, + strategy::BlockDiagonalTruncationStrategy, ) I = findtruncated(diag(D), strategy) return (D[I, I], V[:, I]) diff --git a/src/factorizations/utility.jl b/src/factorizations/utility.jl new file mode 100644 index 0000000..b269c0f --- /dev/null +++ b/src/factorizations/utility.jl @@ -0,0 +1,110 @@ +using MatrixAlgebraKit: MatrixAlgebraKit + +function infimum(r1::AbstractUnitRange, r2::AbstractUnitRange) + (isone(first(r1)) && isone(first(r2))) || + throw(ArgumentError("infimum only defined for ranges starting at 1")) + if length(r1) ≤ length(r2) + return r1 + else + return r1[r2] + end +end + +function supremum(r1::AbstractUnitRange, r2::AbstractUnitRange) + (isone(first(r1)) && isone(first(r2))) || + throw(ArgumentError("supremum only defined for ranges starting at 1")) + if length(r1) ≥ length(r2) + return r1 + else + return r2 + end +end + +transform_rows(A::AbstractMatrix, invrowperm) = A[invrowperm, :] +transform_rows(A::AbstractVector, invrowperm) = A[invrowperm] +transform_cols(A::AbstractMatrix, invcolperm) = A[:, invcolperm] + +function blockdiagonalize(A::AbstractBlockSparseMatrix) + # sort in order to avoid depending on internal details such as dictionary order + bIs = sort!(collect(eachblockstoredindex(A)); by=Int ∘ last ∘ Tuple) + + # obtain permutation for the blocks that are present + rowperm = map(first ∘ Tuple, bIs) + colperm = map(last ∘ Tuple, bIs) + + # These checks might be expensive but are convenient for now + allunique(rowperm) || throw(ArgumentError("input contains more than one block per row")) + allunique(colperm) || + throw(ArgumentError("input contains more than one block per column")) + + # post-process empty rows and columns: this pairs them up in order of occurance, + # putting empty rows and columns due to rectangularity at the end + emptyrows = setdiff(Block.(1:blocksize(A, 1)), rowperm) + append!(rowperm, emptyrows) + emptycols = setdiff(Block.(1:blocksize(A, 2)), colperm) + append!(colperm, emptycols) + + invrowperm = Block.(invperm(Int.(rowperm))) + invcolperm = Block.(invperm(Int.(colperm))) + + return A[rowperm, colperm], (invrowperm, invcolperm) +end + +function isblockdiagonal(A::AbstractBlockSparseMatrix) + for bI in eachblockstoredindex(A) + row, col = Tuple(bI) + row == col || return false + end + # don't need to check for rows and cols appearing only once + # this is guaranteed as long as eachblockstoredindex is unique, which we assume + return true +end + +function isblockpermuteddiagonal(A::AbstractBlockSparseMatrix) + return allunique(first.(eachblockstoredindex(A))) && + allunique(last.(eachblockstoredindex(A))) +end + +""" + BlockDiagonalAlgorithm([f]) <: MatrixAlgebraKit.AbstractAlgorithm + +Type for handling algorithms on a block-by-block basis, which is possible for +block-diagonal input matrices. +Additionally this algorithm may take a function that, given the individual blocks, returns +the algorithm that will be used. This can be leveraged to allow for different algorithms for +each block. +""" +struct BlockDiagonalAlgorithm{F} <: MatrixAlgebraKit.AbstractAlgorithm + falg::F +end + +function block_algorithm(alg::BlockDiagonalAlgorithm, a::AbstractMatrix) + return block_algorithm(alg, typeof(a)) +end +function block_algorithm(alg::BlockDiagonalAlgorithm, A::Type{<:AbstractMatrix}) + return alg.falg(A) +end + +""" + BlockPermutedDiagonalAlgorithm([f]) <: MatrixAlgebraKit.AbstractAlgorithm + +Type for handling algorithms on a block-by-block basis, which is possible for +block-diagonal or block-permuted-diagonal input matrices. The algorithms proceed by first +permuting to a block-diagonal form, and then carrying out the algorithm. +Additionally this algorithm may take a function that, given the individual blocks, returns +the algorithm that will be used. This can be leveraged to allow for different algorithms for +each block. +""" +struct BlockPermutedDiagonalAlgorithm{F} <: MatrixAlgebraKit.AbstractAlgorithm + falg::F +end +function block_algorithm(alg::BlockPermutedDiagonalAlgorithm, a::AbstractMatrix) + return block_algorithm(alg, typeof(a)) +end +function block_algorithm(alg::BlockPermutedDiagonalAlgorithm, A::Type{<:AbstractMatrix}) + return alg.falg(A) +end + +function BlockDiagonalAlgorithm(alg::BlockPermutedDiagonalAlgorithm) + return BlockDiagonalAlgorithm(alg.falg) +end diff --git a/test/Project.toml b/test/Project.toml index 12fb6f0..c084b9f 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -25,7 +25,7 @@ Adapt = "4" Aqua = "0.8" ArrayLayouts = "1" BlockArrays = "1" -BlockSparseArrays = "0.9" +BlockSparseArrays = "0.10" DerivableInterfaces = "0.5" DiagonalArrays = "0.3" GPUArraysCore = "0.2" diff --git a/test/test_factorizations.jl b/test/test_factorizations.jl index 651a8f6..cb60b26 100644 --- a/test/test_factorizations.jl +++ b/test/test_factorizations.jl @@ -119,15 +119,15 @@ end function test_svd(a, (U, S, Vᴴ); full=false) # Check that the SVD is correct - (U * S * Vᴴ ≈ a) || return false - (U' * U ≈ LinearAlgebra.I) || return false - (Vᴴ * Vᴴ' ≈ LinearAlgebra.I) || return false - full || return true + @test (U * S * Vᴴ ≈ a) + @test (U' * U ≈ LinearAlgebra.I) + @test (Vᴴ * Vᴴ' ≈ LinearAlgebra.I) + full || return nothing # Check factors are unitary - (U * U' ≈ LinearAlgebra.I) || return false - (Vᴴ' * Vᴴ ≈ LinearAlgebra.I) || return false - return true + @test (U * U' ≈ LinearAlgebra.I) + @test (Vᴴ' * Vᴴ ≈ LinearAlgebra.I) + return nothing end blockszs = ( @@ -143,7 +143,7 @@ test_params = Iterators.product(blockszs, eltypes) # test empty matrix usv_empty = svd_compact(a) - @test test_svd(a, usv_empty) + test_svd(a, usv_empty) # test blockdiagonal rng = StableRNG(123) @@ -152,13 +152,13 @@ test_params = Iterators.product(blockszs, eltypes) a[Block(I.I...)] = rand(rng, T, size(blocks(a)[i])) end usv = svd_compact(a) - @test test_svd(a, usv) + test_svd(a, usv) rng = StableRNG(123) perm = Random.randperm(rng, length(m)) b = a[Block.(perm), Block.(1:length(n))] usv = svd_compact(b) - @test test_svd(b, usv) + test_svd(b, usv) # test permuted blockdiagonal with missing row/col rng = StableRNG(123) @@ -166,7 +166,7 @@ test_params = Iterators.product(blockszs, eltypes) c = copy(b) delete!(blocks(c).storage, CartesianIndex(Int.(Tuple(I_removed)))) usv = svd_compact(c) - @test test_svd(c, usv) + test_svd(c, usv) end # svd_full! @@ -176,7 +176,7 @@ end # test empty matrix usv_empty = svd_full(a) - @test test_svd(a, usv_empty; full=true) + test_svd(a, usv_empty; full=true) # test blockdiagonal rng = StableRNG(123) @@ -185,13 +185,13 @@ end a[Block(I.I...)] = rand(rng, T, size(blocks(a)[i])) end usv = svd_full(a) - @test test_svd(a, usv; full=true) + test_svd(a, usv; full=true) rng = StableRNG(123) perm = Random.randperm(rng, length(m)) b = a[Block.(perm), Block.(1:length(n))] usv = svd_full(b) - @test test_svd(b, usv; full=true) + test_svd(b, usv; full=true) # test permuted blockdiagonal with missing row/col rng = StableRNG(123) @@ -199,7 +199,7 @@ end c = copy(b) delete!(blocks(c).storage, CartesianIndex(Int.(Tuple(I_removed)))) usv = svd_full(c) - @test test_svd(c, usv; full=true) + test_svd(c, usv; full=true) end # svd_trunc! diff --git a/test/test_issues.jl b/test/test_issues.jl new file mode 100644 index 0000000..0ce92ca --- /dev/null +++ b/test/test_issues.jl @@ -0,0 +1,25 @@ +using BlockArrays +using BlockSparseArrays +using BlockSparseArrays: blocksparse +using MatrixAlgebraKit +using LinearAlgebra: LinearAlgebra +using Test: @test, @testset + +@testset "Issue 162" begin + ax = (blockedrange([2, 2]), blockedrange([2, 2, 2])) + bs = Dict(Block(1, 1) => randn(2, 2), Block(2, 3) => randn(2, 2)) + a = blocksparse(bs, ax) + U, S, Vᴴ = svd_compact(a) + + @test U * S * Vᴴ ≈ a + @test U' * U ≈ LinearAlgebra.I + @test Vᴴ * Vᴴ' ≈ LinearAlgebra.I + + U, S, Vᴴ = svd_full(a); + + @test U * S * Vᴴ ≈ a + @test U' * U ≈ LinearAlgebra.I + @test U * U' ≈ LinearAlgebra.I + @test Vᴴ * Vᴴ' ≈ LinearAlgebra.I + @test Vᴴ' * Vᴴ ≈ LinearAlgebra.I +end