Skip to content

Commit

Permalink
Rename BlockZero to GetUnstoredBlock
Browse files Browse the repository at this point in the history
  • Loading branch information
mtfishman committed Dec 16, 2024
1 parent abe5237 commit 0812e79
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 53 deletions.
2 changes: 1 addition & 1 deletion src/blocksparsearray/defaults.jl
Original file line number Diff line number Diff line change
Expand Up @@ -38,5 +38,5 @@ function default_blocks(
block_data::Dictionary{<:CartesianIndex{N},<:AbstractArray{<:Any,N}},
axes::Tuple{Vararg{AbstractUnitRange,N}},
) where {N}
return SparseArrayDOK(block_data, blocklength.(axes), BlockZero(axes))
return SparseArrayDOK(block_data, blocklength.(axes), GetUnstoredBlock(axes))
end
17 changes: 7 additions & 10 deletions src/blocksparsearrayinterface/blocksparsearrayinterface.jl
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ end
)
# TODO: Use `Block()[]` once https://github.com/JuliaArrays/BlockArrays.jl/issues/430
# is fixed.
return a[BlockIndex{0,Tuple{},Tuple{}}((), ())]
return a[BlockIndex()]
end

# a[1:2, 1:2]
Expand Down Expand Up @@ -135,7 +135,7 @@ end
)
# TODO: Use `Block()[]` once https://github.com/JuliaArrays/BlockArrays.jl/issues/430
# is fixed.
a[BlockIndex{0,Tuple{},Tuple{}}((), ())] = value
a[BlockIndex()] = value
return a
end

Expand Down Expand Up @@ -301,6 +301,8 @@ end
function Base.size(a::SparseSubArrayBlocks)
return length.(axes(a))
end
# TODO: Define `isstored`.
# 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?
return @view a.array[Block(I)]
Expand All @@ -312,9 +314,11 @@ function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) where
## return @view parent_block[blockindices(parent(a.array), block, a.array.indices)...]
end
# TODO: This should be handled by generic `AbstractSparseArray` code.
# TODO: Define `getstoredindex`, `getunstoredindex` instead.
function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::CartesianIndex{N}) where {N}
return a[Tuple(I)...]
end
# TODO: Define `setstoredindex!`, `setunstoredindex!` instead.
function Base.setindex!(a::SparseSubArrayBlocks{<:Any,N}, value, I::Vararg{Int,N}) where {N}
parent_blocks = @view blocks(parent(a.array))[blockrange(a)...]
# TODO: The following line is required to instantiate
Expand Down Expand Up @@ -345,18 +349,11 @@ SparseArraysBase.storedlength(a::SparseSubArrayBlocks) = length(eachstoredindex(
## array::Array
## end

## TODO: Delete.
## TODO: Define `storedvalues` instead.
## function SparseArraysBase.sparse_storage(a::SparseSubArrayBlocks)
## return map(I -> a[I], eachstoredindex(a))
## end

## TODO: Delete.
## function SparseArraysBase.getindex_zero_function(a::SparseSubArrayBlocks)
## # TODO: Base it off of `getindex_zero_function(blocks(parent(a.array))`, but replace the
## # axes with `axes(a.array)`.
## return BlockZero(axes(a.array))
## end

function SparseArraysBase.getunstoredindex(
a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}
) where {N}
Expand Down
60 changes: 18 additions & 42 deletions src/blocksparsearrayinterface/blockzero.jl
Original file line number Diff line number Diff line change
@@ -1,49 +1,25 @@
using BlockArrays: Block, blockedrange
using ArrayLayouts: zero!
using BlockArrays: Block

# Extensions to BlockArrays.jl
blocktuple(b::Block) = Block.(b.n)
inttuple(b::Block) = b.n

# The size of a block
function block_size(axes::Tuple{Vararg{AbstractUnitRange}}, block::Block)
return length.(getindex.(axes, blocktuple(block)))
end

# The size of a block
function block_size(blockinds::Tuple{Vararg{AbstractVector}}, block::Block)
return block_size(blockedrange.(blockinds), block)
end

struct BlockZero{Axes}
struct GetUnstoredBlock{Axes}
axes::Axes
end

function (f::BlockZero)(a::AbstractArray, I...)
return f(eltype(a), I...)
end

function (f::BlockZero)(arraytype::Type{<:SubArray{<:Any,<:Any,P}}, I...) where {P}
return f(P, I...)
end

function (f::BlockZero)(arraytype::Type{<:AbstractArray}, I::CartesianIndex)
return f(arraytype, Tuple(I)...)
end

function (f::BlockZero)(arraytype::Type{<:AbstractArray}, I::Int...)
@inline function (f::GetUnstoredBlock)(
a::AbstractArray{<:Any,N}, I::Vararg{Int,N}
) where {N}
# TODO: Make sure this works for sparse or block sparse blocks, immutable
# blocks, diagonal blocks, etc.!
blck_size = block_size(f.axes, Block(I))
blck_type = similartype(arraytype, blck_size)
return fill!(blck_type(undef, blck_size), false)
end

# Fallback so that `SparseArray` with scalar elements works.
function (f::BlockZero)(blocktype::Type{<:Number}, I...)
return zero(blocktype)
end

# Fallback to Array if it is abstract
function (f::BlockZero)(arraytype::Type{AbstractArray{T,N}}, I...) where {T,N}
return f(Array{T,N}, I...)
b_size = ntuple(ndims(a)) do d
return length(f.axes[d][Block(I[d])])
end
b = similar(eltype(a), b_size)
zero!(b)
return b
end
# TODO: Use `Base.to_indices`.
@inline function (f::GetUnstoredBlock)(
a::AbstractArray{<:Any,N}, I::CartesianIndex{N}
) where {N}
return f(a, Tuple(I)...)
end

0 comments on commit 0812e79

Please sign in to comment.