BlockArrays.jl
BlockArrays.jl copied to clipboard
[WIP] `BlockIndices`
This is an implementation of #346.
TODO:
- [ ] When indexing an
AbstractArray
withBlockIndices
(i.e.getindex(a::AbstractArray, indices::BlockIndices)
) check that the block indices are compatible, say withblockisequal
. - [ ] Implement
view(a::CartesianIndices, indices::BlockIndices)
. - [ ] Implement
getindex(a::AbstractArray, indices::BlockIndices{1}...)
. - [ ] Fix indexing into
BlockIndices
with offset axes (https://github.com/JuliaArrays/BlockArrays.jl/pull/356#issuecomment-2027457046).
Here is a demonstration of what works so far:
julia> using BlockArrays
julia> using BlockArrays: BlockIndices
julia> A = BlockArray(randn(4, 4), [2, 2], [2, 2])
2×2-blocked 4×4 BlockMatrix{Float64}:
-1.74603 0.659762 │ -1.08084 1.75196
0.257445 0.793993 │ -0.922037 -0.802863
──────────────────────┼──────────────────────
1.40111 -0.570552 │ -2.91592 0.415012
-2.30074 1.10324 │ 1.52008 -0.846554
julia> B = BlockIndices(A)
2×2-blocked 4×4 BlockIndices{2, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
Block(1, 1)[1, 1] Block(1, 1)[1, 2] │ Block(1, 2)[1, 1] Block(1, 2)[1, 2]
Block(1, 1)[2, 1] Block(1, 1)[2, 2] │ Block(1, 2)[2, 1] Block(1, 2)[2, 2]
──────────────────────────────────────┼──────────────────────────────────────
Block(2, 1)[1, 1] Block(2, 1)[1, 2] │ Block(2, 2)[1, 1] Block(2, 2)[1, 2]
Block(2, 1)[2, 1] Block(2, 1)[2, 2] │ Block(2, 2)[2, 1] Block(2, 2)[2, 2]
julia> B[Block(2, 1)]
2×2 BlockArrays.BlockIndexRange{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}:
Block(2, 1)[1, 1] Block(2, 1)[1, 2]
Block(2, 1)[2, 1] Block(2, 1)[2, 2]
julia> B[2, 3]
Block(1, 2)[2, 1]
julia> B[Block(2, 1)[1, 2]]
Block(2, 1)[1, 2]
This is the current slicing behavior:
julia> B[Block(1, 1)[1:2, 1:2]]
2×2 Matrix{BlockIndex{2}}:
Block(1, 1)[1, 1] Block(1, 1)[1, 2]
Block(1, 1)[2, 1] Block(1, 1)[2, 2]
julia> B[2:3, 2:3]
2×2 Matrix{BlockIndex{2}}:
Block(1, 1)[2, 2] Block(1, 2)[2, 1]
Block(2, 1)[1, 2] Block(2, 2)[1, 1]
julia> B[Block(1):Block(2), Block(1):Block(2)]
2×2-blocked 4×4 PseudoBlockMatrix{BlockIndex{2}}:
Block(1, 1)[1, 1] Block(1, 1)[1, 2] │ Block(1, 2)[1, 1] Block(1, 2)[1, 2]
Block(1, 1)[2, 1] Block(1, 1)[2, 2] │ Block(1, 2)[2, 1] Block(1, 2)[2, 2]
──────────────────────────────────────┼──────────────────────────────────────
Block(2, 1)[1, 1] Block(2, 1)[1, 2] │ Block(2, 2)[1, 1] Block(2, 2)[1, 2]
Block(2, 1)[2, 1] Block(2, 1)[2, 2] │ Block(2, 2)[2, 1] Block(2, 2)[2, 2]
Ideally, any slicing operation that can get mapped to slicing with a unit range in each dimension (say slicing with AbstractUnitRange
, BlockRange
, BlockIndexRange
, etc.) would output BlockIndices
or BlockIndexRange
, but that isn't implemented yet. I think ideally that would be implemented as something like this:
function Base.getindex(a::BlockIndices, indices::Union{AbstractUnitRange,BlockRange,BlockIndexRange}...)
return BlockIndices(getindex.(a.indices, to_indices(a, indices)))
end
i.e. slicing a BlockIndices
slices the indices in each dimension, similar to the definition for Base.CartesianIndices
: https://github.com/JuliaLang/julia/blob/64de065a183ac70bb049f7f9e30d790f8845dd2b/base/multidimensional.jl#L384-L391. However, slicing blocked unit ranges doesn't always preserve blocking information right now, see for example #347 and #355.
@jishnub let me know what you think of the design of what is implemented, and also what you think the scope of the PR should be, i.e. how much work should go into getting more slicing working vs. iterating on that in future PRs.
Codecov Report
Attention: Patch coverage is 85.00000%
with 9 lines
in your changes are missing coverage. Please review.
Project coverage is 94.16%. Comparing base (
7e31cf8
) to head (43b459b
). Report is 20 commits behind head on master.
Files | Patch % | Lines |
---|---|---|
src/block_indices.jl | 85.00% | 9 Missing :warning: |
Additional details and impacted files
@@ Coverage Diff @@
## master #356 +/- ##
===========================================
+ Coverage 0.00% 94.16% +94.16%
===========================================
Files 16 17 +1
Lines 1480 1559 +79
===========================================
+ Hits 0 1468 +1468
+ Misses 1480 91 -1389
:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.
In the latest commit, I made it so that more slicing operations return either BlockIndices
or BlockIndexRange
, if that can be guaranteed based on the type of the indices (i.e. unit ranges, block ranges, etc.):
julia> using BlockArrays
julia> using BlockArrays: BlockIndices
julia> A = BlockArray(randn(4, 4), [2, 2], [2, 2])
2×2-blocked 4×4 BlockMatrix{Float64}:
-1.38012 -1.38908 │ 0.653502 -0.176679
-0.61462 -1.30911 │ -0.797909 0.083623
─────────────────────┼──────────────────────
-0.889299 1.48537 │ -0.83928 0.470424
1.01574 -1.31194 │ -0.282103 0.863026
julia> B = BlockIndices(A)
2×2-blocked 4×4 BlockIndices{2, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
Block(1, 1)[1, 1] Block(1, 1)[1, 2] │ Block(1, 2)[1, 1] Block(1, 2)[1, 2]
Block(1, 1)[2, 1] Block(1, 1)[2, 2] │ Block(1, 2)[2, 1] Block(1, 2)[2, 2]
──────────────────────────────────────┼──────────────────────────────────────
Block(2, 1)[1, 1] Block(2, 1)[1, 2] │ Block(2, 2)[1, 1] Block(2, 2)[1, 2]
Block(2, 1)[2, 1] Block(2, 1)[2, 2] │ Block(2, 2)[2, 1] Block(2, 2)[2, 2]
julia> B[Block(1, 1)[1:2, 1:2]]
2×2 BlockIndexRange{2, Tuple{UnitRange{Int64}, UnitRange{Int64}}}:
Block(1, 1)[1, 1] Block(1, 1)[1, 2]
Block(1, 1)[2, 1] Block(1, 1)[2, 2]
julia> B[2:3, 2:3]
2×2-blocked 2×2 BlockIndices{2, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
Block(1, 1)[2, 2] │ Block(1, 2)[2, 1]
───────────────────┼───────────────────
Block(2, 1)[1, 2] │ Block(2, 2)[1, 1]
julia> B[Block(2):Block(2), Block(1):Block(2)]
1×2-blocked 2×4 BlockIndices{2, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
Block(2, 1)[1, 1] Block(2, 1)[1, 2] │ Block(2, 2)[1, 1] Block(2, 2)[1, 2]
Block(2, 1)[2, 1] Block(2, 1)[2, 2] │ Block(2, 2)[2, 1] Block(2, 2)[2, 2]
I'm not sure if it is the best design for that code, but it isn't too complicated. Happy to refactor the code, change names and code style, etc. I also realized that the previous definition of BlockIndices
just based on wrapping a tuple of BlockedUnitRange
was too naive, since it couldn't handle the case when the first block wasn't Block(1, 1, ...)
. The new design is based on wrapping a tuple of BlockedUnitRange
for each dimension and then an index offset for each dimension indicating an offset into each of those BlockedUnitRange
s where the BlockIndices
starts. I'm not sure if that is the best design, but the code is fairly simple, and contiguous slices just involve adjusting the end point of the BlockedUnitRange
s and possibly shifting the offset value if the slice starts at a value greater than one.
I'm not fully sure about retaining block axes of the source in the indexing operation because of the reasons discussed earlier. E.g., this seems confusing:
julia> B = BlockIndices(A)
2×2-blocked 4×4 BlockIndices{2, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
Block(1, 1)[1, 1] Block(1, 1)[1, 2] │ Block(1, 2)[1, 1] Block(1, 2)[1, 2]
Block(1, 1)[2, 1] Block(1, 1)[2, 2] │ Block(1, 2)[2, 1] Block(1, 2)[2, 2]
──────────────────────────────────────┼──────────────────────────────────────
Block(2, 1)[1, 1] Block(2, 1)[1, 2] │ Block(2, 2)[1, 1] Block(2, 2)[1, 2]
Block(2, 1)[2, 1] Block(2, 1)[2, 2] │ Block(2, 2)[2, 1] Block(2, 2)[2, 2]
julia> b = BlockArrays._BlockedUnitRange(1, [1,4])
2-blocked 4-element BlockedUnitRange{Vector{Int64}}:
1
─
2
3
4
julia> B[b, b]
2×2-blocked 4×4 BlockIndices{2, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
Block(1, 1)[1, 1] Block(1, 1)[1, 2] │ Block(1, 2)[1, 1] Block(1, 2)[1, 2]
Block(1, 1)[2, 1] Block(1, 1)[2, 2] │ Block(1, 2)[2, 1] Block(1, 2)[2, 2]
──────────────────────────────────────┼──────────────────────────────────────
Block(2, 1)[1, 1] Block(2, 1)[1, 2] │ Block(2, 2)[1, 1] Block(2, 2)[1, 2]
Block(2, 1)[2, 1] Block(2, 1)[2, 2] │ Block(2, 2)[2, 1] Block(2, 2)[2, 2]
In this, the block axes of b
appear to have been ignored. Typically, it goes the other way, where the axes of the indices are retained. For example:
julia> A = BlockArray(randn(4, 4), [2, 2], [2, 2])
2×2-blocked 4×4 BlockMatrix{Float64}:
0.374444 0.888294 │ -1.25349 0.0794558
-0.230626 -0.21222 │ -1.50958 -1.98783
──────────────────────┼───────────────────────
0.846922 -0.493969 │ 0.557274 -1.0908
-1.35255 0.300943 │ -1.35591 -0.201546
julia> b = BlockArrays._BlockedUnitRange(1, [1,4])
2-blocked 4-element BlockedUnitRange{Vector{Int64}}:
1
─
2
3
4
julia> A[b, b]
2×2-blocked 4×4 BlockMatrix{Float64}:
0.374444 │ 0.888294 -1.25349 0.0794558
───────────┼──────────────────────────────────
-0.230626 │ -0.21222 -1.50958 -1.98783
0.846922 │ -0.493969 0.557274 -1.0908
-1.35255 │ 0.300943 -1.35591 -0.201546
I think it makes sense to be consistent with the convention here, where slices constructed using UnitRange
indices lead to Matrix
es, and ask users to use blockedrange
s to preserve the block information. This would also help with consistency in opertions such as
julia> B[b, b][Block(1), Block(1)]
2×2 BlockArrays.BlockIndexRange{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}:
Block(1, 1)[1, 1] Block(1, 1)[1, 2]
Block(1, 1)[2, 1] Block(1, 1)[2, 2]
julia> B[b[Block(1)], b[Block(1)]]
1×1-blocked 1×1 BlockIndices{2, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
Block(1, 1)[1, 1]
where both the two would be equivalent.
Right, I realize the behavior around slicing in this PR is different from that for other AbstractBlockArray
s right now, agreed that I should keep that consistent for the sake of this PR, and change it to the behavior you outlined.
I guess it is a different discussion to decide what slicing with UnitRange
should do, I'll continue that discussion in #347.
In the latest, I made it so the axes of the BlockIndices
resulting from a slice pick up the block structure of the input indices:
julia> using BlockArrays
julia> using BlockArrays: BlockIndices
julia> A = BlockArray(randn(4, 4), [2, 2], [2, 2])
2×2-blocked 4×4 BlockMatrix{Float64}:
1.09381 0.613925 │ 0.275561 -0.0245506
0.586465 3.50312 │ 1.51753 -0.123905
─────────────────────┼───────────────────────
0.613102 0.288321 │ -0.121851 -1.01473
-0.387993 0.871633 │ 1.15325 1.0105
julia> B = BlockIndices(A)
2×2-blocked 4×4 BlockIndices{2, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
Block(1, 1)[1, 1] Block(1, 1)[1, 2] │ Block(1, 2)[1, 1] Block(1, 2)[1, 2]
Block(1, 1)[2, 1] Block(1, 1)[2, 2] │ Block(1, 2)[2, 1] Block(1, 2)[2, 2]
──────────────────────────────────────┼──────────────────────────────────────
Block(2, 1)[1, 1] Block(2, 1)[1, 2] │ Block(2, 2)[1, 1] Block(2, 2)[1, 2]
Block(2, 1)[2, 1] Block(2, 1)[2, 2] │ Block(2, 2)[2, 1] Block(2, 2)[2, 2]
julia> B[Block(1, 2)[1:2, 1:2]]
2×2 BlockArrays.BlockIndexRange{2, Tuple{UnitRange{Int64}, UnitRange{Int64}}}:
Block(1, 2)[1, 1] Block(1, 2)[1, 2]
Block(1, 2)[2, 1] Block(1, 2)[2, 2]
julia> B[Block.(2:2), Block.(1:2)]
1×2-blocked 2×4 BlockIndices{2, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
Block(2, 1)[1, 1] Block(2, 1)[1, 2] │ Block(2, 2)[1, 1] Block(2, 2)[1, 2]
Block(2, 1)[2, 1] Block(2, 1)[2, 2] │ Block(2, 2)[2, 1] Block(2, 2)[2, 2]
julia> B[1:4, 1:4]
1×1-blocked 4×4 BlockIndices{2, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}:
Block(1, 1)[1, 1] Block(1, 1)[1, 2] Block(1, 2)[1, 1] Block(1, 2)[1, 2]
Block(1, 1)[2, 1] Block(1, 1)[2, 2] Block(1, 2)[2, 1] Block(1, 2)[2, 2]
Block(2, 1)[1, 1] Block(2, 1)[1, 2] Block(2, 2)[1, 1] Block(2, 2)[1, 2]
Block(2, 1)[2, 1] Block(2, 1)[2, 2] Block(2, 2)[2, 1] Block(2, 2)[2, 2]
julia> B[2:3, 2:3]
1×1-blocked 2×2 BlockIndices{2, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}:
Block(1, 1)[2, 2] Block(1, 2)[2, 1]
Block(2, 1)[1, 2] Block(2, 2)[1, 1]
julia> B[blockedrange([1, 3]), blockedrange([1, 3])]
2×2-blocked 4×4 BlockIndices{2, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
Block(1, 1)[1, 1] │ Block(1, 1)[1, 2] Block(1, 2)[1, 1] Block(1, 2)[1, 2]
───────────────────┼─────────────────────────────────────────────────────────
Block(1, 1)[2, 1] │ Block(1, 1)[2, 2] Block(1, 2)[2, 1] Block(1, 2)[2, 2]
Block(2, 1)[1, 1] │ Block(2, 1)[1, 2] Block(2, 2)[1, 1] Block(2, 2)[1, 2]
Block(2, 1)[2, 1] │ Block(2, 1)[2, 2] Block(2, 2)[2, 1] Block(2, 2)[2, 2]
This required a slight redesign of the BlockIndices
type, where it now stores one set of (blocked) indices to compute the elements of the BlockIndices
, and another to store the axes, which may or may not be blocked. It seems a bit strange at first, but I think in the end it makes sense is analogous to situations like #354 where a CartesianIndices
may or may not have blocking structure, depending on the axes.
@jishnub let me know what you think.
One behavior I wasn't sure about was indexing with Colon
. Currently it preserves the block structure:
julia> B[:, :]
2×2-blocked 4×4 PseudoBlockMatrix{BlockIndex{2}}:
Block(1, 1)[1, 1] Block(1, 1)[1, 2] │ Block(1, 2)[1, 1] Block(1, 2)[1, 2]
Block(1, 1)[2, 1] Block(1, 1)[2, 2] │ Block(1, 2)[2, 1] Block(1, 2)[2, 2]
──────────────────────────────────────┼──────────────────────────────────────
Block(2, 1)[1, 1] Block(2, 1)[1, 2] │ Block(2, 2)[1, 1] Block(2, 2)[1, 2]
Block(2, 1)[2, 1] Block(2, 1)[2, 2] │ Block(2, 2)[2, 1] Block(2, 2)[2, 2]
(though in principle it should be outputting a BlockIndices
object), which is the same blocking behavior as BlockArray
:
julia> A[:, :]
2×2-blocked 4×4 BlockMatrix{Float64}:
1.09381 0.613925 │ 0.275561 -0.0245506
0.586465 3.50312 │ 1.51753 -0.123905
─────────────────────┼───────────────────────
0.613102 0.288321 │ -0.121851 -1.01473
-0.387993 0.871633 │ 1.15325 1.0105
It seems a bit strange to me that :
would act differently from 1:end
.
Regarding the Colon, this is the standard behavior, which differs from indexing using a UnitRange
. A colon gets lowered to a Base.Slice
object (or perhaps a BlockSlice
here) that inherits the axes from the parent, whereas a UnitRange
has its own axes that are preserved. The fact that they act differently is indeed a bit confusing, but it's consistent with how indexing works elsewhere in Julia.
It looks like the following errors at present:
julia> A = BlockArray(randn(4, 4), [2, 2], [2, 2]);
julia> B = BlockArrays.BlockIndices(A);
julia> A[B]
ERROR: ArgumentError: unable to check bounds for indices of type BlockIndex{2}
Stacktrace:
[1] checkindex(::Type{Bool}, inds::Base.OneTo{Int64}, i::BlockIndex{2})
@ Base ./abstractarray.jl:759
[2] checkindex
@ ./abstractarray.jl:776 [inlined]
[3] checkbounds
@ ./abstractarray.jl:687 [inlined]
[4] checkbounds
@ ./abstractarray.jl:702 [inlined]
[5] _getindex
@ ./multidimensional.jl:888 [inlined]
[6] getindex(A::BlockMatrix{Float64, Matrix{…}, Tuple{…}}, I::BlockIndices{2, Tuple{…}, Tuple{…}})
@ Base ./abstractarray.jl:1291
[7] top-level scope
@ REPL[14]:1
Some type information was truncated. Use `show(err)` to see complete types.
I wonder if this may be made to work?
I don't think this is handling offset axes correctly:
julia> B
2×2-blocked 4×4 BlockIndices{2, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
Block(1, 1)[1, 1] Block(1, 1)[1, 2] │ Block(1, 2)[1, 1] Block(1, 2)[1, 2]
Block(1, 1)[2, 1] Block(1, 1)[2, 2] │ Block(1, 2)[2, 1] Block(1, 2)[2, 2]
──────────────────────────────────────┼──────────────────────────────────────
Block(2, 1)[1, 1] Block(2, 1)[1, 2] │ Block(2, 2)[1, 1] Block(2, 2)[1, 2]
Block(2, 1)[2, 1] Block(2, 1)[2, 2] │ Block(2, 2)[2, 1] Block(2, 2)[2, 2]
julia> @view B[Base.IdentityUnitRange(2:4), Base.IdentityUnitRange(2:4)]
1×1-blocked 3×3 BlockIndices{2, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}, Tuple{Base.IdentityUnitRange{UnitRange{Int64}}, Base.IdentityUnitRange{UnitRange{Int64}}}}:
Block(2, 2)[1, 1] Block(2, 2)[1, 2] #undef
Block(2, 2)[2, 1] Block(2, 2)[2, 2] #undef
#undef #undef #undef
Yeah, I didn't look into indexing arrays with BlockIndices
at all yet, but ideally that would work.
Good point about offset axes, I also didn't have that use case in mind but hopefully is easy to support.
In the latest I added support for indexing with BlockIndices
:
julia> using BlockArrays
julia> using BlockArrays: BlockIndices
julia> A = BlockArray(randn(4, 4), [2, 2], [2, 2])
2×2-blocked 4×4 BlockMatrix{Float64}:
-2.27028 1.18245 │ -0.0100978 0.219669
0.534015 -2.69366 │ 0.131658 -0.131192
─────────────────────────┼───────────────────────
0.0449745 -0.288047 │ 1.71181 0.76368
1.27463 -0.00159753 │ 0.0061755 -0.427065
julia> r = BlockArrays._BlockedUnitRange(2, [2, 3])
2-blocked 2-element BlockedUnitRange{Vector{Int64}}:
2
─
3
julia> B = BlockIndices(A)[r, r]
2×2-blocked 2×2 BlockIndices{2, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
Block(1, 1)[2, 2] │ Block(1, 2)[2, 1]
───────────────────┼───────────────────
Block(2, 1)[1, 2] │ Block(2, 2)[1, 1]
julia> CartesianIndices(A)[B]
CartesianIndices((2:1:3, 2:1:3))
julia> A[B]
2×2-blocked 2×2 BlockMatrix{Float64}:
-2.69366 │ 0.131658
───────────┼──────────
-0.288047 │ 1.71181
though I think I still need to check that the BlockIndices
are compatible with the blocking of the array being indexed into.