diff --git a/Project.toml b/Project.toml index 9e61a2e..b583b07 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.5.1" +version = "0.5.2" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/BlockArraysExtensions/BlockArraysExtensions.jl b/src/BlockArraysExtensions/BlockArraysExtensions.jl index 46d015a..3853304 100644 --- a/src/BlockArraysExtensions/BlockArraysExtensions.jl +++ b/src/BlockArraysExtensions/BlockArraysExtensions.jl @@ -30,13 +30,14 @@ using SparseArraysBase: # A return type for `blocks(array)` when `array` isn't blocked. # Represents a vector with just that single block. -struct SingleBlockView{T,N,Array<:AbstractArray{T,N}} <: AbstractArray{T,N} +struct SingleBlockView{N,Array<:AbstractArray{<:Any,N}} <: AbstractArray{Array,N} array::Array end Base.parent(a::SingleBlockView) = a.array +Base.size(a::SingleBlockView) = ntuple(Returns(1), ndims(a)) blocks_maybe_single(a) = blocks(a) blocks_maybe_single(a::Array) = SingleBlockView(a) -function Base.getindex(a::SingleBlockView{<:Any,N}, index::Vararg{Int,N}) where {N} +function Base.getindex(a::SingleBlockView{N}, index::Vararg{Int,N}) where {N} @assert all(isone, index) return parent(a) end @@ -357,7 +358,11 @@ function blockrange(axis::AbstractUnitRange, r::Base.Slice) end function blockrange(axis::AbstractUnitRange, r::NonBlockedVector) - return Block(1):Block(1) + return Block.(Base.OneTo(1)) +end + +function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:Integer}) + return Block.(Base.OneTo(1)) end function blockrange(axis::AbstractUnitRange, r) diff --git a/src/BlockSparseArrays.jl b/src/BlockSparseArrays.jl index ffbfe93..893bcf0 100644 --- a/src/BlockSparseArrays.jl +++ b/src/BlockSparseArrays.jl @@ -27,6 +27,7 @@ include("abstractblocksparsearray/abstractblocksparsearray.jl") include("abstractblocksparsearray/abstractblocksparsematrix.jl") include("abstractblocksparsearray/abstractblocksparsevector.jl") include("abstractblocksparsearray/wrappedabstractblocksparsearray.jl") +include("abstractblocksparsearray/unblockedsubarray.jl") include("abstractblocksparsearray/views.jl") include("abstractblocksparsearray/arraylayouts.jl") include("abstractblocksparsearray/sparsearrayinterface.jl") diff --git a/src/abstractblocksparsearray/arraylayouts.jl b/src/abstractblocksparsearray/arraylayouts.jl index 8b77120..b875c5d 100644 --- a/src/abstractblocksparsearray/arraylayouts.jl +++ b/src/abstractblocksparsearray/arraylayouts.jl @@ -43,11 +43,20 @@ function ArrayLayouts.sub_materialize(layout::BlockLayout{<:SparseLayout}, a, ax return a_dest end +function _similar(arraytype::Type{<:AbstractArray}, size::Tuple) + return similar(arraytype, size) +end +function _similar( + ::Type{<:SubArray{<:Any,<:Any,<:ArrayType}}, size::Tuple +) where {ArrayType} + return similar(ArrayType, size) +end + # Materialize a SubArray view. function ArrayLayouts.sub_materialize( layout::BlockLayout{<:SparseLayout}, a, axes::Tuple{Vararg{Base.OneTo}} ) - a_dest = blocktype(a)(undef, length.(axes)) + a_dest = _similar(blocktype(a), length.(axes)) a_dest .= a return a_dest end diff --git a/src/abstractblocksparsearray/unblockedsubarray.jl b/src/abstractblocksparsearray/unblockedsubarray.jl new file mode 100644 index 0000000..355ae82 --- /dev/null +++ b/src/abstractblocksparsearray/unblockedsubarray.jl @@ -0,0 +1,86 @@ +using ArrayLayouts: ArrayLayouts, MemoryLayout +using Base.Broadcast: Broadcast, BroadcastStyle +using BlockArrays: BlockArrays, Block, BlockIndexRange, BlockSlice +using TypeParameterAccessors: TypeParameterAccessors, parenttype, similartype + +const UnblockedIndices = Union{ + Vector{<:Integer},BlockSlice{<:Block{1}},BlockSlice{<:BlockIndexRange{1}} +} + +const UnblockedSubArray{T,N} = SubArray{ + T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{UnblockedIndices}} +} + +function BlockArrays.blocks(a::UnblockedSubArray) + return SingleBlockView(a) +end + +function DerivableInterfaces.interface(arraytype::Type{<:UnblockedSubArray}) + return interface(blocktype(parenttype(arraytype))) +end + +function ArrayLayouts.MemoryLayout(arraytype::Type{<:UnblockedSubArray}) + return MemoryLayout(blocktype(parenttype(arraytype))) +end + +function Broadcast.BroadcastStyle(arraytype::Type{<:UnblockedSubArray}) + return BroadcastStyle(blocktype(parenttype(arraytype))) +end + +function TypeParameterAccessors.similartype(arraytype::Type{<:UnblockedSubArray}, elt::Type) + return similartype(blocktype(parenttype(arraytype)), elt) +end + +function Base.similar( + a::UnblockedSubArray, elt::Type, axes::Tuple{Base.OneTo,Vararg{Base.OneTo}} +) + return similar(similartype(blocktype(parenttype(a)), elt), axes) +end +function Base.similar(a::UnblockedSubArray, elt::Type, size::Tuple{Int,Vararg{Int}}) + return similar(a, elt, Base.OneTo.(size)) +end + +function ArrayLayouts.sub_materialize(a::UnblockedSubArray) + a_cpu = adapt(Array, a) + a_cpu′ = similar(a_cpu) + a_cpu′ .= a_cpu + if typeof(a) === typeof(a_cpu) + return a_cpu′ + end + a′ = similar(a) + a′ .= a_cpu′ + return a′ +end + +function Base.map!( + f, a_dest::AbstractArray, a_src1::UnblockedSubArray, a_src_rest::UnblockedSubArray... +) + return invoke( + map!, + Tuple{Any,AbstractArray,AbstractArray,Vararg{AbstractArray}}, + f, + a_dest, + a_src1, + a_src_rest..., + ) +end + +# Fix ambiguity and scalar indexing errors with GPUArrays. +using Adapt: adapt +using GPUArraysCore: GPUArraysCore +function Base.map!( + f, + a_dest::GPUArraysCore.AnyGPUArray, + a_src1::UnblockedSubArray, + a_src_rest::UnblockedSubArray..., +) + a_dest_cpu = adapt(Array, a_dest) + a_srcs_cpu = map(adapt(Array), (a_src1, a_src_rest...)) + map!(f, a_dest_cpu, a_srcs_cpu...) + a_dest .= a_dest_cpu + return a_dest +end + +function Base.iszero(a::UnblockedSubArray) + return invoke(iszero, Tuple{AbstractArray}, adapt(Array, a)) +end diff --git a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index 1b1af5f..a76edbf 100644 --- a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -364,7 +364,19 @@ end function Base.size(a::SparseSubArrayBlocks) return length.(axes(a)) end -# TODO: Define `isstored`. + +# TODO: Make a faster version for when the slice is blockwise. +function SparseArraysBase.isstored( + a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N} +) where {N} + J = Base.reindex(parentindices(a.array), to_indices(a.array, Block.(I))) + # TODO: Try doing this blockwise when possible rather + # than elementwise. + return any(Iterators.product(J...)) do K + return isstored(parent(a.array), K...) + end +end + # TODO: Define `getstoredindex`, `getunstoredindex` instead. function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) where {N} # TODO: Should this be defined as `@view a.array[Block(I)]` instead? @@ -400,9 +412,17 @@ function Base.isassigned(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) whe # TODO: Implement this properly. return true end -function SparseArraysBase.eachstoredindex(a::SparseSubArrayBlocks) - return eachstoredindex(view(blocks(parent(a.array)), blockrange(a)...)) + +function SparseArraysBase.eachstoredindex(::IndexCartesian, a::SparseSubArrayBlocks) + return filter(eachindex(a)) do I + return isstored(a, I) + end + + ## # TODO: This only works for blockwise slices, i.e. slices using + ## # `BlockSliceCollection`. + ## return eachstoredindex(view(blocks(parent(a.array)), blockrange(a)...)) end + # TODO: Either make this the generic interface or define # `SparseArraysBase.sparse_storage`, which is used # to defined this. diff --git a/test/test_basics.jl b/test/test_basics.jl index 0a59f14..bc220d2 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -51,13 +51,6 @@ arrayts = (Array, JLArray) a[Block(2, 2)] = dev(randn(elt, 3, 3)) @test_broken a[:, 4] - # TODO: Fix this and turn it into a proper test. - a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) - a[Block(1, 1)] = dev(randn(elt, 2, 2)) - a[Block(2, 2)] = dev(randn(elt, 3, 3)) - @test_broken a[:, [2, 4]] - @test_broken a[[3, 5], [2, 4]] - # TODO: Fix this and turn it into a proper test. a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) a[Block(1, 1)] = dev(randn(elt, 2, 2)) @@ -713,6 +706,20 @@ arrayts = (Array, JLArray) @test a[Block(2, 2)[1:2, 2:3]] == b @test blockstoredlength(a) == 1 + # Noncontiguous slicing. + a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) + a[Block(1, 1)] = dev(randn(elt, 2, 2)) + a[Block(2, 2)] = dev(randn(elt, 3, 3)) + I = ([3, 5], [2, 4]) + @test Array(a[I...]) == Array(a)[I...] + + # Noncontiguous slicing. + a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) + a[Block(1, 1)] = dev(randn(elt, 2, 2)) + a[Block(2, 2)] = dev(randn(elt, 3, 3)) + I = (:, [2, 4]) + @test Array(a[I...]) == Array(a)[I...] + a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) @views for b in [Block(1, 1), Block(2, 2)] a[b] = randn(elt, size(a[b]))