Nested slicing operations for merging blocks (#466)

Fixes #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 #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]>

Author: mtfishman

src/blockaxis.jl

@@ -426,9 +426,13 @@ end
 @propagate_inbounds getindex(b::AbstractBlockedUnitRange, KR::BlockSlice) = b[KR.block]
 
 getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:Block{1}}) = mortar([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{<:BlockIndex{1}}) = [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, 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{<: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{<:AbstractVector{<:BlockIndex{1}}}}) = mortar([b[K] for K in KR])
 
 _searchsortedfirst(a::AbstractVector, k) = searchsortedfirst(a, k)
 function _searchsortedfirst(a::Tuple, k)

src/views.jl

@@ -26,39 +26,54 @@ to_index(::BlockRange) = throw(ArgumentError("BlockRange must be converted by to
 
 @inline to_indices(A, inds, I::Tuple{Block{1}, Vararg{Any}}) =
     (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
-@inline to_indices(A, inds, I::Tuple{BlockRange{1,R}, Vararg{Any}}) where R =
+@inline to_indices(A, inds, I::Tuple{BlockRange{1}, Vararg{Any}}) =
+    (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
+@inline to_indices(A, inds, I::Tuple{AbstractVector{<:Block{1}}, Vararg{Any}}) =
+    (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
+@inline to_indices(A, inds, I::Tuple{AbstractVector{<:BlockRange{1}}, Vararg{Any}}) =
+    (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
+@inline to_indices(A, inds, I::Tuple{AbstractVector{<:AbstractVector{<:Block{1}}}, Vararg{Any}}) =
     (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
 @inline to_indices(A, inds, I::Tuple{BlockIndex{1}, Vararg{Any}}) =
     (inds[1][I[1]], to_indices(A, _maybetail(inds), tail(I))...)
-@inline to_indices(A, inds, I::Tuple{BlockIndexRange{1,R}, Vararg{Any}}) where R =
-    (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
-@inline to_indices(A, inds, I::Tuple{AbstractVector{Block{1,R}}, Vararg{Any}}) where R =
+@inline to_indices(A, inds, I::Tuple{BlockIndexRange{1}, Vararg{Any}}) =
     (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
 @inline to_indices(A, inds, I::Tuple{AbstractVector{<:BlockIndex{1}}, Vararg{Any}}) =
     (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
 @inline to_indices(A, inds, I::Tuple{AbstractVector{<:BlockIndexRange{1}}, Vararg{Any}}) =
     (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
+@inline to_indices(A, inds, I::Tuple{AbstractVector{<:AbstractVector{<:BlockIndex{1}}}, Vararg{Any}}) =
+    (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
+@inline to_indices(A, inds, I::Tuple{AbstractVector{<:AbstractVector{<:BlockIndexRange{1}}}, Vararg{Any}}) =
+    (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
+@inline to_indices(A, inds, I::Tuple{AbstractVector{<:AbstractVector{<:AbstractVector{<:BlockIndex{1}}}}, Vararg{Any}}) =
+    (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
 
 
 # splat out higher dimensional blocks
 # this mimics view of a CartesianIndex
 @inline to_indices(A, inds, I::Tuple{Block, Vararg{Any}}) =
     to_indices(A, inds, (Block.(I[1].n)..., tail(I)...))
+@inline to_indices(A, inds, I::Tuple{BlockRange, Vararg{Any}}) =
+    to_indices(A, inds, (BlockRange.(tuple.(I[1].indices))..., tail(I)...))
 @inline to_indices(A, inds, I::Tuple{BlockIndex, Vararg{Any}}) =
     to_indices(A, inds, (BlockIndex.(I[1].I, I[1].α)..., tail(I)...))
 @inline to_indices(A, inds, I::Tuple{BlockIndexRange, Vararg{Any}}) =
     to_indices(A, inds, (BlockIndexRange.(Block.(I[1].block.n), tuple.(I[1].indices))..., tail(I)...))
-@inline to_indices(A, inds, I::Tuple{BlockRange, Vararg{Any}}) =
-    to_indices(A, inds, (BlockRange.(tuple.(I[1].indices))..., tail(I)...))
 
 # In 0.7, we need to override to_indices to avoid calling linearindices
 @inline to_indices(A, I::Tuple{BlockIndexRange, Vararg{Any}}) = to_indices(A, axes(A), I)
 @inline to_indices(A, I::Tuple{BlockIndex, Vararg{Any}}) = to_indices(A, axes(A), I)
 @inline to_indices(A, I::Tuple{Block, Vararg{Any}}) = to_indices(A, axes(A), I)
 @inline to_indices(A, I::Tuple{BlockRange, Vararg{Any}}) = to_indices(A, axes(A), I)
 @inline to_indices(A, I::Tuple{AbstractVector{<:Block{1}}, Vararg{Any}}) = to_indices(A, axes(A), I)
+@inline to_indices(A, I::Tuple{AbstractVector{<:BlockRange{1}}, Vararg{Any}}) = to_indices(A, axes(A), I)
+@inline to_indices(A, I::Tuple{AbstractVector{<:AbstractVector{<:Block{1}}}, Vararg{Any}}) = to_indices(A, axes(A), I)
 @inline to_indices(A, I::Tuple{AbstractVector{<:BlockIndex{1}}, Vararg{Any}}) = to_indices(A, axes(A), I)
 @inline to_indices(A, I::Tuple{AbstractVector{<:BlockIndexRange{1}}, Vararg{Any}}) = to_indices(A, axes(A), I)
+@inline to_indices(A, I::Tuple{AbstractVector{<:AbstractVector{<:BlockIndex{1}}}, Vararg{Any}}) = to_indices(A, axes(A), I)
+@inline to_indices(A, I::Tuple{AbstractVector{<:AbstractVector{<:BlockIndexRange{1}}}, Vararg{Any}}) = to_indices(A, axes(A), I)
+@inline to_indices(A, I::Tuple{AbstractVector{<:AbstractVector{<:AbstractVector{<:BlockIndex{1}}}}, Vararg{Any}}) = to_indices(A, axes(A), I)
 
 ## BlockedLogicalIndex
 # Blocked version of `LogicalIndex`:

test/test_blockarrays.jl

@@ -1095,6 +1095,21 @@ end
         @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]]]
         @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]]']
     end
+    @testset "Blocked block-vector indexing (#359)" begin
+        for a in (BlockArray(randn(14, 14), 2:5, 2:5), BlockedArray(randn(14, 14), 2:5, 2:5))
+            for I in (
+                [Block.(1:2), Block.(3:4)],
+                [[Block(1), Block(3)], [Block(2), Block(4)]],
+                [[Block(1)[1:2], Block(3)[1:2]], [Block(2)[1:2], Block(4)[1:2]]],
+                [[[Block(1)[1], Block(1)[2]], [Block(3)[1], Block(3)[2]]], [[Block(2)[1], Block(2)[2]], [Block(4)[1], Block(4)[2]]]],
+            )
+                b = a[I, I]
+                for (i, j) in Iterators.product(1:length(I), 1:length(I))
+                    @test a[I[i], I[j]] == b[Block(i), Block(j)]
+                end
+            end
+        end
+    end
 end
 
 end # module