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..813b7650 --- /dev/null +++ b/src/block_indices.jl @@ -0,0 +1,116 @@ +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 + +# 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}},BS<:NTuple{N,AbstractUnitRange{Int}}} <: AbstractBlockArray{BlockIndex{N},N} + first::NTuple{N,Int} + indices::R + 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 +Base.axes(a::BlockIndices) = a.axes +function BlockIndices(indices::Tuple{Vararg{AbstractUnitRange{Int},N}}) where {N} + first = ntuple(_ -> 1, Val(N)) + axes = map(Base.axes1, blockedgetindex.(indices, Base.OneTo.(last.(indices)))) + return _BlockIndices(first, indices, axes) +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 .+ 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{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(_ -> 1, Val(N)) + firsts = first.(indices) .+ offsets + 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`. +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 + +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 f3beadf2..8bdc7088 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,99 @@ 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 blocklengths.(axes(B)) == ([2, 3], [3, 4]) + @test blocksize(B) == (2, 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] + + @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 BlockIndexRange{2} + @test view(B, Block(1, 2)) == Block(1, 2)[1:2, 1:4] + + @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 BlockIndexRange{2} + @test view(B, Block(1), Block(2)) == Block(1, 2)[1:2, 1:4] + + 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]]]) + 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: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: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} + + 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 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} + + 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