diff --git a/ext/BlockSparseArraysGradedUnitRangesExt/src/BlockSparseArraysGradedUnitRangesExt.jl b/ext/BlockSparseArraysGradedUnitRangesExt/src/BlockSparseArraysGradedUnitRangesExt.jl index e10957e..5846ce5 100644 --- a/ext/BlockSparseArraysGradedUnitRangesExt/src/BlockSparseArraysGradedUnitRangesExt.jl +++ b/ext/BlockSparseArraysGradedUnitRangesExt/src/BlockSparseArraysGradedUnitRangesExt.jl @@ -9,11 +9,14 @@ using BlockArrays: using ..BlockSparseArrays: BlockSparseArrays, AbstractBlockSparseArray, + AbstractBlockSparseArrayInterface, AbstractBlockSparseMatrix, BlockSparseArray, + BlockSparseArrayInterface, BlockSparseMatrix, BlockSparseVector, block_merge +using Derive: @interface using GradedUnitRanges: GradedUnitRanges, AbstractGradedUnitRange, @@ -109,7 +112,7 @@ end # with mixed dual and non-dual axes. This shouldn't be needed once # GradedUnitRanges is rewritten using BlockArrays v1. # TODO: Delete this once GradedUnitRanges is rewritten. -function blocksparse_show( +@interface ::AbstractBlockSparseArrayInterface function Base.show( io::IO, mime::MIME"text/plain", a::AbstractArray, axes_a::Tuple; kwargs... ) println(io, "typeof(axes) = ", typeof(axes_a), "\n") @@ -127,7 +130,7 @@ end function Base.show(io::IO, mime::MIME"text/plain", a::BlockSparseArray; kwargs...) axes_a = axes(a) a_nondual = BlockSparseArray(blocks(a), nondual.(axes(a))) - return blocksparse_show(io, mime, a_nondual, axes_a; kwargs...) + return @interface BlockSparseArrayInterface() show(io, mime, a_nondual, axes_a; kwargs...) end # This is a temporary fix for `show` being broken for BlockSparseArrays @@ -139,7 +142,7 @@ function Base.show( ) axes_a = axes(a) a_nondual = BlockSparseArray(blocks(a'), dual.(nondual.(axes(a'))))' - return blocksparse_show(io, mime, a_nondual, axes_a; kwargs...) + return @interface BlockSparseArrayInterface() show(io, mime, a_nondual, axes_a; kwargs...) end # This is a temporary fix for `show` being broken for BlockSparseArrays @@ -151,6 +154,6 @@ function Base.show( ) axes_a = axes(a) a_nondual = tranpose(BlockSparseArray(transpose(blocks(a)), nondual.(axes(a)))) - return blocksparse_show(io, mime, a_nondual, axes_a; kwargs...) + return @interface BlockSparseArrayInterface() show(io, mime, a_nondual, axes_a; kwargs...) end end diff --git a/src/BlockArraysExtensions/BlockArraysExtensions.jl b/src/BlockArraysExtensions/BlockArraysExtensions.jl index 5dcd190..3208508 100644 --- a/src/BlockArraysExtensions/BlockArraysExtensions.jl +++ b/src/BlockArraysExtensions/BlockArraysExtensions.jl @@ -104,7 +104,7 @@ Base.view(S::BlockIndices, i) = S[i] # @view b[Block(1, 1)[1:2, 2:2]] # ``` # This is similar to the definition: -# blocksparse_to_indices(a, inds, I::Tuple{UnitRange{<:Integer},Vararg{Any}}) +# @interface BlockSparseArrayInterface() to_indices(a, inds, I::Tuple{UnitRange{<:Integer},Vararg{Any}}) function Base.getindex( a::NonBlockedVector{<:Integer,<:BlockIndices}, I::UnitRange{<:Integer} ) diff --git a/src/abstractblocksparsearray/abstractblocksparsearray.jl b/src/abstractblocksparsearray/abstractblocksparsearray.jl index c1fff56..c3297c1 100644 --- a/src/abstractblocksparsearray/abstractblocksparsearray.jl +++ b/src/abstractblocksparsearray/abstractblocksparsearray.jl @@ -19,12 +19,12 @@ end # Specialized in order to fix ambiguity error with `BlockArrays`. function Base.getindex(a::AbstractBlockSparseArray{<:Any,N}, I::Vararg{Int,N}) where {N} - return blocksparse_getindex(a, I...) + return @interface BlockSparseArrayInterface() getindex(a, I...) end # Specialized in order to fix ambiguity error with `BlockArrays`. function Base.getindex(a::AbstractBlockSparseArray{<:Any,0}) - return blocksparse_getindex(a) + return @interface BlockSparseArrayInterface() getindex(a) end ## # Fix ambiguity error with `BlockArrays`. @@ -39,7 +39,7 @@ end ## ## # Fix ambiguity error with `BlockArrays`. ## function Base.getindex(a::AbstractBlockSparseArray, I::Vararg{AbstractVector}) -## ## return blocksparse_getindex(a, I...) +## ## return @interface BlockSparseArrayInterface() getindex(a, I...) ## return ArrayLayouts.layout_getindex(a, I...) ## end @@ -47,13 +47,13 @@ end function Base.setindex!( a::AbstractBlockSparseArray{<:Any,N}, value, I::Vararg{Int,N} ) where {N} - blocksparse_setindex!(a, value, I...) + @interface BlockSparseArrayInterface() setindex!(a, value, I...) return a end # Fix ambiguity error. function Base.setindex!(a::AbstractBlockSparseArray{<:Any,0}, value) - blocksparse_setindex!(a, value) + @interface BlockSparseArrayInterface() setindex!(a, value) return a end diff --git a/src/abstractblocksparsearray/views.jl b/src/abstractblocksparsearray/views.jl index 6e9c376..9c22ab9 100644 --- a/src/abstractblocksparsearray/views.jl +++ b/src/abstractblocksparsearray/views.jl @@ -39,7 +39,7 @@ function Base.view( }, I::Block{N}, ) where {N} - return blocksparse_view(a, I) + return @interface BlockSparseArrayInterface() view(a, I) end function Base.view( a::SubArray{ @@ -47,13 +47,13 @@ function Base.view( }, I::Vararg{Block{1},N}, ) where {N} - return blocksparse_view(a, I...) + return @interface BlockSparseArrayInterface() view(a, I...) end function Base.view( V::SubArray{<:Any,1,<:AnyAbstractBlockSparseArray,<:Tuple{BlockSlice{<:BlockRange{1}}}}, I::Block{1}, ) - return blocksparse_view(a, I) + return @interface BlockSparseArrayInterface() view(a, I) end # Specialized code for getting the view of a block. @@ -63,7 +63,7 @@ function BlockArrays.viewblock( return viewblock(a, Tuple(block)...) end -# TODO: Define `blocksparse_viewblock`. +# TODO: Define `@interface BlockSparseArrayInterface() viewblock`. function BlockArrays.viewblock( a::AbstractBlockSparseArray{<:Any,N}, block::Vararg{Block{1},N} ) where {N} @@ -177,9 +177,9 @@ end # XXX: TODO: Distinguish if a sub-view of the block needs to be taken! # Define a new `SubBlockSlice` which is used in: -# `blocksparse_to_indices(a, inds, I::Tuple{UnitRange{<:Integer},Vararg{Any}})` +# `@interface BlockSparseArrayInterface() to_indices(a, inds, I::Tuple{UnitRange{<:Integer},Vararg{Any}})` # in `blocksparsearrayinterface/blocksparsearrayinterface.jl`. -# TODO: Define `blocksparse_viewblock`. +# TODO: Define `@interface BlockSparseArrayInterface() viewblock`. function BlockArrays.viewblock( a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockSliceCollection,N}}}, block::Vararg{Block{1},N}, @@ -220,7 +220,7 @@ function BlockArrays.viewblock( return @view parent(a)[brs...] end -# TODO: Define `blocksparse_viewblock`. +# TODO: Define `@interface BlockSparseArrayInterface() viewblock`. function BlockArrays.viewblock( a::SubArray{ T, @@ -257,7 +257,7 @@ function BlockArrays.viewblock( ) where {T,N} return viewblock(a, to_tuple(block)...) end -# TODO: Define `blocksparse_viewblock`. +# TODO: Define `@interface BlockSparseArrayInterface() viewblock`. function BlockArrays.viewblock( a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockedSlice,N}}}, I::Vararg{Block{1},N}, @@ -271,7 +271,7 @@ function BlockArrays.viewblock( end return @view parent(a)[brs...] end -# TODO: Define `blocksparse_viewblock`. +# TODO: Define `@interface BlockSparseArrayInterface() viewblock`. function BlockArrays.viewblock( a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockedSlice,N}}}, block::Vararg{BlockIndexRange{1},N}, diff --git a/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl index 63ceabf..730353b 100644 --- a/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl +++ b/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl @@ -9,7 +9,7 @@ using BlockArrays: blockedrange, mortar, unblock -using Derive: Derive +using Derive: Derive, @interface using SplitApplyCombine: groupcount using TypeParameterAccessors: similartype @@ -28,14 +28,14 @@ Derive.interface(::Type{<:AnyAbstractBlockSparseArray}) = BlockSparseArrayInterf function Base.to_indices( a::AnyAbstractBlockSparseArray, inds, I::Tuple{UnitRange{<:Integer},Vararg{Any}} ) - return blocksparse_to_indices(a, inds, I) + return @interface BlockSparseArrayInterface() to_indices(a, inds, I) end # a[[Block(2), Block(1)], [Block(2), Block(1)]] function Base.to_indices( a::AnyAbstractBlockSparseArray, inds, I::Tuple{Vector{<:Block{1}},Vararg{Any}} ) - return blocksparse_to_indices(a, inds, I) + return @interface BlockSparseArrayInterface() to_indices(a, inds, I) end # a[BlockVector([Block(2), Block(1)], [2]), BlockVector([Block(2), Block(1)], [2])] @@ -45,7 +45,7 @@ function Base.to_indices( inds, I::Tuple{AbstractBlockVector{<:Block{1}},Vararg{Any}}, ) - return blocksparse_to_indices(a, inds, I) + return @interface BlockSparseArrayInterface() to_indices(a, inds, I) end # a[mortar([Block(1)[1:2], Block(2)[1:3]])] @@ -54,7 +54,7 @@ function Base.to_indices( inds, I::Tuple{BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}},Vararg{Any}}, ) - return blocksparse_to_indices(a, inds, I) + return @interface BlockSparseArrayInterface() to_indices(a, inds, I) end # a[[Block(1)[1:2], Block(2)[1:2]], [Block(1)[1:2], Block(2)[1:2]]] @@ -65,14 +65,16 @@ function Base.to_indices( end # BlockArrays `AbstractBlockArray` interface -BlockArrays.blocks(a::AnyAbstractBlockSparseArray) = blocksparse_blocks(a) +function BlockArrays.blocks(a::AnyAbstractBlockSparseArray) + @interface BlockSparseArrayInterface() blocks(a) +end # Fix ambiguity error with `BlockArrays` using BlockArrays: BlockSlice function BlockArrays.blocks( a::SubArray{<:Any,<:Any,<:AbstractBlockSparseArray,<:Tuple{Vararg{BlockSlice}}} ) - return blocksparse_blocks(a) + return @interface BlockSparseArrayInterface() blocks(a) end using TypeParameterAccessors: parenttype @@ -105,7 +107,7 @@ function Base.getindex(a::AnyAbstractBlockSparseArray{<:Any,0}) return ArrayLayouts.layout_getindex(a) end -# TODO: Define `blocksparse_isassigned`. +# TODO: Define `@interface BlockSparseArrayInterface() isassigned`. function Base.isassigned( a::AnyAbstractBlockSparseArray{<:Any,N}, index::Vararg{Block{1},N} ) where {N} @@ -121,7 +123,7 @@ function Base.isassigned(a::AnyAbstractBlockSparseArray{<:Any,N}, index::Block{N return isassigned(a, Tuple(index)...) end -# TODO: Define `blocksparse_isassigned`. +# TODO: Define `@interface BlockSparseArrayInterface() isassigned`. function Base.isassigned( a::AnyAbstractBlockSparseArray{<:Any,N}, index::Vararg{BlockIndex{1},N} ) where {N} @@ -133,13 +135,13 @@ function Base.setindex!( a::AnyAbstractBlockSparseArray{<:Any,N}, value, I::BlockIndex{N} ) where {N} # TODO: Use `@interface interface(a) setindex!(...)`. - blocksparse_setindex!(a, value, I) + @interface BlockSparseArrayInterface() setindex!(a, value, I) return a end # Fixes ambiguity error with BlockArrays.jl function Base.setindex!(a::AnyAbstractBlockSparseArray{<:Any,1}, value, I::BlockIndex{1}) # TODO: Use `@interface interface(a) setindex!(...)`. - blocksparse_setindex!(a, value, I) + @interface BlockSparseArrayInterface() setindex!(a, value, I) return a end @@ -212,32 +214,51 @@ function blocksparse_similar(a, elt::Type, axes::Tuple) undef, axes ) end +@interface ::AbstractBlockSparseArrayInterface function Base.similar( + a::AbstractArray, elt::Type, axes::Tuple{Vararg{Int}} +) + return blocksparse_similar(a, elt, axes) +end +@interface ::AbstractBlockSparseArrayInterface function Base.similar( + a::AbstractArray, elt::Type, axes::Tuple +) + return blocksparse_similar(a, elt, axes) +end +@interface ::AbstractBlockSparseArrayInterface function Base.similar( + a::Type{<:AbstractArray}, elt::Type, axes::Tuple{Vararg{Int}} +) + return blocksparse_similar(a, elt, axes) +end +@interface ::AbstractBlockSparseArrayInterface function Base.similar( + a::Type{<:AbstractArray}, elt::Type, axes::Tuple +) + return blocksparse_similar(a, elt, axes) +end # Needed by `BlockArrays` matrix multiplication interface -# TODO: Define a `blocksparse_similar` function. +# TODO: Define a `@interface BlockSparseArrayInterface() similar` function. function Base.similar( arraytype::Type{<:AnyAbstractBlockSparseArray}, elt::Type, axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}}, ) - # TODO: Use `@interface interface(arraytype) similar(...)`. - return blocksparse_similar(arraytype, elt, axes) + return @interface BlockSparseArrayInterface() similar(arraytype, elt, axes) end -# TODO: Define a `blocksparse_similar` function. +# TODO: Define a `@interface BlockSparseArrayInterface() similar` function. function Base.similar( a::AnyAbstractBlockSparseArray, elt::Type, axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}}, ) # TODO: Use `@interface interface(a) similar(...)`. - return blocksparse_similar(a, elt, axes) + return @interface BlockSparseArrayInterface() similar(a, elt, axes) end # Fixes ambiguity error. function Base.similar(a::AnyAbstractBlockSparseArray, elt::Type, axes::Tuple{}) # TODO: Use `@interface interface(a) similar(...)`. - return blocksparse_similar(a, elt, axes) + return @interface BlockSparseArrayInterface() similar(a, elt, axes) end # Fixes ambiguity error with `BlockArrays`. @@ -249,7 +270,7 @@ function Base.similar( }, ) # TODO: Use `@interface interface(a) similar(...)`. - return blocksparse_similar(a, elt, axes) + return @interface BlockSparseArrayInterface() similar(a, elt, axes) end # Fixes ambiguity error with `OffsetArrays`. @@ -259,7 +280,7 @@ function Base.similar( axes::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}, ) # TODO: Use `@interface interface(a) similar(...)`. - return blocksparse_similar(a, elt, axes) + return @interface BlockSparseArrayInterface() similar(a, elt, axes) end # Fixes ambiguity error with `BlockArrays`. @@ -269,7 +290,7 @@ function Base.similar( axes::Tuple{AbstractBlockedUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}, ) # TODO: Use `@interface interface(a) similar(...)`. - return blocksparse_similar(a, elt, axes) + return @interface BlockSparseArrayInterface() similar(a, elt, axes) end # Fixes ambiguity errors with BlockArrays. @@ -283,7 +304,7 @@ function Base.similar( }, ) # TODO: Use `@interface interface(a) similar(...)`. - return blocksparse_similar(a, elt, axes) + return @interface BlockSparseArrayInterface() similar(a, elt, axes) end # Fixes ambiguity error with `StaticArrays`. @@ -291,7 +312,7 @@ function Base.similar( a::AnyAbstractBlockSparseArray, elt::Type, axes::Tuple{Base.OneTo,Vararg{Base.OneTo}} ) # TODO: Use `@interface interface(a) similar(...)`. - return blocksparse_similar(a, elt, axes) + return @interface BlockSparseArrayInterface() similar(a, elt, axes) end # TODO: Implement this in a more generic way using a smarter `copyto!`, diff --git a/src/backup/qr.jl b/src/backup/qr.jl deleted file mode 100644 index c480398..0000000 --- a/src/backup/qr.jl +++ /dev/null @@ -1,143 +0,0 @@ -using .SparseArraysBase: SparseArrayDOK - -# Check if the matrix has 1 or fewer entries -# per row/column. -function is_permutation_matrix(a::SparseMatrixCSC) - return all(col -> length(nzrange(a, col)) ≤ 1, axes(a, 2)) -end - -# Check if the matrix has 1 or fewer entries -# per row/column. -function is_permutation_matrix(a::SparseArrayDOK{<:Any,2}) - keys = collect(Iterators.map(Tuple, nonzero_keys(a))) - I = first.(keys) - J = last.(keys) - return allunique(I) && allunique(J) -end - -function findnonzerorows(a::SparseMatrixCSC, col) - return view(a.rowval, a.colptr[col]:(a.colptr[col + 1] - 1)) -end - -# TODO: Is this already defined? -function SparseArrays.SparseMatrixCSC(a::SparseArrayDOK{<:Any,2}) - # Not defined: - # a_csc = SparseMatrixCSC{eltype(a)}(size(a)) - a_csc = spzeros(eltype(a), size(a)) - for I in nonzero_keys(a) - a_csc[I] = a[I] - end - return a_csc -end - -# TODO: Is this already defined? -# Get the sparse structure of a SparseArray as a SparseMatrixCSC. -function sparse_structure( - structure_type::Type{<:SparseMatrixCSC}, a::SparseArrayDOK{<:Any,2} -) - # Idealy would work but a bit too complicated for `map` right now: - # return SparseMatrixCSC(map(x -> iszero(x) ? false : true, a)) - # TODO: Change to `spzeros(Bool, size(a))`. - a_structure = structure_type(spzeros(Bool, size(a)...)) - for I in nonzero_keys(a) - i, j = Tuple(I) - a_structure[i, j] = true - end - return a_structure -end - -# Get the sparsity structure as a `SparseMatrixCSC` with values -# of `true` where there are structural nonzero blocks and `false` -# otherwise. -function block_sparse_structure(structure_type::Type, a::BlockSparseArray{<:Any,2}) - return sparse_structure(structure_type, blocks(a)) -end - -function is_block_permutation_matrix(a::BlockSparseArray{<:Any,2}) - return is_permutation_matrix(blocks(a)) -end - -qr_rank(alg::Algorithm"thin", a::AbstractArray{<:Any,2}) = minimum(size(a)) - -# m × n → (m × min(m, n)) ⋅ (min(m, n) × n) -function qr_block_sparse_structure(alg::Algorithm"thin", a::BlockSparseArray{<:Any,2}) - axes_row, axes_col = axes(a) - a_csc = block_sparse_structure(SparseMatrixCSC, a) - F = qr(float(a_csc)) - # Outputs full Q - # q_csc = sparse(F.Q[invperm(F.prow), :]) - q_csc = (F.Q * sparse(I, size(a_csc, 1), minimum(size(a_csc))))[invperm(F.prow), :] - r_csc = F.R[:, invperm(F.pcol)] - nblocks = size(q_csc, 2) - @assert nblocks == size(r_csc, 1) - a_sparse = blocks(a) - blocklengths_qr = Vector{Int}(undef, nblocks) - for I in nonzero_keys(a_sparse) - i, k = Tuple(I) - # Get the nonzero columns associated - # with the given row. - j = only(findnonzerorows(r_csc, k)) - # @assert is_structural_nonzero(r, j, k) - # @assert is_structural_nonzero(q, i, j) - blocklengths_qr[j] = qr_rank(alg, @view(a[BlockArrays.Block(i, k)])) - end - axes_qr = blockedrange(blocklengths_qr) - axes_q = (axes(a, 1), axes_qr) - axes_r = (axes_qr, axes(a, 2)) - # TODO: Come up with a better format to ouput. - # TODO: Get `axes_qr` as a permutation of the - # axes of `axes(a, 2)` to preserve sectors - # when using symmetric tensors. - return q_csc, axes_q, r_csc, axes_r -end - -# m × n → (m × m) ⋅ (m × n) -function qr_block_sparse_structure(alg::Algorithm"full", a::BlockSparseArray{<:Any,2}) - return error("Not implemented") -end - -function qr_blocks(a, structure_r, block_a) - i, k = block_a.n - j = only(findnonzerorows(structure_r, k)) - return BlockArrays.Block(i, j), BlockArrays.Block(j, k) -end - -# Block-preserving QR. -function LinearAlgebra.qr(a::BlockSparseArray{<:Any,2}; alg="thin") - return qr(Algorithm(alg), a) -end - -# Block-preserving QR. -function LinearAlgebra.qr(alg::Algorithm, a::BlockSparseArray{<:Any,2}) - if !is_block_permutation_matrix(a) - # Must have 1 or fewer blocks per row/column. - println("Block sparsity structure is:") - display(nonzero_blockkeys(a)) - error("Not a block permutation matrix") - end - eltype_a = eltype(a) - # TODO: `structure_q` isn't needed. - structure_q, axes_q, structure_r, axes_r = qr_block_sparse_structure(alg, a) - # TODO: Make this generic to GPU, use `similar`. - q = BlockSparseArray{eltype_a}(axes_q) - r = BlockSparseArray{eltype_a}(axes_r) - for block_a in nonzero_blockkeys(a) - # TODO: Make thin or full depending on `alg`. - q_b, r_b = qr(a[block_a]) - # Determine the block of Q and R - # TODO: Do the block locations change for `alg="full"`? - block_q, block_r = qr_blocks(a, structure_r, block_a) - - # TODO Make this generic to GPU. - q[block_q] = Matrix(q_b) - r[block_r] = r_b - end - # TODO: If `alg="full"`, fill in blocks of `q` - # with random unitaries. - # Which blocks should be filled? Seems to be based - # on the QNs... - # Maybe fill diagonal blocks. - # TODO: Also store `structure_r` in some way - # since that is needed for permuting the QNs. - return q, r -end diff --git a/src/blocksparsearray/blocksparsearray.jl b/src/blocksparsearray/blocksparsearray.jl index 3b3bbae..93b2bad 100644 --- a/src/blocksparsearray/blocksparsearray.jl +++ b/src/blocksparsearray/blocksparsearray.jl @@ -1,4 +1,5 @@ using BlockArrays: BlockArrays, Block, BlockedUnitRange, blockedrange, blocklength +using Derive: @interface using Dictionaries: Dictionary using SparseArraysBase: SparseArrayDOK @@ -169,7 +170,8 @@ Base.axes(a::BlockSparseArray) = a.axes # BlockArrays `AbstractBlockArray` interface. # This is used by `blocks(::AnyAbstractBlockSparseArray)`. -blocksparse_blocks(a::BlockSparseArray) = a.blocks +@interface ::AbstractBlockSparseArrayInterface BlockArrays.blocks(a::BlockSparseArray) = + a.blocks # TODO: Use `TypeParameterAccessors`. function blockstype( diff --git a/src/blocksparsearrayinterface/arraylayouts.jl b/src/blocksparsearrayinterface/arraylayouts.jl index 3ec48d5..a4a3cab 100644 --- a/src/blocksparsearrayinterface/arraylayouts.jl +++ b/src/blocksparsearrayinterface/arraylayouts.jl @@ -1,21 +1,16 @@ using ArrayLayouts: ArrayLayouts, Dot, MatMulMatAdd, MatMulVecAdd, MulAdd -using BlockArrays: BlockLayout +using BlockArrays: BlockArrays, BlockLayout, muladd! +using Derive: @interface using SparseArraysBase: SparseLayout -using LinearAlgebra: dot, mul! +using LinearAlgebra: LinearAlgebra, dot, mul! -function blocksparse_muladd!( +@interface ::AbstractBlockSparseArrayInterface function BlockArrays.muladd!( α::Number, a1::AbstractArray, a2::AbstractArray, β::Number, a_dest::AbstractArray ) mul!(blocks(a_dest), blocks(a1), blocks(a2), α, β) return a_dest end -function blocksparse_matmul!(m::MulAdd) - α, a1, a2, β, a_dest = m.α, m.A, m.B, m.β, m.C - blocksparse_muladd!(α, a1, a2, β, a_dest) - return a_dest -end - function ArrayLayouts.materialize!( m::MatMulMatAdd{ <:BlockLayout{<:SparseLayout}, @@ -23,7 +18,8 @@ function ArrayLayouts.materialize!( <:BlockLayout{<:SparseLayout}, }, ) - blocksparse_matmul!(m) + α, a1, a2, β, a_dest = m.α, m.A, m.B, m.β, m.C + @interface BlockSparseArrayInterface() muladd!(m.α, m.A, m.B, m.β, m.C) return m.C end function ArrayLayouts.materialize!( @@ -33,16 +29,18 @@ function ArrayLayouts.materialize!( <:BlockLayout{<:SparseLayout}, }, ) - blocksparse_matmul!(m) + @interface BlockSparseArrayInterface() matmul!(m) return m.C end -function blocksparse_dot(a1::AbstractArray, a2::AbstractArray) +@interface ::AbstractBlockSparseArrayInterface function LinearAlgebra.dot( + a1::AbstractArray, a2::AbstractArray +) # TODO: Add a check that the blocking of `a1` and `a2` are # the same, or the same up to a reshape. return dot(blocks(a1), blocks(a2)) end function Base.copy(d::Dot{<:BlockLayout{<:SparseLayout},<:BlockLayout{<:SparseLayout}}) - return blocksparse_dot(d.A, d.B) + return @interface BlockSparseArrayInterface() dot(d.A, d.B) end diff --git a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index f04d65d..a3e391e 100644 --- a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -1,5 +1,6 @@ using ArrayLayouts: ArrayLayouts, zero! using BlockArrays: + BlockArrays, AbstractBlockVector, Block, BlockIndex, @@ -39,17 +40,22 @@ Derive.arraytype(::AbstractBlockSparseArrayInterface, T::Type) = BlockSparseArra struct BlockSparseArrayInterface <: AbstractBlockSparseArrayInterface end -blocksparse_blocks(a::AbstractArray) = error("Not implemented") +@interface ::AbstractBlockSparseArrayInterface BlockArrays.blocks(a::AbstractArray) = + error("Not implemented") blockstype(a::AbstractArray) = blockstype(typeof(a)) -function blocksparse_getindex(a::AbstractArray{<:Any,N}, I::Vararg{Int,N}) where {N} +@interface ::AbstractBlockSparseArrayInterface function Base.getindex( + a::AbstractArray{<:Any,N}, I::Vararg{Int,N} +) where {N} @boundscheck checkbounds(a, I...) return a[findblockindex.(axes(a), I)...] end # Fix ambiguity error. -function blocksparse_getindex(a::AbstractArray{<:Any,0}) +@interface ::AbstractBlockSparseArrayInterface function Base.getindex( + a::AbstractArray{<:Any,0} +) # TODO: Use `Block()[]` once https://github.com/JuliaArrays/BlockArrays.jl/issues/430 # is fixed. return a[BlockIndex{0,Tuple{},Tuple{}}((), ())] @@ -61,14 +67,16 @@ end # and make that explicit with `@blocked a[1:2, 1:2]`. See the discussion in # https://github.com/JuliaArrays/BlockArrays.jl/issues/347 and also # https://github.com/ITensor/ITensors.jl/issues/1336. -function blocksparse_to_indices(a, inds, I::Tuple{UnitRange{<:Integer},Vararg{Any}}) +@interface ::AbstractBlockSparseArrayInterface function Base.to_indices( + a, inds, I::Tuple{UnitRange{<:Integer},Vararg{Any}} +) bs1 = to_blockindices(inds[1], I[1]) I1 = BlockSlice(bs1, blockedunitrange_getindices(inds[1], I[1])) return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) end # Special case when there is no blocking. -function blocksparse_to_indices( +@interface ::AbstractBlockSparseArrayInterface function Base.to_indices( a, inds::Tuple{Base.OneTo{<:Integer},Vararg{Any}}, I::Tuple{UnitRange{<:Integer},Vararg{Any}}, @@ -77,14 +85,16 @@ function blocksparse_to_indices( end # a[[Block(2), Block(1)], [Block(2), Block(1)]] -function blocksparse_to_indices(a, inds, I::Tuple{Vector{<:Block{1}},Vararg{Any}}) +@interface ::AbstractBlockSparseArrayInterface function Base.to_indices( + a, inds, I::Tuple{Vector{<:Block{1}},Vararg{Any}} +) I1 = BlockIndices(I[1], blockedunitrange_getindices(inds[1], I[1])) return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) end # a[mortar([Block(1)[1:2], Block(2)[1:3]]), mortar([Block(1)[1:2], Block(2)[1:3]])] # a[[Block(1)[1:2], Block(2)[1:3]], [Block(1)[1:2], Block(2)[1:3]]] -function blocksparse_to_indices( +@interface ::AbstractBlockSparseArrayInterface function Base.to_indices( a, inds, I::Tuple{BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}},Vararg{Any}} ) I1 = BlockIndices(I[1], blockedunitrange_getindices(inds[1], I[1])) @@ -94,7 +104,7 @@ end # a[BlockVector([Block(2), Block(1)], [2]), BlockVector([Block(2), Block(1)], [2])] # Permute and merge blocks. # TODO: This isn't merging blocks yet, that needs to be implemented that. -function blocksparse_to_indices( +@interface ::AbstractBlockSparseArrayInterface function Base.to_indices( a, inds, I::Tuple{AbstractBlockVector{<:Block{1}},Vararg{Any}} ) I1 = BlockIndices(I[1], blockedunitrange_getindices(inds[1], I[1])) @@ -104,21 +114,27 @@ end # TODO: Need to implement this! function block_merge end -function blocksparse_setindex!(a::AbstractArray{<:Any,N}, value, I::Vararg{Int,N}) where {N} +@interface ::AbstractBlockSparseArrayInterface function Base.setindex!( + a::AbstractArray{<:Any,N}, value, I::Vararg{Int,N} +) where {N} @boundscheck checkbounds(a, I...) a[findblockindex.(axes(a), I)...] = value return a end # Fix ambiguity error. -function blocksparse_setindex!(a::AbstractArray{<:Any,0}, value) +@interface ::AbstractBlockSparseArrayInterface function Base.setindex!( + a::AbstractArray{<:Any,0}, value +) # TODO: Use `Block()[]` once https://github.com/JuliaArrays/BlockArrays.jl/issues/430 # is fixed. a[BlockIndex{0,Tuple{},Tuple{}}((), ())] = value return a end -function blocksparse_setindex!(a::AbstractArray{<:Any,N}, value, I::BlockIndex{N}) where {N} +@interface ::AbstractBlockSparseArrayInterface function Base.setindex!( + a::AbstractArray{<:Any,N}, value, I::BlockIndex{N} +) where {N} i = Int.(Tuple(block(I))) a_b = blocks(a)[i...] a_b[I.α...] = value @@ -128,7 +144,9 @@ function blocksparse_setindex!(a::AbstractArray{<:Any,N}, value, I::BlockIndex{N end # Fix ambiguity error. -function blocksparse_setindex!(a::AbstractArray{<:Any,0}, value, I::BlockIndex{0}) +@interface ::AbstractBlockSparseArrayInterface function Base.setindex!( + a::AbstractArray{<:Any,0}, value, I::BlockIndex{0} +) a_b = blocks(a)[] a_b[] = value # Set the block, required if it is structurally zero. @@ -177,7 +195,9 @@ struct SparsePermutedDimsArrayBlocks{ } <: AbstractSparseArray{BlockType,N} array::Array end -function blocksparse_blocks(a::PermutedDimsArray) +@interface ::AbstractBlockSparseArrayInterface function BlockArrays.blocks( + a::PermutedDimsArray +) return SparsePermutedDimsArrayBlocks{eltype(a),ndims(a),blocktype(parent(a)),typeof(a)}(a) end function Base.size(a::SparsePermutedDimsArrayBlocks) @@ -207,8 +227,10 @@ end reverse_index(index) = reverse(index) reverse_index(index::CartesianIndex) = CartesianIndex(reverse(Tuple(index))) -blocksparse_blocks(a::Transpose) = transpose(blocks(parent(a))) -blocksparse_blocks(a::Adjoint) = adjoint(blocks(parent(a))) +@interface ::AbstractBlockSparseArrayInterface BlockArrays.blocks(a::Transpose) = + transpose(blocks(parent(a))) +@interface ::AbstractBlockSparseArrayInterface BlockArrays.blocks(a::Adjoint) = + adjoint(blocks(parent(a))) # Represents the array of arrays of a `SubArray` # wrapping a block spare array, i.e. `blocks(array)` where `a` is a `SubArray`. @@ -216,7 +238,7 @@ struct SparseSubArrayBlocks{T,N,BlockType<:AbstractArray{T,N},Array<:SubArray{T, AbstractSparseArray{BlockType,N} array::Array end -function blocksparse_blocks(a::SubArray) +@interface ::AbstractBlockSparseArrayInterface function BlockArrays.blocks(a::SubArray) return SparseSubArrayBlocks{eltype(a),ndims(a),blocktype(parent(a)),typeof(a)}(a) end # TODO: Define this as `blockrange(a::AbstractArray, indices::Tuple{Vararg{AbstractUnitRange}})`. @@ -295,7 +317,7 @@ end to_blocks_indices(I::BlockSlice{<:BlockRange{1}}) = Int.(I.block) to_blocks_indices(I::BlockIndices{<:Vector{<:Block{1}}}) = Int.(I.blocks) -function blocksparse_blocks( +@interface ::AbstractBlockSparseArrayInterface function BlockArrays.blocks( a::SubArray{<:Any,<:Any,<:Any,<:Tuple{Vararg{BlockSliceCollection}}} ) return @view blocks(parent(a))[map(to_blocks_indices, parentindices(a))...] diff --git a/src/blocksparsearrayinterface/linearalgebra.jl b/src/blocksparsearrayinterface/linearalgebra.jl index ac7f566..702e761 100644 --- a/src/blocksparsearrayinterface/linearalgebra.jl +++ b/src/blocksparsearrayinterface/linearalgebra.jl @@ -1,6 +1,6 @@ -using LinearAlgebra: mul! +using LinearAlgebra: LinearAlgebra, mul! -function blocksparse_mul!( +@interface ::AbstractBlockSparseArrayInterface function LinearAlgebra.mul!( a_dest::AbstractMatrix, a1::AbstractMatrix, a2::AbstractMatrix, diff --git a/src/blocksparsearrayinterface/views.jl b/src/blocksparsearrayinterface/views.jl index 8e43f26..56d17b0 100644 --- a/src/blocksparsearrayinterface/views.jl +++ b/src/blocksparsearrayinterface/views.jl @@ -1,3 +1,3 @@ -function blocksparse_view(a, I...) +@interface ::AbstractBlockSparseArrayInterface function Base.view(a, I...) return Base.invoke(view, Tuple{AbstractArray,Vararg{Any}}, a, I...) end