Skip to content

[RFC] Non-contiguous block slicing operations #462

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 17 commits into
base: master
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/src/lib/internals.md
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,9 @@ BlockedOneTo
BlockedUnitRange
BlockRange
BlockIndexRange
BlockIndices
BlockSlice
NoncontiguousBlockSlice
unblock
SubBlockIterator
blockcheckbounds_indices
Expand Down
2 changes: 1 addition & 1 deletion src/BlockArrays.jl
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ using LinearAlgebra, ArrayLayouts, FillArrays
# AbstractBlockArray interface exports
export AbstractBlockArray, AbstractBlockMatrix, AbstractBlockVector, AbstractBlockVecOrMat
export Block, getblock, getblock!, setblock!, eachblock, blocks
export blockaxes, blocksize, blocklength, blockcheckbounds, BlockBoundsError, BlockIndex, BlockIndexRange
export blockaxes, blocksize, blocklength, blockcheckbounds, BlockBoundsError, BlockIndex, BlockIndexRange, BlockIndices
export blocksizes, blocklengths, blocklasts, blockfirsts, blockisequal, blockequals, blockisapprox
export eachblockaxes
export BlockRange, blockedrange, BlockedUnitRange, BlockedOneTo
Expand Down
14 changes: 7 additions & 7 deletions src/blockaxis.jl
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,9 @@
@propagate_inbounds getindex(b::AbstractArray, K::BlockIndex{1}, J::BlockIndex{1}...) =
b[BlockIndex(tuple(K, J...))]

@propagate_inbounds getindex(b::AbstractArray{T,N}, K::BlockIndexRange{N}) where {T,N} = b[block(K)][K.indices...]
@propagate_inbounds getindex(b::LayoutArray{T,N}, K::BlockIndexRange{N}) where {T,N} = b[block(K)][K.indices...]
@propagate_inbounds getindex(b::LayoutArray{T,1}, K::BlockIndexRange{1}) where {T} = b[block(K)][K.indices...]
@propagate_inbounds getindex(b::AbstractArray{T,N}, K::BlockIndices{N}) where {T,N} = b[block(K)][K.indices...]
@propagate_inbounds getindex(b::LayoutArray{T,N}, K::BlockIndices{N}) where {T,N} = b[block(K)][K.indices...]
@propagate_inbounds getindex(b::LayoutArray{T,1}, K::BlockIndices{1}) where {T} = b[block(K)][K.indices...]

function findblockindex(b::AbstractVector, k::Integer)
@boundscheck k in b || throw(BoundsError())
Expand Down Expand Up @@ -539,11 +539,11 @@ end

getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:Block{1}}) = mortar([b[K] for K in KR])
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:AbstractVector{<:Block{1}}}) = mortar([b[K] for K in KR])
getindex(b::AbstractBlockedUnitRange, Kkr::BlockIndexRange{1}) = b[block(Kkr)][Kkr.indices...]
getindex(b::AbstractBlockedUnitRange, Kkr::BlockIndices{1}) = b[block(Kkr)][Kkr.indices...]
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:BlockIndex{1}}) = [b[K] for K in KR]
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:BlockIndexRange{1}}) = mortar([b[K] for K in KR])
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:BlockIndices{1}}) = mortar([b[K] for K in KR])
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:AbstractVector{<:BlockIndex{1}}}) = mortar([b[K] for K in KR])
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:AbstractVector{<:BlockIndexRange{1}}}) = mortar([b[K] for K in KR])
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:AbstractVector{<:BlockIndices{1}}}) = mortar([b[K] for K in KR])
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:AbstractVector{<:AbstractVector{<:BlockIndex{1}}}}) = mortar([b[K] for K in KR])

_searchsortedfirst(a::AbstractVector, k) = searchsortedfirst(a, k)
Expand All @@ -567,7 +567,7 @@ Base.dataids(b::AbstractBlockedUnitRange) = Base.dataids(blocklasts(b))
Base.checkindex(::Type{Bool}, b::BlockRange, K::Integer) = checkindex(Bool, Integer.(b), K)
Base.checkindex(::Type{Bool}, b::AbstractUnitRange{<:Integer}, K::Block{1}) = checkindex(Bool, blockaxes(b,1), Integer(K))

