Skip to content

Commit

Permalink
[BlockSparseArrys] Zero dimensional block sparse array
Browse files Browse the repository at this point in the history
  • Loading branch information
mtfishman committed Nov 8, 2024
1 parent f9b6309 commit 1701179
Show file tree
Hide file tree
Showing 5 changed files with 62 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -22,11 +22,16 @@ function blockstype(arraytype::Type{<:AbstractBlockSparseArray{T,N}}) where {T,N
return SparseArrayDOK{AbstractArray{T,N},N}
end

## # Specialized in order to fix ambiguity error with `BlockArrays`.
# 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...)
end

# Specialized in order to fix ambiguity error with `BlockArrays`.
function Base.getindex(a::AbstractBlockSparseArray{<:Any,0})
return blocksparse_getindex(a)
end

## # Fix ambiguity error with `BlockArrays`.
## function Base.getindex(a::AbstractBlockSparseArray{<:Any,N}, I::Block{N}) where {N}
## return ArrayLayouts.layout_getindex(a, I)
Expand All @@ -51,6 +56,12 @@ function Base.setindex!(
return a
end

# Fix ambiguity error.
function Base.setindex!(a::AbstractBlockSparseArray{<:Any,0}, value)
blocksparse_setindex!(a, value)
return a
end

function Base.setindex!(
a::AbstractBlockSparseArray{<:Any,N}, value, I::Vararg{Block{1},N}
) where {N}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,10 @@ function Base.getindex(
)
return ArrayLayouts.layout_getindex(a, I...)
end
# Fixes ambiguity error.
function Base.getindex(a::BlockSparseArrayLike{<:Any,0})
return ArrayLayouts.layout_getindex(a)
end

# TODO: Define `blocksparse_isassigned`.
function Base.isassigned(
Expand All @@ -100,6 +104,11 @@ function Base.isassigned(
return isassigned(blocks(a), Int.(index)...)
end

# Fix ambiguity error.
function Base.isassigned(a::BlockSparseArrayLike{<:Any,0})
return isassigned(blocks(a))
end

function Base.isassigned(a::BlockSparseArrayLike{<:Any,N}, index::Block{N}) where {N}
return isassigned(a, Tuple(index)...)
end
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,10 @@ function BlockSparseArray{T,N}(axes::Tuple{Vararg{AbstractUnitRange,N}}) where {
return BlockSparseArray{T,N,default_arraytype(T, axes)}(axes)
end

function BlockSparseArray{T,0}(axes::Tuple{}) where {T}
return BlockSparseArray{T,0,default_arraytype(T, axes)}(axes)
end

function BlockSparseArray{T,N}(dims::Tuple{Vararg{Vector{Int},N}}) where {T,N}
return BlockSparseArray{T,N}(blockedrange.(dims))
end
Expand All @@ -84,6 +88,10 @@ function BlockSparseArray{T}(axes::Tuple{Vararg{AbstractUnitRange}}) where {T}
return BlockSparseArray{T,length(axes)}(axes)
end

function BlockSparseArray{T}(axes::Tuple{}) where {T}
return BlockSparseArray{T,length(axes)}(axes)
end

function BlockSparseArray{T}(dims::Vararg{Vector{Int}}) where {T}
return BlockSparseArray{T}(dims)
end
Expand All @@ -92,6 +100,10 @@ function BlockSparseArray{T}(axes::Vararg{AbstractUnitRange}) where {T}
return BlockSparseArray{T}(axes)
end

function BlockSparseArray{T}() where {T}
return BlockSparseArray{T}(())
end

function BlockSparseArray{T,N,A}(
::UndefInitializer, dims::Tuple
) where {T,N,A<:AbstractArray{T,N}}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,13 @@ function blocksparse_getindex(a::AbstractArray{<:Any,N}, I::Vararg{Int,N}) where
return a[findblockindex.(axes(a), I)...]
end

# Fix ambiguity error.
function blocksparse_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{}}((), ())]
end

# a[1:2, 1:2]
# TODO: This definition means that the result of slicing a block sparse array
# with a non-blocked unit range is blocked. We may want to change that behavior,
Expand Down Expand Up @@ -77,6 +84,14 @@ function blocksparse_setindex!(a::AbstractArray{<:Any,N}, value, I::Vararg{Int,N
return a
end

# Fix ambiguity error.
function blocksparse_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}
i = Int.(Tuple(block(I)))
a_b = blocks(a)[i...]
Expand All @@ -86,6 +101,15 @@ function blocksparse_setindex!(a::AbstractArray{<:Any,N}, value, I::BlockIndex{N
return a
end

# Fix ambiguity error.
function blocksparse_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.
blocks(a)[] = a_b
return a
end

function blocksparse_fill!(a::AbstractArray, value)
for b in BlockRange(a)
# We can't use:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,11 @@ function sparse_getindex(a::AbstractArray, I::Vararg{Int})
return sparse_getindex(a, CartesianIndex(I))
end

# Fix ambiguity error.
function sparse_getindex(a::AbstractArray{<:Any,0})
return sparse_getindex(a, CartesianIndex())
end

# Linear indexing
function sparse_getindex(a::AbstractArray, I::CartesianIndex{1})
return sparse_getindex(a, CartesianIndices(a)[I])
Expand Down

0 comments on commit 1701179

Please sign in to comment.