BlockArrays.jl
BlockArrays.jl copied to clipboard
[WIP] `BlockIndices`
This is an implementation of #346.
TODO:
- [ ] When indexing an
AbstractArraywithBlockIndices(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
BlockIndiceswith 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 BlockedUnitRanges 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 BlockedUnitRanges 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 Matrixes, and ask users to use blockedranges 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 AbstractBlockArrays 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.