function Base.checkindex(::Type{Bool}, axis::AbstractBlockedUnitRange, ind::BlockIndexRange{1})
function Base.checkindex(::Type{Bool}, axis::AbstractBlockedUnitRange, ind::BlockIndices{1})
checkindex(Bool, axis, first(ind)) && checkindex(Bool, axis, last(ind))
end
function Base.checkindex(::Type{Bool}, axis::AbstractBlockedUnitRange, ind::BlockIndex{1})
Expand Down
154 changes: 109 additions & 45 deletions src/blockindices.jl
Original file line number Diff line number Diff line change
Expand Up @@ -163,23 +163,23 @@
20
```
"""
struct BlockIndex{N,TI<:Tuple{Vararg{Integer,N}},Tα<:Tuple{Vararg{Integer,N}}}
struct BlockIndex{N,TI<:Tuple{Vararg{Any,N}},Tα<:Tuple{Vararg{Any,N}}}
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This generalization enables constructing indexing objects such as Block(1)[Block(2)] or Block(1)[Block(2)[3]], for example in the case when a block array has blocks that are themselves blocked. I have use cases for that where I want to construct a block array where the blocks have extra structure and custom non-integer indexing (in particular, the blocks have a Kronecker product structure, so I want to be able to slice the blocks while preserving that structure).

I::TI
α::Tα
end

@inline BlockIndex(a::NTuple{N,Block{1}}, b::Tuple) where N = BlockIndex(Int.(a), b)
@inline BlockIndex(::Tuple{}, b::Tuple{}) = BlockIndex{0,Tuple{},Tuple{}}((), ())

@inline BlockIndex(a::Integer, b::Integer) = BlockIndex((a,), (b,))
@inline BlockIndex(a::Tuple, b::Integer) = BlockIndex(a, (b,))
@inline BlockIndex(a::Integer, b::Tuple) = BlockIndex((a,), b)
@inline BlockIndex(a, b) = BlockIndex((a,), (b,))
@inline BlockIndex(a::Tuple, b) = BlockIndex(a, (b,))
@inline BlockIndex(a, b::Tuple) = BlockIndex((a,), b)
@inline BlockIndex() = BlockIndex((), ())

@inline BlockIndex(a::Block, b::Tuple) = BlockIndex(a.n, b)
@inline BlockIndex(a::Block, b::Integer) = BlockIndex(a, (b,))
@inline BlockIndex(a::Block, b) = BlockIndex(a, (b,))

@inline function BlockIndex(I::Tuple{Vararg{Integer,N}}, α::Tuple{Vararg{Integer,M}}) where {M,N}
@inline function BlockIndex(I::Tuple{Vararg{Any,N}}, α::Tuple{Vararg{Any,M}}) where {M,N}
M <= N || throw(ArgumentError("number of indices must not exceed the number of blocks"))
α2 = ntuple(k -> k <= M ? α[k] : 1, N)
BlockIndex(I, α2)
Expand Down Expand Up @@ -207,19 +207,71 @@
checkbounds(::Type{Bool}, A::AbstractArray{<:Any,N}, I::AbstractArray{<:BlockIndex{N}}) where N =
all(i -> checkbounds(Bool, A, i), I)

struct BlockIndexRange{N,R<:Tuple{Vararg{AbstractUnitRange{<:Integer},N}},I<:Tuple{Vararg{Integer,N}},BI<:Integer} <: AbstractArray{BlockIndex{N,NTuple{N,BI},I},N}
struct BlockIndices{N,R<:Tuple{Vararg{AbstractVector,N}},I<:Tuple{Vararg{Any,N}},BI} <: AbstractArray{BlockIndex{N,NTuple{N,BI},I},N}
block::Block{N,BI}
indices::R
function BlockIndexRange(block::Block{N,BI}, inds::R) where {N,BI<:Integer,R<:Tuple{Vararg{AbstractUnitRange{<:Integer},N}}}
function BlockIndices(block::Block{N,BI}, inds::R) where {N,BI<:Integer,R<:Tuple{Vararg{AbstractVector,N}}}
I = Tuple{eltype.(inds)...}
return new{N,R,I,BI}(block,inds)
end
end

"""
BlockIndices(block, startind:stopind)

Represents a cartesian product of indices inside a block.

It can be constructed and used to index into `BlockArrays` in the following manner:

```jldoctest
julia> BlockIndices(Block(1,2), ([1,3],[2,4]))
Block(1, 2)[[1, 3], [2, 4]]

julia> Block(1)[[1,3]] == BlockIndices(Block(1), [1,3])
true

