Skip to content

Nested block indexing #482

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open

Nested block indexing #482

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
47f5936
Nested block indexing
mtfishman Jul 15, 2025
56b40a0
Improve test coverage
mtfishman Jul 15, 2025
521cb33
Merge branch 'JuliaArrays:master' into mf/blockindex_generalization
mtfishman Jul 15, 2025
cf14db5
More tests
mtfishman Jul 15, 2025
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
View file
Open in desktop
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
View file
Open in desktop
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
View file
Open in desktop
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
View file
Open in desktop
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