There and back again: a block indexing tale

[mtfishman]

I think a helpful abstraction could be a BlockIndices type, analogous to the CartesianIndices type in Base Julia, to aid with converting back and forth between cartesian indices and block indices. I would find that very helpful for advanced slicing operations, for example when performing slicing operations across block arrays with mismatched blocking structures.

As a minimal demonstration:

using BlockArrays

struct BlockIndices{N,R<:NTuple{N,OrdinalRange{Int,Int}}} <: AbstractArray{BlockIndex{N},N}
  indices::R
end

BlockIndices(a::AbstractArray) = BlockIndices(axes(a))

function Base.getindex(a::BlockIndices{N}, index::Vararg{Integer,N}) where {N}
  return BlockIndex(findblockindex.(a.indices, index))
end

which would allow for this syntax for converting back and forth between cartesian indices and block indices:

julia> a = BlockArray(randn(4, 4), [2, 2], [2, 2])
2×2-blocked 4×4 BlockMatrix{Float64}:
  1.05776   -0.19332   │   0.804285  -0.0377738
 -1.00564    0.368672  │   1.26819    1.02757  
 ──────────────────────┼───────────────────────
  0.37596   -1.43335   │   0.280084  -1.06666  
  0.996146  -1.04132   │  -1.21102    0.661257 

julia> CartesianIndices(a)[Block(1, 2)[2, 1]]
CartesianIndex(2, 3)

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

I found it interesting that indexing into CartesianIndices with a BlockIndex already works, I assume it is because the CartesianIndices object gets the same blocked axes as the BlockArray.

Something that doesn't work right now is indexing into CartesianIndices with a Block and getting back out a CartesianIndices, for example:

julia> CartesianIndices(a)[Block(1, 2)]
2×2 Matrix{CartesianIndex{2}}:
 CartesianIndex(1, 3)  CartesianIndex(1, 4)
 CartesianIndex(2, 3)  CartesianIndex(2, 4)

but that should be simple to fix.

We may be able to make BlockIndices a subtype of AbstractBlockArray and automatically inherit a lot of AbstractBlockArray behavior, since it will have blocked axes.

[jishnub]

I quite like this idea, and I think this would certainly lead to code that is easier to reason about. Would you mind submitting a PR implementing this?

[mtfishman]

Sure, happy to make a PR with a basic implementation of BlockIndices.