Nested block indexing

.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])