Nested block indexing (JuliaArrays#482)

This enables nested block indexing by generalizing `BlockIndex` to allow
indices that aren't integers, for example:
```julia
julia> using BlockArrays

julia> a = mortar([[1, 2], [3, 4]])
2-blocked 4-element BlockVector{Int64}:
 1
 2
 ─
 3
 4

julia> b = mortar([[5, 6], [7, 8]])
2-blocked 4-element BlockVector{Int64}:
 5
 6
 ─
 7
 8

julia> A = mortar([a, b])
2-blocked 8-element BlockVector{Int64, Vector{BlockVector{Int64, Vector{Vector{Int64}}, Tuple{BlockedOneTo{Int64, Vector{Int64}}}}}, Tuple{BlockedOneTo{Int64, Vector{Int64}}}}:
 1
 2
 3
 4
 ─
 5
 6
 7
 8

julia> A[Block(2)[Block(1)]]
2-element Vector{Int64}:
 5
 6
```

As an alternative to just generalizing `BlockIndex`, maybe we could
define a new type such as `GenericBlockIndex` and then make `BlockIndex`
a type alias for that with the indices restricted to integers. I don't
have a strong opinion about that but this design was simpler and I don't
see much advantage one way or another.

Related to JuliaArrays#49 and JuliaArrays#446, though note that in this PR the nested block
structure isn't reflected in the axes, which would be required for more
systematic support of nested block structures. This PR is split off from
JuliaArrays#462.

---------

Co-authored-by: Sheehan Olver <[email protected]>

Author: mtfishman

.github/workflows/downstream.yml

@@ -34,10 +34,11 @@ jobs:
         julia-version: ['1']
         os: [ubuntu-latest]
         package:
-          - {repo: BlockBandedMatrices.jl, group: JuliaLinearAlgebra}
           - {repo: ApproxFunBase.jl, group: JuliaApproximation}
           - {repo: ContinuumArrays.jl, group: JuliaApproximation}
+          - {repo: MultivariateOrthogonalPolynomials.jl, group: JuliaApproximation}
           - {repo: LazyArrays.jl, group: JuliaArrays}
+          - {repo: BlockBandedMatrices.jl, group: JuliaLinearAlgebra}
           - {repo: LazyBandedMatrices.jl, group: JuliaLinearAlgebra}
           - {repo: InfiniteLinearAlgebra.jl, group: JuliaLinearAlgebra}
 

Project.toml

@@ -1,6 +1,6 @@
 name = "BlockArrays"
 uuid = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
-version = "1.8"
+version = "1.9"
 
 
 [deps]

src/abstractblockarray.jl

@@ -218,6 +218,10 @@ end
 @inline view(block_arr::AbstractBlockVector, block::Block{1}) = viewblock(block_arr, block)
 @propagate_inbounds view(block_arr::AbstractBlockArray, block::Block{1}...) = view(block_arr, Block(block))
 
+@inline view(block_arr::AbstractBlockArray{<:Any,N}, b::BlockIndex{N}) where N = view(view(block_arr, block(b)), blockindices(b)...)
+@inline view(block_arr::AbstractBlockVector, b::BlockIndex{1}) = view(view(block_arr, block(b)), blockindices(b)...)
+@propagate_inbounds view(block_arr::AbstractBlockArray{<:Any,N}, b::Vararg{BlockIndex{1},N}) where N = view(block_arr, BlockIndex(b))
+
 """
     eachblock(A::AbstractBlockArray)
 

src/blockindices.jl

@@ -163,31 +163,32 @@ julia> a[Block(2,2)[2,3]]
 20
 ```
 """
-struct BlockIndex{N,TI<:Tuple{Vararg{Integer,N}},Tα<:Tuple{Vararg{Integer,N}}}
+struct BlockIndex{N,TI<:Tuple{Vararg{Integer,N}},Tα<:Tuple{Vararg{Any,N}}}
     I::TI
     α::Tα
 end
 
 @inline BlockIndex(a::NTuple{N,Block{1}}, b::Tuple) where N = BlockIndex(Int.(a), b)
 @inline BlockIndex(::Tuple{}, b::Tuple{}) = BlockIndex{0,Tuple{},Tuple{}}((), ())
 
-@inline BlockIndex(a::Integer, b::Integer) = BlockIndex((a,), (b,))
-@inline BlockIndex(a::Tuple, b::Integer) = BlockIndex(a, (b,))
+@inline BlockIndex(a::Integer, b) = BlockIndex((a,), (b,))
+@inline BlockIndex(a::Tuple, b) = BlockIndex(a, (b,))
 @inline BlockIndex(a::Integer, b::Tuple) = BlockIndex((a,), b)
 @inline BlockIndex() = BlockIndex((), ())
 
 @inline BlockIndex(a::Block, b::Tuple) = BlockIndex(a.n, b)
-@inline BlockIndex(a::Block, b::Integer) = BlockIndex(a, (b,))
+@inline BlockIndex(a::Block, b) = BlockIndex(a, (b,))
 
-@inline function BlockIndex(I::Tuple{Vararg{Integer,N}}, α::Tuple{Vararg{Integer,M}}) where {M,N}
+@inline function BlockIndex(I::Tuple{Vararg{Integer,N}}, α::Tuple{Vararg{Any,M}}) where {M,N}
     M <= N || throw(ArgumentError("number of indices must not exceed the number of blocks"))
     α2 = ntuple(k -> k <= M ? α[k] : 1, N)
     BlockIndex(I, α2)
 end
 
 block(b::BlockIndex) = Block(b.I...)
-blockindex(b::BlockIndex{1}) = b.α[1]
-blockindex(b::BlockIndex) = CartesianIndex(b.α)
+blockindices(b::BlockIndex) = b.α
+blockindex(b::BlockIndex{1}) = blockindices(b)[1]
+blockindex(b::BlockIndex) = merge_indices(blockindices(b))
 
 BlockIndex(indcs::Tuple{Vararg{BlockIndex{1},N}}) where N = BlockIndex(block.(indcs), blockindex.(indcs))
 
@@ -318,15 +319,34 @@ function BlockIndexRange(inds::Tuple{BlockIndexRange{1},Vararg{BlockIndexRange{1
     BlockIndexRange(Block(block.(inds)), map(ind -> ind.indices[1], inds))
 end
 
+blockindices(b::BlockIndices) = b.indices
 block(R::BlockIndices) = R.block
 
 copy(R::BlockIndices) = BlockIndices(R.block, map(copy, R.indices))
 
+split_index(i::CartesianIndex) = Tuple(i)
+split_index(i::Block) = Tuple(i)
+split_index(i::BlockIndex) = map(BlockIndex, Tuple(block(i)), blockindices(i))
+split_index(i::BlockIndexRange) = map(BlockIndexRange, Tuple(block(i)), blockindices(i))
+
+merge_indices(i::Tuple{Vararg{Integer}}) = CartesianIndex(i)
+merge_indices(i::Tuple{Vararg{Block{1}}}) = Block(i)
+merge_indices(i::Tuple{Vararg{BlockIndex{1}}}) = BlockIndex(i)
+merge_indices(i::Tuple{Vararg{BlockIndexRange{1}}}) = BlockIndexRange(i)
+
 getindex(::Block{0}) = BlockIndex()
-getindex(B::Block{N}, inds::Vararg{Integer,N}) where N = BlockIndex(B,inds)
+getindex(B::Block{N}, inds::Vararg{Any,N}) where N = BlockIndex(B,inds)
 getindex(B::Block{N}, inds::Vararg{AbstractVector,N}) where N = BlockIndices(B,inds)
+getindex(B::Block{1}, inds) = BlockIndex(B,inds)
+getindex(B::Block{N}, inds::Vararg{AbstractUnitRange{<:Integer},N}) where N = BlockIndices(B,inds)
 getindex(B::Block{1}, inds::Colon) = B
 getindex(B::Block{1}, inds::Base.Slice) = B
+getindex(B::Block{N}, inds::Block{N}) where N = B[split_index(inds)...]
+getindex(B::Block{1}, inds::Block{1}) = BlockIndex(B,inds)
+getindex(B::Block{N}, inds::BlockIndex{N}) where N = B[split_index(inds)...]
+getindex(B::Block{1}, inds::BlockIndex{1}) = BlockIndex(B,inds)
+getindex(B::Block{N}, inds::BlockIndexRange{N}) where N = B[split_index(inds)...]
+getindex(B::Block{1}, inds::BlockIndexRange{1}) = BlockIndex(B,inds)
 
 getindex(B::BlockIndices{0}) = B.block[]
 @propagate_inbounds getindex(B::BlockIndices{N}, kr::Vararg{AbstractVector,N}) where {N} = BlockIndices(B.block, map(getindex, B.indices, kr))

test/test_blockarrays.jl

@@ -996,17 +996,36 @@ end
     @testset "Blockindex" begin
         a = BlockedArray(randn(3), [1,2])
         @test a[Block(1)[1]] == a[1]
-        @test view(a, Block(1)[1]) ≡ view(a, 1)
+        @test view(a, Block(1)[1]) ≡ view(parent(a), 1)
+        @test view(a, Block(1)[1]) == view(a, 1)
         @test a[Block(1)[1:1]] == a[1:1]
+
         A = BlockedArray(randn(3,3), [1,2], [1,2])
         @test A[Block(1)[1], Block(1)[1]] == A[Block(1,1)[1,1]] == A[1,1]
         # Regression test for #442
-        @test view(A, Block(1)[1], Block(1)[1]) ≡ view(A, Block(1,1)[1,1]) ≡ view(A, 1, 1)
+        @test view(A, Block(1)[1], Block(1)[1]) ≡ view(A, Block(1,1)[1,1]) ≡ view(parent(A), 1, 1)
         @test A[Block(1)[1:1], Block(1)[1:1]] == A[Block(1,1)[1:1,1:1]] == A[1:1,1:1]
         @test A[Block(1)[1:1], Block(1)[1]] == BlockArray(A)[Block(1)[1:1], Block(1)[1]] == A[1:1,1]
         @test A[Block(1)[1], Block(1)[1:1]] == BlockArray(A)[Block(1)[1], Block(1)[1:1]] == A[1,1:1]
     end
 
+    @testset "Nested block indexing" begin
+        va = BlockedArray(randn(4), [2,2])
+        vb = BlockedArray(randn(4), [2,2])
+        V = mortar([va,vb])
+        @test V[Block(2)[Block(1)]] == view(V, Block(2)[Block(1)]) == V[Block(2)][Block(1)] == vb[Block(1)]
+        @test V[Block(2)[Block(1)[2]]] == view(V, Block(2)[Block(1)[2]])[] == V[Block(2)][Block(1)[2]] == vb[Block(1)[2]]
+        @test V[Block(2)[Block(1)[1:2]]] == view(V, Block(2)[Block(1)[1:2]]) == V[Block(2)][Block(1)[1:2]] == vb[Block(1)[1:2]]
+
+        ma = BlockedArray(randn(4,4), [2,2], [2,2])
+        mb = BlockedArray(randn(4,4), [2,2], [2,2])
+        mc = BlockedArray(randn(4,4), [2,2], [2,2])
+        md = BlockedArray(randn(4,4), [2,2], [2,2])
+        M = mortar([[ma] [mc]; [mb] [md]])
+        @test M[Block(2,2)[Block(1,1)]] == view(M, Block(2,2)[Block(1,1)]) == M[Block(2,2)][Block(1,1)] == md[Block(1,1)]
+        @test M[Block(2,2)[Block(1,1)[2,2]]] == view(M, Block(2,2)[Block(1,1)[2,2]])[] == M[Block(2,2)][Block(1,1)[2,2]] == md[Block(1,1)[2,2]]
+        @test M[Block(2,2)[Block(1,1)[1:2,2:2]]] == view(M, Block(2,2)[Block(1,1)[1:2,2:2]]) == M[Block(2,2)][Block(1,1)[1:2,2:2]] == md[Block(1,1)[1:2,2:2]]
+    end
     @testset "BlockIndices" begin
         a = BlockedArray(randn(5), [2,3])
         @test a[Block(2)[[1,3]]] == a[[3,5]]

test/test_blockindices.jl

@@ -3,6 +3,7 @@ module TestBlockIndices
 using BlockArrays, FillArrays, Test, StaticArrays, ArrayLayouts
 using OffsetArrays
 import BlockArrays: BlockIndex, BlockIndexRange, BlockSlice, NoncontiguousBlockSlice
+import BlockArrays: split_index, merge_indices
 
 @testset "Blocks" begin
     @test Int(Block(2)) === Integer(Block(2)) === Number(Block(2)) === 2
@@ -86,6 +87,7 @@ import BlockArrays: BlockIndex, BlockIndexRange, BlockSlice, NoncontiguousBlockS
     @testset "BlockIndex" begin
         @test Block()[] == BlockIndex()
         @test Block(1)[1] == BlockIndex((1,),(1,))
+        @test Block(1)[Block(1)] == BlockIndex((1,),(Block(1),))
         @test Block(1)[1:2] == BlockIndexRange(Block(1),(1:2,))
         @test Block(1)[[1,3]] == BlockIndices(Block(1),([1,3],))
         @test Block(1,1)[1,1] == BlockIndex((1,1),(1,1)) == BlockIndex((1,1),(1,))
@@ -107,6 +109,7 @@ import BlockArrays: BlockIndex, BlockIndexRange, BlockSlice, NoncontiguousBlockS
         @test BlockIndex(UInt(2),(2,)) === BlockIndex((UInt(2),),(2,))
         @test BlockIndex(Block(2),2) === BlockIndex(Block(2),(2,))
         @test BlockIndex(Block(2),UInt(2)) === BlockIndex(Block(2),(UInt(2),))
+        @test BlockIndex(Block(2),Block(2)) === BlockIndex(Block(2),(Block(2),))
         @test copy(Block(1)[1:2]) === BlockIndexRange(Block(1),1:2)
         @test copy(Block(1)[[1,3]]) == BlockIndices(Block(1),[1,3])
         @test copy(Block(1)[[1,3]]) ≢ BlockIndices(Block(1),[1,3])
@@ -1021,6 +1024,19 @@ end
     @test a[[Block(1),Block(3)]] == a[Block.(1:2:3)] == [1,3]
 end
 
+@testset "split_index" begin
+    @test split_index(CartesianIndex(1, 2)) ≡ (1, 2)
+    @test split_index(Block(1, 2)) ≡ (Block(1), Block(2))
+    @test split_index(Block(1, 2)[3, 4]) ≡ (Block(1)[3], Block(2)[4])
+    @test split_index(Block(1, 2)[3:4, 4:5]) ≡ (Block(1)[3:4], Block(2)[4:5])
+end
+
+@testset "merge_indices" begin
+    @test merge_indices((1, 2)) ≡ CartesianIndex(1, 2)
+    @test merge_indices((Block(1), Block(2))) ≡ Block(1, 2)
+    @test merge_indices((Block(1)[3], Block(2)[4])) ≡ Block(1, 2)[3, 4]
+    @test merge_indices((Block(1)[3:4], Block(2)[4:5])) ≡ Block(1, 2)[3:4, 4:5]
+end
 @testset "to_index" begin
     @test_throws ArgumentError Base.to_index(Block(3))
     @test_throws ArgumentError Base.to_index(Block(3)[2])