Skip to content

Commit 2555908

Browse files
dlfivefiftydlfivefifty
authored
Generalise A[::BlockIndex{1}, ::BlockIndex{1}] (#107)
* Generalise A[::BlockIndex{1}, ::BlockIndex{1}] * mixes of BlockIndex and BlockIndexRange
1 parent 5b8f5ec commit 2555908

File tree

7 files changed

+100
-92
lines changed

7 files changed

+100
-92
lines changed

‎Project.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
name = "BlockArrays"
22
uuid = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
3-
version = "0.12.1"
3+
version = "0.12.2"
44

55
[deps]
66
ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

‎src/blockarray.jl

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -331,9 +331,6 @@ end
331331
_blockindex_getindex(block_arr, blockindex)
332332
@inline Base.getindex(block_arr::BlockVector{T}, blockindex::BlockIndex{1}) where {T} =
333333
_blockindex_getindex(block_arr, blockindex)
334-
@inline Base.getindex(block_arr::BlockArray{T,N}, blockindex::Vararg{BlockIndex{1},N}) where {T,N} =
335-
block_arr[BlockIndex(blockindex)]
336-
337334

338335
###########################
339336
# AbstractArray Interface #

‎src/blockaxis.jl

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,15 @@
11

22
# interface
3-
getindex(b::AbstractVector, K::BlockIndex{1}) = b[Block(K.I[1])][K.α[1]]
4-
getindex(b::AbstractVector, K::BlockIndexRange{1}) = b[K.block][K.indices[1]]
3+
4+
@inline getindex(b::AbstractVector, K::BlockIndex{1}) = b[Block(K.I[1])][K.α[1]]
5+
@inline getindex(b::AbstractArray{T,N}, K::BlockIndex{N}) where {T,N} =
6+
b[block(K)][K.α...]
7+
@inline getindex(b::AbstractArray{T,N}, K::Vararg{BlockIndex{1},N}) where {T,N} =
8+
b[BlockIndex(K)]
9+
10+
@inline getindex(b::AbstractArray{T,N}, K::BlockIndexRange{N}) where {T,N} =
11+
b[block(K)][K.indices...]
12+
513

614
function findblockindex(b::AbstractVector, k::Integer)
715
K = findblock(b, k)

‎src/blockindices.jl

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,8 @@ BlockIndexRange(block::Block{N}, inds::NTuple{N,AbstractUnitRange{Int}}) where {
173173
BlockIndexRange(block::Block{N}, inds::Vararg{AbstractUnitRange{Int},N}) where {N} =
174174
BlockIndexRange(block,inds)
175175

176+
block(R::BlockIndexRange) = R.block
177+
176178
getindex(B::Block{N}, inds::Vararg{Int,N}) where N = BlockIndex(B,inds)
177179
getindex(B::Block{N}, inds::Vararg{AbstractUnitRange{Int},N}) where N = BlockIndexRange(B,inds)
178180
getindex(B::Block{1}, inds::Colon) = B

‎src/pseudo_blockarray.jl

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -169,13 +169,20 @@ end
169169
# Indexing #
170170
############
171171

172-
@inline function Base.getindex(block_arr::PseudoBlockArray{T,N}, blockindex::BlockIndex{N}) where {T,N}
172+
@inline function _pseudoblockindex_getindex(block_arr, blockindex)
173173
I = getindex.(axes(block_arr), getindex.(Block.(blockindex.I), blockindex.α))
174174
@boundscheck checkbounds(block_arr.blocks, I...)
175175
@inbounds v = block_arr.blocks[I...]
176176
return v
177177
end
178178

179+
@inline Base.getindex(block_arr::PseudoBlockArray{T,N}, blockindex::BlockIndex{N}) where {T,N} =
180+
_pseudoblockindex_getindex(block_arr, blockindex)
181+
182+
183+
@inline Base.getindex(block_arr::PseudoBlockVector{T}, blockindex::BlockIndex{1}) where T =
184+
_pseudoblockindex_getindex(block_arr, blockindex)
185+
179186
@inline function getblock(block_arr::PseudoBlockArray{T,N}, block::Vararg{Integer, N}) where {T,N}
180187
range = getindex.(axes(block_arr), Block.(block))
181188
return view(block_arr.blocks, range...)

‎src/views.jl

Lines changed: 68 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,51 @@
11
### Views
2+
3+
"""
4+
unblock(block_sizes, inds, I)
5+
6+
Returns the indices associated with a block as a `BlockSlice`.
7+
"""
8+
function unblock(A, inds, I)
9+
B = first(I)
10+
if length(inds) == 0
11+
# Allow `ones(2)[Block(1)[1:1], Block(1)[1:1]]` which is
12+
# similar to `ones(2)[1:1, 1:1]`.
13+
BlockSlice(B,Base.OneTo(1))
14+
else
15+
BlockSlice(B,inds[1][B])
16+
end
17+
end
18+
19+
to_index(::Block) = throw(ArgumentError("Block must be converted by to_indices(...)"))
20+
to_index(::BlockIndex) = throw(ArgumentError("BlockIndex must be converted by to_indices(...)"))
221
to_index(::BlockIndexRange) = throw(ArgumentError("BlockIndexRange must be converted by to_indices(...)"))
22+
to_index(::BlockRange) = throw(ArgumentError("BlockRange must be converted by to_indices(...)"))
323

24+
25+
@inline to_indices(A, inds, I::Tuple{Block{1}, Vararg{Any}}) =
26+
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
27+
@inline to_indices(A, inds, I::Tuple{BlockRange{1,R}, Vararg{Any}}) where R =
28+
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
29+
@inline to_indices(A, inds, I::Tuple{BlockIndex{1}, Vararg{Any}}) where R =
30+
(inds[1][I[1]], to_indices(A, _maybetail(inds), tail(I))...)
431
@inline to_indices(A, inds, I::Tuple{BlockIndexRange{1,R}, Vararg{Any}}) where R =
532
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
633

734
# splat out higher dimensional blocks
835
# this mimics view of a CartesianIndex
36+
@inline to_indices(A, inds, I::Tuple{Block, Vararg{Any}}) =
37+
to_indices(A, inds, (Block.(I[1].n)..., tail(I)...))
38+
@inline to_indices(A, inds, I::Tuple{BlockIndex, Vararg{Any}}) =
39+
to_indices(A, inds, (BlockIndex.(I[1].I, I[1].α)..., tail(I)...))
940
@inline to_indices(A, inds, I::Tuple{BlockIndexRange, Vararg{Any}}) =
10-
to_indices(A, inds, (BlockRange.(Block.(I[1].block.n), tuple.(I[1].indices))..., tail(I)...))
11-
41+
to_indices(A, inds, (BlockIndexRange.(Block.(I[1].block.n), tuple.(I[1].indices))..., tail(I)...))
42+
@inline to_indices(A, inds, I::Tuple{BlockRange, Vararg{Any}}) =
43+
to_indices(A, inds, (BlockRange.(tuple.(I[1].indices))..., tail(I)...))
1244

1345
# In 0.7, we need to override to_indices to avoid calling linearindices
14-
@inline to_indices(A, I::Tuple{BlockIndexRange, Vararg{Any}}) =
15-
to_indices(A, axes(A), I)
46+
@inline to_indices(A, I::Tuple{BlockIndexRange, Vararg{Any}}) = to_indices(A, axes(A), I)
47+
@inline to_indices(A, I::Tuple{Block, Vararg{Any}}) = to_indices(A, axes(A), I)
48+
@inline to_indices(A, I::Tuple{BlockRange, Vararg{Any}}) = to_indices(A, axes(A), I)
1649

1750
if VERSION >= v"1.2-" # See also `reindex` definitions in views.jl
1851
reindex(idxs::Tuple{BlockSlice{<:BlockRange}, Vararg{Any}},
@@ -68,96 +101,46 @@ end
68101
end
69102

70103

71-
"""
72-
unblock(block_sizes, inds, I)
73-
74-
Returns the indices associated with a block as a `BlockSlice`.
75-
"""
76-
function unblock(A, inds, I)
77-
B = first(I)
78-
if length(inds) == 0
79-
# Allow `ones(2)[Block(1)[1:1], Block(1)[1:1]]` which is
80-
# similar to `ones(2)[1:1, 1:1]`.
81-
BlockSlice(B,Base.OneTo(1))
82-
else
83-
BlockSlice(B,inds[1][B])
84-
end
85-
end
86-
87-
88-
to_index(::Block) = throw(ArgumentError("Block must be converted by to_indices(...)"))
89-
to_index(::BlockIndex) = throw(ArgumentError("BlockIndex must be converted by to_indices(...)"))
90-
to_index(::BlockRange) = throw(ArgumentError("BlockRange must be converted by to_indices(...)"))
91-
92-
@inline to_indices(A, inds, I::Tuple{Block{1}, Vararg{Any}}) =
93-
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
94-
95-
# splat out higher dimensional blocks
96-
# this mimics view of a CartesianIndex
97-
@inline to_indices(A, inds, I::Tuple{Block, Vararg{Any}}) =
98-
to_indices(A, inds, (Block.(I[1].n)..., tail(I)...))
99-
100-
@inline to_indices(A, inds, I::Tuple{BlockRange{1,R}, Vararg{Any}}) where R =
101-
(unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
102-
103-
# splat out higher dimensional blocks
104-
# this mimics view of a CartesianIndex
105-
@inline to_indices(A, inds, I::Tuple{BlockRange, Vararg{Any}}) =
106-
to_indices(A, inds, (BlockRange.(tuple.(I[1].indices))..., tail(I)...))
107-
108-
109-
# In 0.7, we need to override to_indices to avoid calling linearindices
110-
@inline to_indices(A, I::Tuple{Block, Vararg{Any}}) =
111-
to_indices(A, axes(A), I)
112-
113-
@inline to_indices(A, I::Tuple{BlockRange, Vararg{Any}}) =
114-
to_indices(A, axes(A), I)
115-
116-
117104
# The first argument for `reindex` is removed as of
118105
# https://github.com/JuliaLang/julia/pull/30789 in Julia `Base`. So,
119106
# we define 2-arg `reindex` for Julia 1.2 and later.
120107
if VERSION >= v"1.2-"
108+
# BlockSlices map the blocks and the indices
109+
# this is loosely based on Slice reindex in subarray.jl
110+
reindex(idxs::Tuple{BlockSlice{<:BlockRange}, Vararg{Any}},
111+
subidxs::Tuple{BlockSlice{<:BlockRange}, Vararg{Any}}) =
112+
(@_propagate_inbounds_meta; (BlockSlice(BlockRange(idxs[1].block.indices[1][Int.(subidxs[1].block)]),
113+
idxs[1].indices[subidxs[1].indices]),
114+
reindex(tail(idxs), tail(subidxs))...))
121115

122-
# BlockSlices map the blocks and the indices
123-
# this is loosely based on Slice reindex in subarray.jl
124-
reindex(idxs::Tuple{BlockSlice{<:BlockRange}, Vararg{Any}},
125-
subidxs::Tuple{BlockSlice{<:BlockRange}, Vararg{Any}}) =
126-
(@_propagate_inbounds_meta; (BlockSlice(BlockRange(idxs[1].block.indices[1][Int.(subidxs[1].block)]),
127-
idxs[1].indices[subidxs[1].indices]),
128-
reindex(tail(idxs), tail(subidxs))...))
129-
130-
reindex(idxs::Tuple{BlockSlice{BlockRange{1,Tuple{UnitRange{Int}}}}, Vararg{Any}},
131-
subidxs::Tuple{BlockSlice{Block{1,Int}}, Vararg{Any}}) =
132-
(@_propagate_inbounds_meta; (BlockSlice(Block(idxs[1].block.indices[1][Int(subidxs[1].block)]),
133-
idxs[1].indices[subidxs[1].indices]),
134-
reindex(tail(idxs), tail(subidxs))...))
135-
136-
function reindex(idxs::Tuple{BlockSlice{Block{1,Int}}, Vararg{Any}},
137-
subidxs::Tuple{BlockSlice{Block{1,Int}}, Vararg{Any}})
138-
(idxs[1], reindex(tail(idxs), tail(subidxs))...)
139-
end
116+
reindex(idxs::Tuple{BlockSlice{BlockRange{1,Tuple{UnitRange{Int}}}}, Vararg{Any}},
117+
subidxs::Tuple{BlockSlice{Block{1,Int}}, Vararg{Any}}) =
118+
(@_propagate_inbounds_meta; (BlockSlice(Block(idxs[1].block.indices[1][Int(subidxs[1].block)]),
119+
idxs[1].indices[subidxs[1].indices]),
120+
reindex(tail(idxs), tail(subidxs))...))
140121

122+
function reindex(idxs::Tuple{BlockSlice{Block{1,Int}}, Vararg{Any}},
123+
subidxs::Tuple{BlockSlice{Block{1,Int}}, Vararg{Any}})
124+
(idxs[1], reindex(tail(idxs), tail(subidxs))...)
125+
end
141126
else # if VERSION >= v"1.2-"
127+
reindex(V, idxs::Tuple{BlockSlice{<:BlockRange}, Vararg{Any}},
128+
subidxs::Tuple{BlockSlice{<:BlockRange}, Vararg{Any}}) =
129+
(@_propagate_inbounds_meta; (BlockSlice(BlockRange(idxs[1].block.indices[1][Int.(subidxs[1].block)]),
130+
idxs[1].indices[subidxs[1].indices]),
131+
reindex(V, tail(idxs), tail(subidxs))...))
142132

143-
reindex(V, idxs::Tuple{BlockSlice{<:BlockRange}, Vararg{Any}},
144-
subidxs::Tuple{BlockSlice{<:BlockRange}, Vararg{Any}}) =
145-
(@_propagate_inbounds_meta; (BlockSlice(BlockRange(idxs[1].block.indices[1][Int.(subidxs[1].block)]),
146-
idxs[1].indices[subidxs[1].indices]),
147-
reindex(V, tail(idxs), tail(subidxs))...))
148-
149-
reindex(V, idxs::Tuple{BlockSlice{BlockRange{1,Tuple{UnitRange{Int}}}}, Vararg{Any}},
150-
subidxs::Tuple{BlockSlice{Block{1,Int}}, Vararg{Any}}) =
151-
(@_propagate_inbounds_meta; (BlockSlice(Block(idxs[1].block.indices[1][Int(subidxs[1].block)]),
152-
idxs[1].indices[subidxs[1].indices]),
153-
reindex(V, tail(idxs), tail(subidxs))...))
154-
155-
function reindex(V, idxs::Tuple{BlockSlice{Block{1,Int}}, Vararg{Any}},
156-
subidxs::Tuple{BlockSlice{Block{1,Int}}, Vararg{Any}})
157-
subidxs[1].block == Block(1) || throw(BoundsError(V, subidxs[1].block))
158-
(idxs[1], reindex(V, tail(idxs), tail(subidxs))...)
159-
end
133+
reindex(V, idxs::Tuple{BlockSlice{BlockRange{1,Tuple{UnitRange{Int}}}}, Vararg{Any}},
134+
subidxs::Tuple{BlockSlice{Block{1,Int}}, Vararg{Any}}) =
135+
(@_propagate_inbounds_meta; (BlockSlice(Block(idxs[1].block.indices[1][Int(subidxs[1].block)]),
136+
idxs[1].indices[subidxs[1].indices]),
137+
reindex(V, tail(idxs), tail(subidxs))...))
160138

139+
function reindex(V, idxs::Tuple{BlockSlice{Block{1,Int}}, Vararg{Any}},
140+
subidxs::Tuple{BlockSlice{Block{1,Int}}, Vararg{Any}})
141+
subidxs[1].block == Block(1) || throw(BoundsError(V, subidxs[1].block))
142+
(idxs[1], reindex(V, tail(idxs), tail(subidxs))...)
143+
end
161144
end # if VERSION >= v"1.2-"
162145

163146

‎test/test_blockarrays.jl

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -468,4 +468,15 @@ end
468468
@test A*b isa PseudoBlockVector{Float64}
469469
@test Ã*b isa PseudoBlockVector{Float64}
470470
@test A*b ≈ Ã*b ≈ Matrix(A)*b
471+
end
472+
473+
@testset "Blockindex" begin
474+
a = PseudoBlockArray(randn(3), [1,2])
475+
@test a[Block(1)[1]] == a[1]
476+
@test a[Block(1)[1:1]] == a[1:1]
477+
A = PseudoBlockArray(randn(3,3), [1,2], [1,2])
478+
@test A[Block(1)[1], Block(1)[1]] == A[Block(1,1)[1,1]] == A[1,1]
479+
@test A[Block(1)[1:1], Block(1)[1:1]] == A[Block(1,1)[1:1,1:1]] == A[1:1,1:1]
480+
@test A[Block(1)[1:1], Block(1)[1]] == BlockArray(A)[Block(1)[1:1], Block(1)[1]] == A[1:1,1]
481+
@test A[Block(1)[1], Block(1)[1:1]] == BlockArray(A)[Block(1)[1], Block(1)[1:1]] == A[1,1:1]
471482
end

0 commit comments

Comments
 (0)