julia> Block(1,2)[[1,3],[2,4]] == BlockIndices(Block(1,2), ([1,3],[2,4]))
true

julia> BlockIndices((Block(1)[[1,3]], Block(2)[[2,4]]))
Block(1, 2)[[1, 3], [2, 4]]

julia> arr = Array(reshape(1:25, (5,5)));

julia> a = BlockedArray(arr, [3,2], [1,4])
2×2-blocked 5×5 BlockedMatrix{Int64}:
1 │ 6 11 16 21
2 │ 7 12 17 22
3 │ 8 13 18 23
───┼────────────────
4 │ 9 14 19 24
5 │ 10 15 20 25

julia> a[Block(1,2)[[1,3],[2,4]]]
2×2 Matrix{Int64}:
11 21
13 23

julia> a[Block(2,2)[[2],[2,4]]]
1×2 Matrix{Int64}:
15 25
```
"""
BlockIndices

BlockIndices(block::Block{N}, inds::Vararg{AbstractVector,N}) where {N} =
BlockIndices(block,inds)
function BlockIndices(inds::Tuple{BlockIndices{1},Vararg{BlockIndices{1}}})
BlockIndices(Block(block.(inds)), map(ind -> ind.indices[1], inds))
end

const BlockIndexRange{N,R<:Tuple{Vararg{AbstractUnitRange{<:Integer},N}},I<:Tuple{Vararg{Any,N}},BI} = BlockIndices{N,R,I,BI}

"""
BlockIndexRange(block, startind:stopind)

Represents a cartesian range inside a block.
Represents a cartesian range inside a block. Type alias for `BlockIndices` with
the indices constrained to ranges.

It can be constructed and used to index into `BlockArrays` in the following manner:

Expand Down Expand Up @@ -260,60 +312,60 @@
"""
BlockIndexRange

BlockIndexRange(block::Block{N}, inds::Tuple{Vararg{AbstractUnitRange{<:Integer},N}}) where {N} =
BlockIndices(block, inds)
BlockIndexRange(block::Block{N}, inds::Vararg{AbstractUnitRange{<:Integer},N}) where {N} =
BlockIndexRange(block,inds)

BlockIndices(block, inds)
function BlockIndexRange(inds::Tuple{BlockIndexRange{1},Vararg{BlockIndexRange{1}}})
BlockIndexRange(Block(block.(inds)), map(ind -> ind.indices[1], inds))
end

block(R::BlockIndexRange) = R.block
block(R::BlockIndices) = R.block

copy(R::BlockIndexRange) = BlockIndexRange(R.block, map(copy, R.indices))
copy(R::BlockIndices) = BlockIndices(R.block, map(copy, R.indices))

getindex(::Block{0}) = BlockIndex()
getindex(B::Block{N}, inds::Vararg{Integer,N}) where N = BlockIndex(B,inds)
getindex(B::Block{N}, inds::Vararg{AbstractUnitRange{<:Integer},N}) where N = BlockIndexRange(B,inds)
getindex(B::Block{N}, inds::Vararg{Any,N}) where N = BlockIndex(B,inds)
getindex(B::Block{N}, inds::Vararg{AbstractVector,N}) where N = BlockIndices(B,inds)
getindex(B::Block{1}, inds::Colon) = B
getindex(B::Block{1}, inds::Base.Slice) = B
getindex(B::BlockIndices{0}) = B.block[]
@propagate_inbounds getindex(B::BlockIndices{N}, kr::Vararg{AbstractVector,N}) where N = BlockIndices(B.block, map(getindex, B.indices, kr))
@propagate_inbounds getindex(B::BlockIndices{N}, inds::Vararg{Int,N}) where N = B.block[Base.reindex(B.indices, inds)...]

Check warning on line 334 in src/blockindices.jl

View check run for this annotation

Codecov / codecov/patch

src/blockindices.jl#L334

Added line #L334 was not covered by tests

getindex(B::BlockIndexRange{0}) = B.block[]
@propagate_inbounds getindex(B::BlockIndexRange{N}, kr::Vararg{AbstractUnitRange{<:Integer},N}) where {N} = BlockIndexRange(B.block, map(getindex, B.indices, kr))
@propagate_inbounds getindex(B::BlockIndexRange{N}, inds::Vararg{Int,N}) where N = B.block[Base.reindex(B.indices, inds)...]

eltype(R::BlockIndexRange) = eltype(typeof(R))
eltype(::Type{BlockIndexRange{N}}) where {N} = BlockIndex{N}
eltype(::Type{BlockIndexRange{N,R,I,BI}}) where {N,R,I,BI} = BlockIndex{N,NTuple{N,BI},I}
IteratorSize(::Type{<:BlockIndexRange}) = Base.HasShape{1}()
eltype(R::BlockIndices) = eltype(typeof(R))
eltype(::Type{BlockIndices{N}}) where {N} = BlockIndex{N}

Check warning on line 337 in src/blockindices.jl

View check run for this annotation

Codecov / codecov/patch

src/blockindices.jl#L337

Added line #L337 was not covered by tests
eltype(::Type{BlockIndices{N,R,I,BI}}) where {N,R,I,BI} = BlockIndex{N,NTuple{N,BI},I}
IteratorSize(::Type{<:BlockIndices}) = Base.HasShape{1}()

Check warning on line 339 in src/blockindices.jl

View check run for this annotation

Codecov / codecov/patch

src/blockindices.jl#L339

Added line #L339 was not covered by tests


first(iter::BlockIndexRange) = BlockIndex(iter.block.n, map(first, iter.indices))
last(iter::BlockIndexRange) = BlockIndex(iter.block.n, map(last, iter.indices))
first(iter::BlockIndices) = BlockIndex(iter.block.n, map(first, iter.indices))
last(iter::BlockIndices) = BlockIndex(iter.block.n, map(last, iter.indices))

@inline function iterate(iter::BlockIndexRange)
@inline function iterate(iter::BlockIndices)
iterfirst, iterlast = first(iter), last(iter)
if any(map(>, iterfirst.α, iterlast.α))
return nothing
end
iterfirst, iterfirst
end
@inline function iterate(iter::BlockIndexRange, state)
@inline function iterate(iter::BlockIndices, state)
nextstate = BlockIndex(state.I, inc(state.α, first(iter).α, last(iter).α))
nextstate.α[end] > last(iter.indices[end]) && return nothing
nextstate, nextstate
end

size(iter::BlockIndexRange) = map(dimlength, first(iter).α, last(iter).α)
length(iter::BlockIndexRange) = prod(size(iter))
size(iter::BlockIndices) = map(dimlength, first(iter).α, last(iter).α)
length(iter::BlockIndices) = prod(size(iter))


Block(bs::BlockIndexRange) = bs.block
Block(bs::BlockIndices) = bs.block

##
# checkindex
##

function checkbounds(::Type{Bool}, A::AbstractArray{<:Any,N}, I::BlockIndexRange{N}) where N
function checkbounds(::Type{Bool}, A::AbstractArray{<:Any,N}, I::BlockIndices{N}) where N
bl = block(I)
checkbounds(Bool, A, bl) || return false
# TODO: Replace with `eachblockaxes(A)[bl]` once that is defined.
Expand All @@ -334,10 +386,11 @@
"""
BlockSlice(block, indices)

Represent an AbstractUnitRange{<:Integer} of indices that attaches a block.
Represents an AbstractUnitRange{<:Integer} of indices attached to a block,
a subblock, or a range of blocks.

Upon calling `to_indices()`, Blocks are converted to BlockSlice objects to represent
the indices over which the Block spans.
the indices over which the block, subblock, or range of blocks spans.

This mimics the relationship between `Colon` and `Base.Slice`.
"""
Expand All @@ -347,6 +400,7 @@
end

Block(bs::BlockSlice{<:Block}) = bs.block
Block(bs::BlockSlice{<:BlockIndices}) = Block(bs.block)


for f in (:axes, :unsafe_indices, :axes1, :first, :last, :size, :length,
Expand All @@ -366,36 +420,46 @@
# Avoid creating a SubArray wrapper in certain non-allocating cases
@propagate_inbounds view(C::CartesianIndices{N}, bs::Vararg{BlockSlice,N}) where {N} = view(C, map(x->x.indices, bs)...)

Block(bs::BlockSlice{<:BlockIndexRange}) = Block(bs.block)

"""
BlockedSlice(blocks, indices)
NoncontiguousBlockSlice(blocks, indices)

Represents blocked indices attached to a collection of corresponding blocks.
Represents an AbstractVector of indices attached to a (potentially non-contiguous) subblock,
set of blocks, or set of subblocks. This is the generalization of `BlockSlice` to
non-contiguous slices.

Upon calling `to_indices()`, a collection of blocks are converted to BlockedSlice objects to represent
Upon calling `to_indices()`, a collection of blocks are converted to NoncontiguousBlockSlice objects to represent
the indices over which the blocks span.

This mimics the relationship between `Colon` and `Base.Slice`, `Block` and `BlockSlice`, etc.
"""
struct BlockedSlice{BB,T<:Integer,INDS<:AbstractVector{T}} <: AbstractVector{T}
blocks::BB
struct NoncontiguousBlockSlice{BB,T,INDS<:AbstractVector{T}} <: AbstractVector{T}
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a relaunch of #459, hopefully with a better name, more context, and a better docstring. The idea is that BlockSlice is analogous to Slice for contiguous blockwise slices, for example when you slice with Block(2), Block(2)[1:2], and Block.(1:2) those get converted to BlockSlice by to_indices in order to store both the contiguous absolute range but also preserve the original blockwise information. Here, a NoncontiguousBlockSlice gets constructed by to_indices when you slice with non-contiguous blockwise slices such as [Block(1), Block(3)], [Block(1)[1:2], Block(3)[2:3]], and Block(2)[[2, 4]].

The reason for having two types is basically for the purpose of subyping: BlockSlice is an AbstractUnitRange and that information is used in a various parts of code, while this covers more general cases where the slices are not equivalent to unit ranges.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change is split off in #479.

block::BB
indices::INDS
end

for f in (:axes, :size)
@eval $f(S::BlockedSlice) = $f(S.indices)
Block(bs::NoncontiguousBlockSlice{<:Block}) = bs.block
Block(bs::NoncontiguousBlockSlice{<:BlockIndices}) = Block(bs.block)

for f in (:axes, :unsafe_indices, :axes1, :first, :last, :size, :length,
:unsafe_length, :start)
@eval $f(S::NoncontiguousBlockSlice) = $f(S.indices)
end

@propagate_inbounds getindex(S::BlockedSlice, i::Integer) = getindex(S.indices, i)
@propagate_inbounds getindex(S::BlockedSlice, k::Block{1}) = BlockSlice(S.blocks[Int(k)], getindex(S.indices, k))
_indices(B::NoncontiguousBlockSlice) = B.indices

@propagate_inbounds getindex(S::NoncontiguousBlockSlice, i::Integer) = getindex(S.indices, i)
@propagate_inbounds getindex(S::NoncontiguousBlockSlice{<:Block{1}}, k::AbstractVector{<:Integer}) =
NoncontiguousBlockSlice(S.block[_indices(k)], S.indices[_indices(k)])
@propagate_inbounds getindex(S::NoncontiguousBlockSlice{<:BlockIndices{1,<:Tuple{AbstractVector}}}, k::AbstractVector{<:Integer}) =
NoncontiguousBlockSlice(S.block[_indices(k)], S.indices[_indices(k)])
@propagate_inbounds getindex(S::NoncontiguousBlockSlice{<:AbstractVector{<:Block{1}}}, k::Block{1}) =
BlockSlice(S.block[Int(k)], getindex(S.indices, k))

struct BlockRange{N,R<:NTuple{N,AbstractUnitRange{<:Integer}}} <: AbstractArray{Block{N,Int},N}
indices::R
BlockRange{N,R}(inds::R) where {N,R} = new{N,R}(inds)
end


# The following is adapted from Julia v0.7 base/multidimensional.jl
# definition of CartesianRange

Expand Down
4 changes: 2 additions & 2 deletions src/show.jl
Original file line number Diff line number Diff line change
Expand Up @@ -186,13 +186,13 @@ function Base.show(io::IO, B::BlockIndex)
print(io, "]")
end

function Base.show(io::IO, B::BlockIndexRange)
function Base.show(io::IO, B::BlockIndices)
show(io, Block(B))
print(io, "[")
print_tuple_elements(io, B.indices)
print(io, "]")
end
Base.show(io::IO, ::MIME"text/plain", B::BlockIndexRange) = show(io, B)
Base.show(io::IO, ::MIME"text/plain", B::BlockIndices) = show(io, B)

Base.show(io::IO, mimetype::MIME"text/plain", a::AbstractBlockedUnitRange) =
Base.invoke(show, Tuple{typeof(io),MIME"text/plain",AbstractArray},io, mimetype, a)
Expand Down
Loading
Loading