Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/abstractblockarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
35 changes: 27 additions & 8 deletions src/blockindices.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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))

Expand Down Expand Up @@ -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))
Expand Down
24 changes: 22 additions & 2 deletions test/test_blockarrays.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
17 changes: 17 additions & 0 deletions test/test_blockindices.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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))
Expand All @@ -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

Expand Down Expand Up @@ -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
Loading