Support views involving vectors of blocks, Fixes #184 and #358 (#445)

JuliaArrays · Feb 6, 2025 · a4e544e · a4e544e · dlfivefifty · Feb 6, 2025
1 parent 6e54660
commit a4e544e
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 4 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "BlockArrays"
 uuid = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
-version = "1.3.0"
+version = "1.4.0"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

diff --git a/src/blockaxis.jl b/src/blockaxis.jl
@@ -397,6 +397,11 @@ end
 
 @propagate_inbounds getindex(b::AbstractBlockedUnitRange, KR::BlockSlice) = b[KR.block]
 
+getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:Block{1}}) = mortar([b[K] for K in KR])
+getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:BlockIndexRange{1}}) = mortar([b[K] for K in KR])
+getindex(b::AbstractBlockedUnitRange, KR::AbstractVector{<:BlockIndex{1}}) = [b[K] for K in KR]
+getindex(b::AbstractBlockedUnitRange, Kkr::BlockIndexRange{1}) =  b[block(Kkr)][Kkr.indices...]
+
 _searchsortedfirst(a::AbstractVector, k) = searchsortedfirst(a, k)
 function _searchsortedfirst(a::Tuple, k)
     k ≤ first(a) && return 1

diff --git a/src/blockindices.jl b/src/blockindices.jl
@@ -252,9 +252,9 @@ Block(bs::BlockIndexRange) = bs.block
 
 
 """
-    BlockSlice(indices)
+    BlockSlice(block, indices)
 
-Represent an AbstractUnitRange of indices that attaches a block.
+Represent an AbstractUnitRange{<:Integer} of indices that attaches a block.
 
 Upon calling `to_indices()`, Blocks are converted to BlockSlice objects to represent
 the indices over which the Block spans.

diff --git a/src/views.jl b/src/views.jl
@@ -7,8 +7,12 @@ Returns the indices associated with a block as a `BlockSlice`.
 """
 function unblock(A, inds, I)
     B = first(I)
-    BlockSlice(B,inds[1][B])
+    _blockslice(B, inds[1][B])
 end
+
+_blockslice(B, a::AbstractUnitRange) = BlockSlice(B, a)
+_blockslice(B, a) = a # drop block structure
+
 # Allow `ones(2)[Block(1)[1:1], Block(1)[1:1]]` which is
 # similar to `ones(2)[1:1, 1:1]`.
 # Need to check the length of I in case its empty
@@ -28,6 +32,13 @@ to_index(::BlockRange) = throw(ArgumentError("BlockRange must be converted by to
     (inds[1][I[1]], to_indices(A, _maybetail(inds), tail(I))...)
 @inline to_indices(A, inds, I::Tuple{BlockIndexRange{1,R}, Vararg{Any}}) where R =
     (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
+@inline to_indices(A, inds, I::Tuple{AbstractVector{Block{1,R}}, Vararg{Any}}) where R =
+    (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
+@inline to_indices(A, inds, I::Tuple{AbstractVector{<:BlockIndex{1}}, Vararg{Any}}) =
+    (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
+@inline to_indices(A, inds, I::Tuple{AbstractVector{<:BlockIndexRange{1}}, Vararg{Any}}) =
+    (unblock(A, inds, I), to_indices(A, _maybetail(inds), tail(I))...)
+
 
 # splat out higher dimensional blocks
 # this mimics view of a CartesianIndex
@@ -44,6 +55,9 @@ to_index(::BlockRange) = throw(ArgumentError("BlockRange must be converted by to
 @inline to_indices(A, I::Tuple{BlockIndexRange, Vararg{Any}}) = to_indices(A, axes(A), I)
 @inline to_indices(A, I::Tuple{Block, Vararg{Any}}) = to_indices(A, axes(A), I)
 @inline to_indices(A, I::Tuple{BlockRange, Vararg{Any}}) = to_indices(A, axes(A), I)
+@inline to_indices(A, I::Tuple{AbstractVector{<:Block{1}}, Vararg{Any}}) = to_indices(A, axes(A), I)
+@inline to_indices(A, I::Tuple{AbstractVector{<:BlockIndex{1}}, Vararg{Any}}) = to_indices(A, axes(A), I)
+@inline to_indices(A, I::Tuple{AbstractVector{<:BlockIndexRange{1}}, Vararg{Any}}) = to_indices(A, axes(A), I)
 
 @propagate_inbounds reindex(idxs::Tuple{BlockSlice{<:BlockRange}, Vararg{Any}},
         subidxs::Tuple{BlockSlice{<:BlockIndexRange}, Vararg{Any}}) =

diff --git a/test/test_blockarrays.jl b/test/test_blockarrays.jl
@@ -924,6 +924,21 @@ end
         @test_throws BoundsError a[:,Block.(1:2)]
         @test size(a[:,1]) == (8,)
     end
+
+    @testset "Block-vector indexing (#184)" begin
+        a = BlockArray(1:6, [1,2,2,1])
+        @test a[Block.(2:3)] == a[collect(Block.(2:3))]
+
+        for B in (BlockArray(I, fill(2,4), fill(2,5)), BlockedArray(I, fill(2,4), fill(2,5)))
+            @test B[[Block(1),Block(2)], [Block(1),Block(2)]] == view(B, [Block(1),Block(2)], [Block(1), Block(2)]) == view(B, Block.(1:2), [Block(1), Block(2)])  == view(B, Block.(1:2), Block.(1:2))
+            @test B[[Block(1),Block(3)], [Block(1), Block(5)]] == view(B, [Block(1),Block(3)], [Block(1), Block(5)])
+        end
+    end
+    @testset "BlockIndex-vector indexing (#358)" begin
+        a = BlockArray(randn(6, 6), [3, 3], [3, 3])
+        @test a[[Block(1)[1:2], Block(2)[1:2]], [Block(1)[1:2], Block(2)[1:2]]] == [a[Block(1,1)[1:2,1:2]] a[Block(1,2)[1:2,1:2]]; a[Block(2,1)[1:2,1:2]] a[Block(2,2)[1:2,1:2]]]
+        @test a[[Block(1)[1], Block(2)[2]], [Block(1)[1:2], Block(2)[1:2]]] == [a[Block(1)[1],Block(1)[1:2]]' a[Block(1)[1], Block(2)[1:2]]'; a[Block(2)[2],Block(1)[1:2]]' a[Block(2)[2], Block(2)[1:2]]']
+    end
 end
 
 end # module
diff --git a/test/test_blockindices.jl b/test/test_blockindices.jl
@@ -888,4 +888,9 @@ end
     @test !blockisequal(axv, axB)
 end
 
+@testset "BlockIndices" begin
+    a = BlockedOneTo(1:3)
+    @test a[[Block(1),Block(3)]] == a[Block.(1:2:3)] == [1,3]
+end
+
 end # module