diff --git a/src/abstractblockarray.jl b/src/abstractblockarray.jl index 89ec0c1d..fad86696 100644 --- a/src/abstractblockarray.jl +++ b/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) diff --git a/src/blockindices.jl b/src/blockindices.jl index d2ce8540..fdd5e863 100644 --- a/src/blockindices.jl +++ b/src/blockindices.jl @@ -163,7 +163,7 @@ 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 @@ -171,23 +171,24 @@ 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)) @@ -268,14 +269,32 @@ function BlockIndexRange(inds::Tuple{BlockIndexRange{1},Vararg{BlockIndexRange{1 end block(R::BlockIndexRange) = R.block +blockindices(b::BlockIndexRange) = b.indices copy(R::BlockIndexRange) = BlockIndexRange(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{1}, inds) = BlockIndex(B,inds) getindex(B::Block{N}, inds::Vararg{AbstractUnitRange{<:Integer},N}) where N = BlockIndexRange(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::BlockIndexRange{0}) = B.block[] @propagate_inbounds getindex(B::BlockIndexRange{N}, kr::Vararg{AbstractUnitRange{<:Integer},N}) where {N} = BlockIndexRange(B.block, map(getindex, B.indices, kr)) diff --git a/test/test_blockarrays.jl b/test/test_blockarrays.jl index a2a73c5d..c97f8109 100644 --- a/test/test_blockarrays.jl +++ b/test/test_blockarrays.jl @@ -987,17 +987,37 @@ 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 "BlockIndexRange blocks" begin a = mortar([Block(1)[1:2], Block(3)[2:3]]) @test a[Block(1)] === Block(1)[1:2] diff --git a/test/test_blockindices.jl b/test/test_blockindices.jl index d39e0c50..3e0f929a 100644 --- a/test/test_blockindices.jl +++ b/test/test_blockindices.jl @@ -3,6 +3,7 @@ module TestBlockIndices using BlockArrays, FillArrays, Test, StaticArrays, ArrayLayouts using OffsetArrays import BlockArrays: BlockIndex, BlockIndexRange, BlockSlice, BlockedSlice +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, BlockedSlice @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)[1,1] == BlockIndex((1,1),(1,1)) == BlockIndex((1,1),(1,)) @test Block(1,1)[1:2,1:2] == BlockIndexRange(Block(1,1),(1:2,1:2)) @@ -98,6 +100,7 @@ import BlockArrays: BlockIndex, BlockIndexRange, BlockSlice, BlockedSlice @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) end @@ -987,4 +990,18 @@ 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 + end # module