From aca3c701c07f9ca5fe676468fb1e2184ce7575d3 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Tue, 26 Mar 2024 13:43:43 -0400
Subject: [PATCH 1/8] BlockIndices
---
src/BlockArrays.jl | 1 +
src/block_indices.jl | 18 ++++++++++++++++++
test/test_blockindices.jl | 31 ++++++++++++++++++++++++++++++-
3 files changed, 49 insertions(+), 1 deletion(-)
create mode 100644 src/block_indices.jl
diff --git a/src/BlockArrays.jl b/src/BlockArrays.jl
index c05a66dc..b728a7ba 100644
--- a/src/BlockArrays.jl
+++ b/src/BlockArrays.jl
@@ -56,6 +56,7 @@ end
include("blockindices.jl")
include("blockaxis.jl")
include("abstractblockarray.jl")
+include("block_indices.jl")
include("blockarray.jl")
include("pseudo_blockarray.jl")
include("views.jl")
diff --git a/src/block_indices.jl b/src/block_indices.jl
new file mode 100644
index 00000000..ca030662
--- /dev/null
+++ b/src/block_indices.jl
@@ -0,0 +1,18 @@
+struct BlockIndices{N,R<:NTuple{N,OrdinalRange{Int,Int}}} <: AbstractBlockArray{BlockIndex{N},N}
+ indices::R
+end
+Base.axes(a::BlockIndices) = map(Base.axes1, a.indices)
+BlockIndices(a::AbstractArray) = BlockIndices(axes(a))
+
+function Base.getindex(a::BlockIndices{N}, index::Vararg{Integer,N}) where {N}
+ return BlockIndex(findblockindex.(axes(a), index))
+end
+
+function Base.view(a::BlockIndices{N}, block::Block{N}) where {N}
+ return viewblock(a, block)
+end
+
+function viewblock(a::BlockIndices, block)
+ range = Base.OneTo.(getindex.(blocklengths.(axes(a)), Int.(Tuple(block))))
+ return block[range...]
+end
diff --git a/test/test_blockindices.jl b/test/test_blockindices.jl
index f3beadf2..d2540c6c 100644
--- a/test/test_blockindices.jl
+++ b/test/test_blockindices.jl
@@ -1,6 +1,6 @@
using BlockArrays, FillArrays, Test, StaticArrays, ArrayLayouts
using OffsetArrays
-import BlockArrays: BlockIndex, BlockIndexRange, BlockSlice
+import BlockArrays: BlockIndex, BlockIndexRange, BlockIndices, BlockSlice
@testset "Blocks" begin
@test Int(Block(2)) === Integer(Block(2)) === Number(Block(2)) === 2
@@ -446,3 +446,32 @@ end
first(eachblock(B))[1,2] = 0
@test B[1,2] == 0
end
+
+@testset "BlockIndices" begin
+ v = Array(reshape(1:35, (5, 7)))
+ A = BlockArray(v, [2,3], [3,4])
+ B = BlockIndices(A)
+ @test B == BlockIndices((blockedrange([2,3]), blockedrange([3,4])))
+ @test eltype(B) === BlockIndex{2}
+ @test size(B) == (5, 7)
+ @test axes(B) == (1:5, 1:7)
+ @test blocksize(B) == (2, 2)
+ @test blockaxes(B) == (Block(1):Block(2), Block(1):Block(2))
+ @test B[1, 1] == Block(1, 1)[1, 1]
+ @test B[4, 6] == Block(2, 2)[2, 3]
+
+ @test B[Block(1, 1)] isa BlockArrays.BlockIndexRange{2}
+ @test B[Block(1, 1)] == Block(1, 1)[1:2, 1:3]
+ @test B[Block(2, 1)] == Block(2, 1)[1:3, 1:3]
+ @test B[Block(1, 2)] == Block(1, 2)[1:2, 1:4]
+ @test B[Block(2, 2)] == Block(2, 2)[1:3, 1:4]
+
+ @test view(B, Block(1, 2)) isa BlockArrays.BlockIndexRange{2}
+ @test view(B, Block(1, 2)) == Block(1, 2)[1:2, 1:4]
+
+ @test B[Block(1), Block(2)] isa BlockArrays.BlockIndexRange{2}
+ @test B[Block(1), Block(2)] == Block(1, 2)[1:2, 1:4]
+
+ @test view(B, Block(1), Block(2)) isa BlockArrays.BlockIndexRange{2}
+ @test view(B, Block(1), Block(2)) == Block(1, 2)[1:2, 1:4]
+end
From 010bfed912746a1f8df2438453d1ebee805c8ef9 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Tue, 26 Mar 2024 16:43:49 -0400
Subject: [PATCH 2/8] Fix for non-one-based ranges
---
src/block_indices.jl | 14 +++++++++++++-
test/test_blockindices.jl | 22 ++++++++++++++++++----
2 files changed, 31 insertions(+), 5 deletions(-)
diff --git a/src/block_indices.jl b/src/block_indices.jl
index ca030662..9e28975e 100644
--- a/src/block_indices.jl
+++ b/src/block_indices.jl
@@ -4,14 +4,26 @@ end
Base.axes(a::BlockIndices) = map(Base.axes1, a.indices)
BlockIndices(a::AbstractArray) = BlockIndices(axes(a))
+function to_blockindex(a::AbstractUnitRange, index::Integer)
+ axis_blockindex = findblockindex(only(axes(a)), index)
+ if !isone(first(a)) && block(axis_blockindex) == Block(1)
+ axis_blockindex = block(axis_blockindex)[blockindex(axis_blockindex) + first(a) - one(eltype(a))]
+ end
+ return axis_blockindex
+end
+
function Base.getindex(a::BlockIndices{N}, index::Vararg{Integer,N}) where {N}
- return BlockIndex(findblockindex.(axes(a), index))
+ return BlockIndex(to_blockindex.(a.indices, index))
end
function Base.view(a::BlockIndices{N}, block::Block{N}) where {N}
return viewblock(a, block)
end
+function Base.view(a::BlockIndices, block::Vararg{Block{1}})
+ return view(a, Block(block))
+end
+
function viewblock(a::BlockIndices, block)
range = Base.OneTo.(getindex.(blocklengths.(axes(a)), Int.(Tuple(block))))
return block[range...]
diff --git a/test/test_blockindices.jl b/test/test_blockindices.jl
index d2540c6c..2cbdbe3a 100644
--- a/test/test_blockindices.jl
+++ b/test/test_blockindices.jl
@@ -460,18 +460,32 @@ end
@test B[1, 1] == Block(1, 1)[1, 1]
@test B[4, 6] == Block(2, 2)[2, 3]
- @test B[Block(1, 1)] isa BlockArrays.BlockIndexRange{2}
+ @test B[Block(1, 1)] isa BlockIndexRange{2}
@test B[Block(1, 1)] == Block(1, 1)[1:2, 1:3]
@test B[Block(2, 1)] == Block(2, 1)[1:3, 1:3]
@test B[Block(1, 2)] == Block(1, 2)[1:2, 1:4]
@test B[Block(2, 2)] == Block(2, 2)[1:3, 1:4]
- @test view(B, Block(1, 2)) isa BlockArrays.BlockIndexRange{2}
+ @test view(B, Block(1, 2)) isa BlockIndexRange{2}
@test view(B, Block(1, 2)) == Block(1, 2)[1:2, 1:4]
- @test B[Block(1), Block(2)] isa BlockArrays.BlockIndexRange{2}
+ @test B[Block(1), Block(2)] isa BlockIndexRange{2}
@test B[Block(1), Block(2)] == Block(1, 2)[1:2, 1:4]
- @test view(B, Block(1), Block(2)) isa BlockArrays.BlockIndexRange{2}
+ @test view(B, Block(1), Block(2)) isa BlockIndexRange{2}
@test view(B, Block(1), Block(2)) == Block(1, 2)[1:2, 1:4]
+
+ r1 = BlockArrays._BlockedUnitRange(2, [2, 3])
+ r2 = BlockArrays._BlockedUnitRange(2, [3, 4])
+ @test B[2:3, 2:4] == BlockIndices((r1, r2))
+ @test_broken B[2:3, 2:4] isa BlockIndices{2}
+
+ @test B[Block(2, 2)[2:3, 2:3]] == Block(2, 2)[2:3, 2:3]
+ @test_broken B[Block(2, 2)[2:3, 2:3]] isa BlockedIndexRange{2}
+
+ @test B[Block(1):Block(2), Block(1):Block(2)] == B
+ @test_broken B[Block(1):Block(2), Block(1):Block(2)] isa BlockedIndices{2}
+
+ @test B[Block(2), Block(1):Block(2)] == mortar([[Block(2, 1)[1:3, 1:3]] [Block(2, 2)[1:3, 1:4]]])
+ @test_broken B[Block(2), Block(1):Block(2)] isa BlockedIndices{2}
end
From ea2e171d960ace3332ed1700bec97ab94472cc31 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Wed, 27 Mar 2024 11:10:32 -0400
Subject: [PATCH 3/8] Preserve BlockIndices when possible
---
src/block_indices.jl | 88 +++++++++++++++++++++++++++++++++++----
test/test_blockindices.jl | 38 +++++++++++++----
2 files changed, 109 insertions(+), 17 deletions(-)
diff --git a/src/block_indices.jl b/src/block_indices.jl
index 9e28975e..a10f908c 100644
--- a/src/block_indices.jl
+++ b/src/block_indices.jl
@@ -1,9 +1,3 @@
-struct BlockIndices{N,R<:NTuple{N,OrdinalRange{Int,Int}}} <: AbstractBlockArray{BlockIndex{N},N}
- indices::R
-end
-Base.axes(a::BlockIndices) = map(Base.axes1, a.indices)
-BlockIndices(a::AbstractArray) = BlockIndices(axes(a))
-
function to_blockindex(a::AbstractUnitRange, index::Integer)
axis_blockindex = findblockindex(only(axes(a)), index)
if !isone(first(a)) && block(axis_blockindex) == Block(1)
@@ -12,19 +6,95 @@ function to_blockindex(a::AbstractUnitRange, index::Integer)
return axis_blockindex
end
+function blockedunitrange_getindex(a::BlockedUnitRange, indices::AbstractUnitRange)
+ first_block = block(to_blockindex(a, first(indices)))
+ last_block = block(to_blockindex(a, last(indices)))
+ lasts = [blocklasts(a)[Int(first_block):(Int(last_block) - 1)]; last(indices)]
+ return BlockArrays._BlockedUnitRange(first(indices), lasts)
+end
+
+function _BlockIndices end
+
+struct BlockIndices{N,R<:NTuple{N,AbstractUnitRange{Int}}} <: AbstractBlockArray{BlockIndex{N},N}
+ first::NTuple{N,Int}
+ indices::R
+ global function _BlockIndices(first::NTuple{N,Int}, indices::NTuple{N,AbstractUnitRange{Int}}) where {N}
+ return new{N,typeof(indices)}(first, indices)
+ end
+end
+function Base.axes(a::BlockIndices)
+ return map(Base.axes1, blockedunitrange_getindex.(a.indices, range.(a.first, last.(a.indices))))
+end
+function BlockIndices(indices::Tuple{Vararg{AbstractUnitRange{Int}}})
+ first = ntuple(Returns(1), length(indices))
+ return _BlockIndices(first, indices)
+end
+BlockIndices(a::AbstractArray) = BlockIndices(axes(a))
+
function Base.getindex(a::BlockIndices{N}, index::Vararg{Integer,N}) where {N}
- return BlockIndex(to_blockindex.(a.indices, index))
+ return BlockIndex(to_blockindex.(a.indices, index .+ a.first .- 1))
end
function Base.view(a::BlockIndices{N}, block::Block{N}) where {N}
return viewblock(a, block)
end
-function Base.view(a::BlockIndices, block::Vararg{Block{1}})
- return view(a, Block(block))
+function Base.view(a::BlockIndices{1}, block::Block{1})
+ return viewblock(a, block)
+end
+
+function Base.view(a::BlockIndices{N}, block::Vararg{Block{1},N}) where {N}
+ return view(a, Block(block))
end
function viewblock(a::BlockIndices, block)
range = Base.OneTo.(getindex.(blocklengths.(axes(a)), Int.(Tuple(block))))
return block[range...]
end
+
+function Base.view(a::BlockIndices{N}, indices::Vararg{BlockIndexRange{1},N}) where {N}
+ return view(a, BlockIndexRange(Block(block.(indices)), only.(getfield.(indices, :indices))))
+end
+
+function Base.view(a::BlockIndices{N}, indices::BlockIndexRange{N}) where {N}
+ a_block = a[block(indices)]
+ return block(a_block)[getindex.(a_block.indices, indices.indices)...]
+end
+
+# Circumvent that this is getting hijacked to call `a[block(indices)][indices.indices...]`,
+# which hits the bug https://github.com/JuliaArrays/BlockArrays.jl/issues/355.
+function Base.getindex(a::BlockIndices{N}, indices::BlockIndexRange{N}) where {N}
+ return view(a, indices)
+end
+
+function Base.view(a::BlockIndices{N}, indices::Vararg{AbstractUnitRange,N}) where {N}
+ offsets = a.first .- ntuple(Returns(1), Val(N))
+ firsts = first.(indices) .+ offsets
+ inds = blockedunitrange_getindex.(a.indices, Base.OneTo.(last.(indices) .+ offsets))
+ return _BlockIndices(firsts, inds)
+end
+
+# Ranges that result in contiguous slices, and therefore preserve `BlockIndices`.
+const BlockOrUnitRanges = Union{AbstractUnitRange,CartesianIndices{1},Block{1},BlockRange{1},BlockIndexRange{1}}
+
+function Base.view(a::BlockIndices{N}, indices::Vararg{BlockOrUnitRanges,N}) where {N}
+ return view(a, to_indices(a, indices)...)
+end
+
+function Base.view(a::BlockIndices{N}, indices::CartesianIndices{N}) where {N}
+ return view(a, to_indices(a, (indices,))...)
+end
+
+# For some reason this doesn't call `view` automatically.
+function Base.getindex(a::BlockIndices{N}, indices::CartesianIndices{N}) where {N}
+ return view(a, to_indices(a, (indices,))...)
+end
+
+function Base.view(a::BlockIndices{N}, indices::BlockRange{N}) where {N}
+ return view(a, to_indices(a, (indices,))...)
+end
+
+# For some reason this doesn't call `view` automatically.
+function Base.getindex(a::BlockIndices{N}, indices::BlockRange{N}) where {N}
+ return view(a, to_indices(a, (indices,))...)
+end
diff --git a/test/test_blockindices.jl b/test/test_blockindices.jl
index 2cbdbe3a..6a7523b0 100644
--- a/test/test_blockindices.jl
+++ b/test/test_blockindices.jl
@@ -449,12 +449,13 @@ end
@testset "BlockIndices" begin
v = Array(reshape(1:35, (5, 7)))
- A = BlockArray(v, [2,3], [3,4])
+ A = BlockArray(v, [2, 3], [3, 4])
B = BlockIndices(A)
@test B == BlockIndices((blockedrange([2,3]), blockedrange([3,4])))
@test eltype(B) === BlockIndex{2}
@test size(B) == (5, 7)
@test axes(B) == (1:5, 1:7)
+ @test blocklengths.(axes(B)) == ([2, 3], [3, 4])
@test blocksize(B) == (2, 2)
@test blockaxes(B) == (Block(1):Block(2), Block(1):Block(2))
@test B[1, 1] == Block(1, 1)[1, 1]
@@ -475,17 +476,38 @@ end
@test view(B, Block(1), Block(2)) isa BlockIndexRange{2}
@test view(B, Block(1), Block(2)) == Block(1, 2)[1:2, 1:4]
- r1 = BlockArrays._BlockedUnitRange(2, [2, 3])
- r2 = BlockArrays._BlockedUnitRange(2, [3, 4])
- @test B[2:3, 2:4] == BlockIndices((r1, r2))
- @test_broken B[2:3, 2:4] isa BlockIndices{2}
+ B23_24 = mortar([[Block(1, 1)[2:2, 2:3]] [Block(1, 2)[2:2, 1:1]]
+ [Block(2, 1)[1:1, 2:3]] [Block(2, 2)[1:1, 1:1]]])
+ @test B[2:3, 2:4] == B23_24
+ @test B[2:3, 2:4] isa BlockIndices{2}
+ @test B[CartesianIndices((2:3,)), CartesianIndices((2:4,))] == B23_24
+ @test B[CartesianIndices((2:3,)), CartesianIndices((2:4,))] isa BlockIndices{2}
+ @test B[CartesianIndices((2:3, 2:4))] == B23_24
+ @test B[CartesianIndices((2:3, 2:4))] isa BlockIndices{2}
@test B[Block(2, 2)[2:3, 2:3]] == Block(2, 2)[2:3, 2:3]
- @test_broken B[Block(2, 2)[2:3, 2:3]] isa BlockedIndexRange{2}
+ @test B[Block(2, 2)[2:3, 2:3]] isa BlockIndexRange{2}
@test B[Block(1):Block(2), Block(1):Block(2)] == B
- @test_broken B[Block(1):Block(2), Block(1):Block(2)] isa BlockedIndices{2}
+ @test B[Block(1):Block(2), Block(1):Block(2)] isa BlockIndices{2}
+
+ @test B[BlockRange(1:2, 1:2)] == B
+ @test B[BlockRange(1:2, 1:2)] isa BlockIndices{2}
@test B[Block(2), Block(1):Block(2)] == mortar([[Block(2, 1)[1:3, 1:3]] [Block(2, 2)[1:3, 1:4]]])
- @test_broken B[Block(2), Block(1):Block(2)] isa BlockedIndices{2}
+ @test B[Block(2), Block(1):Block(2)] isa BlockIndices{2}
+
+ @test B[2:4, 2:5][2:3, 2:3] == mortar([[Block(2, 1)[1:2, 3:3]] [Block(2, 2)[1:2, 1:1]]])
+ @test B[2:4, 2:5][2:3, 2:3] isa BlockIndices{2}
+
+ B = BlockIndices(blockedrange([2, 3]))
+ @test B == [Block(1)[1], Block(1)[2], Block(2)[1], Block(2)[2], Block(2)[3]]
+ @test blocklengths(only(axes(B))) == [2, 3]
+ @test B[Block(1)] == Block(1)[1:2]
+ @test B[Block(1)] isa BlockIndexRange{1}
+ @test B[Block(2)] == Block(2)[1:3]
+ @test B[Block(2)] isa BlockIndexRange{1}
+ @test B[2:4] == [Block(1)[2], Block(2)[1], Block(2)[2]]
+ @test blocklengths(only(axes(B[2:4]))) == [1, 2]
+ @test B[2:4] isa BlockIndices{1}
end
From 027130e3a226c16b755b729bb210839c7ab7a6ce Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Wed, 27 Mar 2024 12:00:11 -0400
Subject: [PATCH 4/8] Julia 1.6 compatibility
---
src/block_indices.jl | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/src/block_indices.jl b/src/block_indices.jl
index a10f908c..99c07552 100644
--- a/src/block_indices.jl
+++ b/src/block_indices.jl
@@ -25,8 +25,8 @@ end
function Base.axes(a::BlockIndices)
return map(Base.axes1, blockedunitrange_getindex.(a.indices, range.(a.first, last.(a.indices))))
end
-function BlockIndices(indices::Tuple{Vararg{AbstractUnitRange{Int}}})
- first = ntuple(Returns(1), length(indices))
+function BlockIndices(indices::Tuple{Vararg{AbstractUnitRange{Int},N}}) where {N}
+ first = ntuple(_ -> 1, Val(N))
return _BlockIndices(first, indices)
end
BlockIndices(a::AbstractArray) = BlockIndices(axes(a))
@@ -68,7 +68,7 @@ function Base.getindex(a::BlockIndices{N}, indices::BlockIndexRange{N}) where {N
end
function Base.view(a::BlockIndices{N}, indices::Vararg{AbstractUnitRange,N}) where {N}
- offsets = a.first .- ntuple(Returns(1), Val(N))
+ offsets = a.first .- ntuple(_ -> 1, Val(N))
firsts = first.(indices) .+ offsets
inds = blockedunitrange_getindex.(a.indices, Base.OneTo.(last.(indices) .+ offsets))
return _BlockIndices(firsts, inds)
From d17f762f1cee8033c7302f186455bddc35112beb Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Wed, 27 Mar 2024 12:14:51 -0400
Subject: [PATCH 5/8] Another fix for Julia 1.6
---
src/block_indices.jl | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/src/block_indices.jl b/src/block_indices.jl
index 99c07552..1ce351ef 100644
--- a/src/block_indices.jl
+++ b/src/block_indices.jl
@@ -23,7 +23,7 @@ struct BlockIndices{N,R<:NTuple{N,AbstractUnitRange{Int}}} <: AbstractBlockArray
end
end
function Base.axes(a::BlockIndices)
- return map(Base.axes1, blockedunitrange_getindex.(a.indices, range.(a.first, last.(a.indices))))
+ return map(Base.axes1, blockedunitrange_getindex.(a.indices, (:).(a.first, last.(a.indices))))
end
function BlockIndices(indices::Tuple{Vararg{AbstractUnitRange{Int},N}}) where {N}
first = ntuple(_ -> 1, Val(N))
From 328cce99c9778a991a168c37d81ff795b236669e Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Thu, 28 Mar 2024 17:57:11 -0400
Subject: [PATCH 6/8] Axes of slices match axes of input indices
---
src/block_indices.jl | 24 ++++++++++++++----------
test/test_blockindices.jl | 35 +++++++++++++++++++++++------------
2 files changed, 37 insertions(+), 22 deletions(-)
diff --git a/src/block_indices.jl b/src/block_indices.jl
index 1ce351ef..a2ad0b61 100644
--- a/src/block_indices.jl
+++ b/src/block_indices.jl
@@ -6,28 +6,32 @@ function to_blockindex(a::AbstractUnitRange, index::Integer)
return axis_blockindex
end
-function blockedunitrange_getindex(a::BlockedUnitRange, indices::AbstractUnitRange)
+# https://github.com/JuliaArrays/BlockArrays.jl/issues/347
+function blockedgetindex(a::BlockedUnitRange, indices::AbstractUnitRange)
first_block = block(to_blockindex(a, first(indices)))
last_block = block(to_blockindex(a, last(indices)))
lasts = [blocklasts(a)[Int(first_block):(Int(last_block) - 1)]; last(indices)]
return BlockArrays._BlockedUnitRange(first(indices), lasts)
end
+function blockedgetindex(a::AbstractUnitRange, indices::AbstractUnitRange)
+ return a[indices]
+end
function _BlockIndices end
-struct BlockIndices{N,R<:NTuple{N,AbstractUnitRange{Int}}} <: AbstractBlockArray{BlockIndex{N},N}
+struct BlockIndices{N,R<:NTuple{N,AbstractUnitRange{Int}},BS<:NTuple{N,AbstractUnitRange{Int}}} <: AbstractBlockArray{BlockIndex{N},N}
first::NTuple{N,Int}
indices::R
- global function _BlockIndices(first::NTuple{N,Int}, indices::NTuple{N,AbstractUnitRange{Int}}) where {N}
- return new{N,typeof(indices)}(first, indices)
+ axes::BS
+ global function _BlockIndices(first::NTuple{N,Int}, indices::R, axes::BS) where {N,R<:NTuple{N,AbstractUnitRange{Int}},BS<:NTuple{N,AbstractUnitRange{Int}}}
+ return new{N,R,BS}(first, indices, axes)
end
end
-function Base.axes(a::BlockIndices)
- return map(Base.axes1, blockedunitrange_getindex.(a.indices, (:).(a.first, last.(a.indices))))
-end
+Base.axes(a::BlockIndices) = a.axes
function BlockIndices(indices::Tuple{Vararg{AbstractUnitRange{Int},N}}) where {N}
first = ntuple(_ -> 1, Val(N))
- return _BlockIndices(first, indices)
+ axes = map(Base.axes1, blockedgetindex.(indices, Base.OneTo.(last.(indices))))
+ return _BlockIndices(first, indices, axes)
end
BlockIndices(a::AbstractArray) = BlockIndices(axes(a))
@@ -70,8 +74,8 @@ end
function Base.view(a::BlockIndices{N}, indices::Vararg{AbstractUnitRange,N}) where {N}
offsets = a.first .- ntuple(_ -> 1, Val(N))
firsts = first.(indices) .+ offsets
- inds = blockedunitrange_getindex.(a.indices, Base.OneTo.(last.(indices) .+ offsets))
- return _BlockIndices(firsts, inds)
+ inds = blockedgetindex.(a.indices, Base.OneTo.(last.(indices) .+ offsets))
+ return _BlockIndices(firsts, inds, Base.axes1.(indices))
end
# Ranges that result in contiguous slices, and therefore preserve `BlockIndices`.
diff --git a/test/test_blockindices.jl b/test/test_blockindices.jl
index 6a7523b0..ce6d955e 100644
--- a/test/test_blockindices.jl
+++ b/test/test_blockindices.jl
@@ -457,7 +457,7 @@ end
@test axes(B) == (1:5, 1:7)
@test blocklengths.(axes(B)) == ([2, 3], [3, 4])
@test blocksize(B) == (2, 2)
- @test blockaxes(B) == (Block(1):Block(2), Block(1):Block(2))
+ @test blockaxes(B) == (Block.(1:2), Block.(1:2))
@test B[1, 1] == Block(1, 1)[1, 1]
@test B[4, 6] == Block(2, 2)[2, 3]
@@ -478,24 +478,33 @@ end
B23_24 = mortar([[Block(1, 1)[2:2, 2:3]] [Block(1, 2)[2:2, 1:1]]
[Block(2, 1)[1:1, 2:3]] [Block(2, 2)[1:1, 1:1]]])
- @test B[2:3, 2:4] == B23_24
- @test B[2:3, 2:4] isa BlockIndices{2}
- @test B[CartesianIndices((2:3,)), CartesianIndices((2:4,))] == B23_24
- @test B[CartesianIndices((2:3,)), CartesianIndices((2:4,))] isa BlockIndices{2}
- @test B[CartesianIndices((2:3, 2:4))] == B23_24
- @test B[CartesianIndices((2:3, 2:4))] isa BlockIndices{2}
+ Br = B[2:3, 2:4]
+ @test Br == B23_24
+ @test blocksize(Br) == (1, 1)
+ @test Br isa AbstractMatrix{<:BlockIndex{2}}
+ @test Br isa BlockIndices{2}
+ Br = B[CartesianIndices((2:3,)), CartesianIndices((2:4,))]
+ @test Br == B23_24
+ @test blocksize(Br) == (1, 1)
+ @test Br isa AbstractMatrix{<:BlockIndex{2}}
+ @test Br isa BlockIndices{2}
+ Br = B[CartesianIndices((2:3, 2:4))]
+ @test Br == B23_24
+ @test Br isa AbstractMatrix{<:BlockIndex{2}}
+ @test Br isa BlockIndices{2}
@test B[Block(2, 2)[2:3, 2:3]] == Block(2, 2)[2:3, 2:3]
@test B[Block(2, 2)[2:3, 2:3]] isa BlockIndexRange{2}
- @test B[Block(1):Block(2), Block(1):Block(2)] == B
- @test B[Block(1):Block(2), Block(1):Block(2)] isa BlockIndices{2}
+ @test B[Block.(1:2), Block.(1:2)] == B
+ @test B[Block.(1:2), Block.(1:2)] isa BlockIndices{2}
@test B[BlockRange(1:2, 1:2)] == B
@test B[BlockRange(1:2, 1:2)] isa BlockIndices{2}
- @test B[Block(2), Block(1):Block(2)] == mortar([[Block(2, 1)[1:3, 1:3]] [Block(2, 2)[1:3, 1:4]]])
- @test B[Block(2), Block(1):Block(2)] isa BlockIndices{2}
+ # TODO: Should this make a `BlockIndices{1}`?
+ @test B[Block(2), Block.(1:2)] == mortar([[Block(2, 1)[1:3, 1:3]] [Block(2, 2)[1:3, 1:4]]])
+ @test B[Block(2), Block.(1:2)] isa BlockIndices{2}
@test B[2:4, 2:5][2:3, 2:3] == mortar([[Block(2, 1)[1:2, 3:3]] [Block(2, 2)[1:2, 1:1]]])
@test B[2:4, 2:5][2:3, 2:3] isa BlockIndices{2}
@@ -508,6 +517,8 @@ end
@test B[Block(2)] == Block(2)[1:3]
@test B[Block(2)] isa BlockIndexRange{1}
@test B[2:4] == [Block(1)[2], Block(2)[1], Block(2)[2]]
- @test blocklengths(only(axes(B[2:4]))) == [1, 2]
+ @test blocklength(only(axes(B[2:4]))) == 1
+ @test blocklengths(only(axes(B[2:4]))) == [3]
+ @test B[2:4] isa AbstractVector{<:BlockIndex{1}}
@test B[2:4] isa BlockIndices{1}
end
From fbe6c0bd3384e9bda19c7a64dc79ffbfc56782e0 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Thu, 28 Mar 2024 18:09:09 -0400
Subject: [PATCH 7/8] Fix test
---
test/test_blockindices.jl | 5 ++---
1 file changed, 2 insertions(+), 3 deletions(-)
diff --git a/test/test_blockindices.jl b/test/test_blockindices.jl
index ce6d955e..d1ca5481 100644
--- a/test/test_blockindices.jl
+++ b/test/test_blockindices.jl
@@ -502,9 +502,8 @@ end
@test B[BlockRange(1:2, 1:2)] == B
@test B[BlockRange(1:2, 1:2)] isa BlockIndices{2}
- # TODO: Should this make a `BlockIndices{1}`?
- @test B[Block(2), Block.(1:2)] == mortar([[Block(2, 1)[1:3, 1:3]] [Block(2, 2)[1:3, 1:4]]])
- @test B[Block(2), Block.(1:2)] isa BlockIndices{2}
+ @test B[Block.(2:2), Block.(1:2)] == mortar([[Block(2, 1)[1:3, 1:3]] [Block(2, 2)[1:3, 1:4]]])
+ @test B[Block.(2:2), Block.(1:2)] isa BlockIndices{2}
@test B[2:4, 2:5][2:3, 2:3] == mortar([[Block(2, 1)[1:2, 3:3]] [Block(2, 2)[1:2, 1:1]]])
@test B[2:4, 2:5][2:3, 2:3] isa BlockIndices{2}
From 43b459bddb5a9ea6d77e2f507e33f14d999523ce Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Fri, 29 Mar 2024 19:12:06 -0400
Subject: [PATCH 8/8] Add support for indexing with BlockIndices
---
src/block_indices.jl | 12 ++++++++++++
test/test_blockindices.jl | 21 +++++++++++++++++++++
2 files changed, 33 insertions(+)
diff --git a/src/block_indices.jl b/src/block_indices.jl
index a2ad0b61..813b7650 100644
--- a/src/block_indices.jl
+++ b/src/block_indices.jl
@@ -102,3 +102,15 @@ end
function Base.getindex(a::BlockIndices{N}, indices::BlockRange{N}) where {N}
return view(a, to_indices(a, (indices,))...)
end
+
+function Base.getindex(a::CartesianIndices{N}, indices::BlockIndices{N}) where {N}
+ # TODO: Is there a better way to write this?
+ new_axes = (:).(Tuple(a[first(indices)]), Tuple(a[last(indices)]))
+ firsts = first.(new_axes)
+ blocklasts = ntuple(i -> accumulate(+, blocklengths(axes(indices, i)); init=firsts[i] - one(firsts[i])), Val(N))
+ return CartesianIndices(_BlockedUnitRange.(firsts, blocklasts))
+end
+
+function Base.getindex(a::AbstractArray{<:Any,N}, indices::BlockIndices{N}) where {N}
+ return a[CartesianIndices(a)[indices]]
+end
diff --git a/test/test_blockindices.jl b/test/test_blockindices.jl
index d1ca5481..8bdc7088 100644
--- a/test/test_blockindices.jl
+++ b/test/test_blockindices.jl
@@ -520,4 +520,25 @@ end
@test blocklengths(only(axes(B[2:4]))) == [3]
@test B[2:4] isa AbstractVector{<:BlockIndex{1}}
@test B[2:4] isa BlockIndices{1}
+
+ A = BlockArray(randn(5, 5), [2, 3], [2, 3])
+ BA = BlockIndices(A)
+ r = BlockArrays._BlockedUnitRange(2, [2, 4])
+ B = BA[r, r]
+ @test B == BA[2:4, 2:4]
+ @test size(B) == (3, 3)
+ @test blocksize(B) == (2, 2)
+ @test blocklengths.(axes(B)) == ([1, 2], [1, 2])
+ CB = CartesianIndices(A)[B]
+ @test CB == CartesianIndices((2:4, 2:4))
+ @test size(CB) == (3, 3)
+ @test blocksize(CB) == (2, 2)
+ @test blocklengths.(axes(CB)) == ([1, 2], [1, 2])
+ @test all(blockisequal.(axes(CB), axes(B)))
+ AB = A[B]
+ @test AB == A[2:4, 2:4]
+ @test size(AB) == (3, 3)
+ @test blocksize(AB) == (2, 2)
+ @test blocklengths.(axes(AB)) == ([1, 2], [1, 2])
+ @test all(blockisequal.(axes(AB), axes(B)))
end