Skip to content

Commit a4a0b93

Browse files
mtfishmandlfivefifty
mtfishman
and
dlfivefifty
authored
Nested slicing operations for merging blocks (JuliaArrays#466)
Fixes JuliaArrays#359. This introduces syntax for merging blocks by indexing with nested vectors of `Block{1}`, `BlockRange{1}`, `BlockIndexRange{1}`, etc. The implementation is very similar to JuliaArrays#445. As an example: ```julia julia> using BlockArrays julia> A = BlockArray(randn(8, 8), [2, 2, 2, 2], [2, 2, 2, 2]) 4×4-blocked 8×8 BlockMatrix{Float64}: 0.163738 -0.859246 │ -0.399256 -0.143974 │ -0.812783 0.469877 │ 0.0631785 1.56266 -0.209473 1.23701 │ 0.479233 0.219572 │ -0.548844 1.363 │ 0.795119 -0.169581 ───────────────────────┼─────────────────────────┼────────────────────────┼──────────────────────── 1.34493 -0.846004 │ 1.32566 0.470043 │ 0.650647 0.617614 │ -0.0926643 0.550304 1.20779 -0.528143 │ 1.09621 -0.0868135 │ -0.178719 -0.105456 │ 0.243687 -1.31676 ───────────────────────┼─────────────────────────┼────────────────────────┼──────────────────────── 1.201 1.48691 │ 0.264342 0.696739 │ -0.86067 -0.608839 │ -0.082433 -0.0931981 0.674171 -1.23494 │ 0.124443 1.94983 │ 0.929236 -1.62257 │ -0.693029 -0.0212493 ───────────────────────┼─────────────────────────┼────────────────────────┼──────────────────────── 0.0174839 1.65394 │ -0.851706 -1.23722 │ 1.63834 0.158663 │ 1.27422 -2.37923 1.28377 0.364039 │ 0.564611 -1.346 │ -0.619627 0.240933 │ 0.270664 -0.593572 julia> A[[[Block(1), Block(2)], [Block(3), Block(4)]], [[Block(1), Block(2)], [Block(3), Block(4)]]] 2×2-blocked 8×8 BlockedMatrix{Float64}: 0.163738 -0.859246 -0.399256 -0.143974 │ -0.812783 0.469877 0.0631785 1.56266 -0.209473 1.23701 0.479233 0.219572 │ -0.548844 1.363 0.795119 -0.169581 1.34493 -0.846004 1.32566 0.470043 │ 0.650647 0.617614 -0.0926643 0.550304 1.20779 -0.528143 1.09621 -0.0868135 │ -0.178719 -0.105456 0.243687 -1.31676 ──────────────────────────────────────────────┼────────────────────────────────────────────── 1.201 1.48691 0.264342 0.696739 │ -0.86067 -0.608839 -0.082433 -0.0931981 0.674171 -1.23494 0.124443 1.94983 │ 0.929236 -1.62257 -0.693029 -0.0212493 0.0174839 1.65394 -0.851706 -1.23722 │ 1.63834 0.158663 1.27422 -2.37923 1.28377 0.364039 0.564611 -1.346 │ -0.619627 0.240933 0.270664 -0.593572 julia> A[[[Block(1)[2:2], Block(2)[2:2]], [Block(3)[2:2], Block(4)[2:2]]], [[Block(1)[2:2], Block(2)[2:2]], [Block(3)[2:2], Block(4)[2:2]]]] 2×2-blocked 4×4 BlockedMatrix{Float64}: 1.23701 0.219572 │ 1.363 -0.169581 -0.528143 -0.0868135 │ -0.105456 -1.31676 ───────────────────────┼─────────────────────── -1.23494 1.94983 │ -1.62257 -0.0212493 0.364039 -1.346 │ 0.240933 -0.593572 ``` It feels unfortunate that the implementation has to be done in this way by catching the various cases manually with dispatch, but I can't think a better way (besides introducing a custom slicing operation). Co-authored-by: Sheehan Olver <[email protected]>
1 parent ecb8f8f commit a4a0b93

File tree

3 files changed

+42
-8
lines changed

3 files changed

+42
-8
lines changed

‎src/blockaxis.jl‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -426,9 +426,13 @@ end
426426
@propagate_inbounds getindex(b::AbstractBlockedUnitRange, KR::BlockSlice) = b[KR.block]
427427

428428
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:Block{1}}) = mortar([b[K] for K in KR])
429-
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:BlockIndexRange{1}}) = mortar([b[K] for K in KR])
430-
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:BlockIndex{1}}) = [b[K] for K in KR]
429+
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:AbstractVector{<:Block{1}}}) = mortar([b[K] for K in KR])
431430
getindex(b::AbstractBlockedUnitRange, Kkr::BlockIndexRange{1}) = b[block(Kkr)][Kkr.indices...]
431+
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:BlockIndex{1}}) = [b[K] for K in KR]
432+
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:BlockIndexRange{1}}) = mortar([b[K] for K in KR])
433+
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:AbstractVector{<:BlockIndex{1}}}) = mortar([b[K] for K in KR])
434+
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:AbstractVector{<:BlockIndexRange{1}}}) = mortar([b[K] for K in KR])
435+
getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:AbstractVector{<:AbstractVector{<:BlockIndex{1}}}}) = mortar([b[K] for K in KR])
432436

433437
_searchsortedfirst(a::AbstractVector, k) = searchsortedfirst(a, k)
434438
function _searchsortedfirst(a::Tuple, k)

‎src/views.jl‎

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -26,39 +26,54 @@ to_index(::BlockRange) = throw(ArgumentError("BlockRange must be converted by to
2626

2727
@inline to_indices(A, inds, I::Tuple{Block{1}, Vararg{Any}}) =
2828
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
29-
@inline to_indices(A, inds, I::Tuple{BlockRange{1,R}, Vararg{Any}}) where R =
29+
@inline to_indices(A, inds, I::Tuple{BlockRange{1}, Vararg{Any}}) =
30+
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
31+
@inline to_indices(A, inds, I::Tuple{AbstractVector{<:Block{1}}, Vararg{Any}}) =
32+
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
33+
@inline to_indices(A, inds, I::Tuple{AbstractVector{<:BlockRange{1}}, Vararg{Any}}) =
34+
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
35+
@inline to_indices(A, inds, I::Tuple{AbstractVector{<:AbstractVector{<:Block{1}}}, Vararg{Any}}) =
3036
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
3137
@inline to_indices(A, inds, I::Tuple{BlockIndex{1}, Vararg{Any}}) =
3238
(inds[1][I[1]], to_indices(A, _maybetail(inds), tail(I))...)
33-
@inline to_indices(A, inds, I::Tuple{BlockIndexRange{1,R}, Vararg{Any}}) where R =
34-
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
35-
@inline to_indices(A, inds, I::Tuple{AbstractVector{Block{1,R}}, Vararg{Any}}) where R =
39+
@inline to_indices(A, inds, I::Tuple{BlockIndexRange{1}, Vararg{Any}}) =
3640
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
3741
@inline to_indices(A, inds, I::Tuple{AbstractVector{<:BlockIndex{1}}, Vararg{Any}}) =
3842
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
3943
@inline to_indices(A, inds, I::Tuple{AbstractVector{<:BlockIndexRange{1}}, Vararg{Any}}) =
4044
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
45+
@inline to_indices(A, inds, I::Tuple{AbstractVector{<:AbstractVector{<:BlockIndex{1}}}, Vararg{Any}}) =
46+
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
47+
@inline to_indices(A, inds, I::Tuple{AbstractVector{<:AbstractVector{<:BlockIndexRange{1}}}, Vararg{Any}}) =
48+
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
49+
@inline to_indices(A, inds, I::Tuple{AbstractVector{<:AbstractVector{<:AbstractVector{<:BlockIndex{1}}}}, Vararg{Any}}) =
50+
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
4151

4252

4353
# splat out higher dimensional blocks
4454
# this mimics view of a CartesianIndex
4555
@inline to_indices(A, inds, I::Tuple{Block, Vararg{Any}}) =
4656
to_indices(A, inds, (Block.(I[1].n)..., tail(I)...))
57+
@inline to_indices(A, inds, I::Tuple{BlockRange, Vararg{Any}}) =
58+
to_indices(A, inds, (BlockRange.(tuple.(I[1].indices))..., tail(I)...))
4759
@inline to_indices(A, inds, I::Tuple{BlockIndex, Vararg{Any}}) =
4860
to_indices(A, inds, (BlockIndex.(I[1].I, I[1].α)..., tail(I)...))
4961
@inline to_indices(A, inds, I::Tuple{BlockIndexRange, Vararg{Any}}) =
5062
to_indices(A, inds, (BlockIndexRange.(Block.(I[1].block.n), tuple.(I[1].indices))..., tail(I)...))
51-
@inline to_indices(A, inds, I::Tuple{BlockRange, Vararg{Any}}) =
52-
to_indices(A, inds, (BlockRange.(tuple.(I[1].indices))..., tail(I)...))
5363

5464
# In 0.7, we need to override to_indices to avoid calling linearindices
5565
@inline to_indices(A, I::Tuple{BlockIndexRange, Vararg{Any}}) = to_indices(A, axes(A), I)
5666
@inline to_indices(A, I::Tuple{BlockIndex, Vararg{Any}}) = to_indices(A, axes(A), I)
5767
@inline to_indices(A, I::Tuple{Block, Vararg{Any}}) = to_indices(A, axes(A), I)
5868
@inline to_indices(A, I::Tuple{BlockRange, Vararg{Any}}) = to_indices(A, axes(A), I)
5969
@inline to_indices(A, I::Tuple{AbstractVector{<:Block{1}}, Vararg{Any}}) = to_indices(A, axes(A), I)
70+
@inline to_indices(A, I::Tuple{AbstractVector{<:BlockRange{1}}, Vararg{Any}}) = to_indices(A, axes(A), I)
71+
@inline to_indices(A, I::Tuple{AbstractVector{<:AbstractVector{<:Block{1}}}, Vararg{Any}}) = to_indices(A, axes(A), I)
6072
@inline to_indices(A, I::Tuple{AbstractVector{<:BlockIndex{1}}, Vararg{Any}}) = to_indices(A, axes(A), I)
6173
@inline to_indices(A, I::Tuple{AbstractVector{<:BlockIndexRange{1}}, Vararg{Any}}) = to_indices(A, axes(A), I)
74+
@inline to_indices(A, I::Tuple{AbstractVector{<:AbstractVector{<:BlockIndex{1}}}, Vararg{Any}}) = to_indices(A, axes(A), I)
75+
@inline to_indices(A, I::Tuple{AbstractVector{<:AbstractVector{<:BlockIndexRange{1}}}, Vararg{Any}}) = to_indices(A, axes(A), I)
76+
@inline to_indices(A, I::Tuple{AbstractVector{<:AbstractVector{<:AbstractVector{<:BlockIndex{1}}}}, Vararg{Any}}) = to_indices(A, axes(A), I)
6277

6378
## BlockedLogicalIndex
6479
# Blocked version of `LogicalIndex`:

‎test/test_blockarrays.jl‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1095,6 +1095,21 @@ end
10951095
@test a[[Block(1)[1:2], Block(2)[1:2]], [Block(1)[1:2], Block(2)[1:2]]] == [a[Block(1,1)[1:2,1:2]] a[Block(1,2)[1:2,1:2]]; a[Block(2,1)[1:2,1:2]] a[Block(2,2)[1:2,1:2]]]
10961096
@test a[[Block(1)[1], Block(2)[2]], [Block(1)[1:2], Block(2)[1:2]]] == [a[Block(1)[1],Block(1)[1:2]]' a[Block(1)[1], Block(2)[1:2]]'; a[Block(2)[2],Block(1)[1:2]]' a[Block(2)[2], Block(2)[1:2]]']
10971097
end
1098+
@testset "Blocked block-vector indexing (#359)" begin
1099+
for a in (BlockArray(randn(14, 14), 2:5, 2:5), BlockedArray(randn(14, 14), 2:5, 2:5))
1100+
for I in (
1101+
[Block.(1:2), Block.(3:4)],
1102+
[[Block(1), Block(3)], [Block(2), Block(4)]],
1103+
[[Block(1)[1:2], Block(3)[1:2]], [Block(2)[1:2], Block(4)[1:2]]],
1104+
[[[Block(1)[1], Block(1)[2]], [Block(3)[1], Block(3)[2]]], [[Block(2)[1], Block(2)[2]], [Block(4)[1], Block(4)[2]]]],
1105+
)
1106+
b = a[I, I]
1107+
for (i, j) in Iterators.product(1:length(I), 1:length(I))
1108+
@test a[I[i], I[j]] == b[Block(i), Block(j)]
1109+
end
1110+
end
1111+
end
1112+
end
10981113
end
10991114

11001115
end # module

0 commit comments

Comments
 (0)