Generalize blockwise slicing

Project.toml

@@ -1,7 +1,7 @@
 name = "BlockSparseArrays"
 uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.7.20"
+version = "0.7.21"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -167,8 +167,14 @@ const BlockIndexRangeSlices = BlockIndices{
 const BlockIndexVectorSlices = BlockIndices{
   <:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexVector}}
 }
+const GenericBlockIndexVectorSlices = BlockIndices{
+  <:BlockVector{<:GenericBlockIndex{1},<:Vector{<:BlockIndexVector}}
+}
 const SubBlockSliceCollection = Union{
-  BlockIndexRangeSlice,BlockIndexRangeSlices,BlockIndexVectorSlices
+  BlockIndexRangeSlice,
+  BlockIndexRangeSlices,
+  BlockIndexVectorSlices,
+  GenericBlockIndexVectorSlices,
 }
 
 # TODO: This is type piracy. This is used in `reindex` when making
@@ -392,6 +398,13 @@ function blockrange(
   return map(Block, blocks(r))
 end
 
+function blockrange(
+  axis::AbstractUnitRange,
+  r::BlockVector{<:GenericBlockIndex{1},<:AbstractVector{<:BlockIndexVector}},
+)
+  return map(Block, blocks(r))
+end
+
 function blockrange(axis::AbstractUnitRange, r)
   return error("Slicing not implemented for range of type `$(typeof(r))`.")
 end

src/BlockArraysExtensions/blockedunitrange.jl

@@ -244,6 +244,15 @@ BlockArrays.blockindex(b::GenericBlockIndex{1}) = b.α[1]
 function GenericBlockIndex(indcs::Tuple{Vararg{GenericBlockIndex{1},N}}) where {N}
   GenericBlockIndex(block.(indcs), blockindex.(indcs))
 end
+
+function Base.checkindex(
+  ::Type{Bool}, axis::AbstractBlockedUnitRange, ind::GenericBlockIndex{1}
+)
+  return checkindex(Bool, axis, block(ind)) &&
+         checkbounds(Bool, axis[block(ind)], blockindex(ind))
+end
+Base.to_index(i::GenericBlockIndex) = i
+
 function print_tuple_elements(io::IO, @nospecialize(t))
   if !isempty(t)
     print(io, t[1])
@@ -261,6 +270,13 @@ function Base.show(io::IO, B::GenericBlockIndex)
   return nothing
 end
 
+# https://github.com/JuliaArrays/BlockArrays.jl/blob/v1.6.3/src/views.jl#L31-L32
+_maybetail(::Tuple{}) = ()
+_maybetail(t::Tuple) = Base.tail(t)
+@inline function Base.to_indices(A, inds, I::Tuple{GenericBlockIndex{1},Vararg{Any}})
+  return (inds[1][I[1]], to_indices(A, _maybetail(inds), Base.tail(I))...)
+end
+
 using Base: @propagate_inbounds
 @propagate_inbounds function Base.getindex(b::AbstractVector, K::GenericBlockIndex{1})
   return b[Block(K.I[1])][K.α[1]]
@@ -316,6 +332,19 @@ using ArrayLayouts: LayoutArray
   K
 )][K.indices...]
 
+function blockedunitrange_getindices(
+  a::AbstractBlockedUnitRange,
+  indices::BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexVector{1}}},
+)
+  return mortar(map(b -> a[b], blocks(indices)))
+end
+function blockedunitrange_getindices(
+  a::AbstractBlockedUnitRange,
+  indices::BlockVector{<:GenericBlockIndex{1},<:Vector{<:BlockIndexVector{1}}},
+)
+  return mortar(map(b -> a[b], blocks(indices)))
+end
+
 function to_blockindices(a::AbstractBlockedUnitRange{<:Integer}, I::AbstractArray{Bool})
   I_blocks = blocks(BlockedVector(I, blocklengths(a)))
   I′_blocks = map(eachindex(I_blocks)) do b

src/abstractblocksparsearray/views.jl

@@ -166,7 +166,7 @@ function BlockArrays.viewblock(
     <:AbstractBlockSparseArray{T,N},
     <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}},
   },
-  block::Union{Block{N},BlockIndexRange{N}},
+  block::Union{Block{N},BlockIndexRange{N},BlockIndexVector{N}},
 ) where {T,N}
   return viewblock(a, to_tuple(block)...)
 end
@@ -228,6 +228,14 @@ function to_blockindexrange(
   # work right now.
   return blocks(a.blocks)[Int(I)]
 end
+function to_blockindexrange(
+  a::BlockIndices{<:BlockVector{<:GenericBlockIndex{1},<:Vector{<:BlockIndexVector}}},
+  I::Block{1},
+)
+  # TODO: Ideally we would just use `a.blocks[I]` but that doesn't
+  # work right now.
+  return blocks(a.blocks)[Int(I)]
+end
 function to_blockindexrange(
   a::Base.Slice{<:AbstractBlockedUnitRange{<:Integer}}, I::Block{1}
 )

src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -78,13 +78,36 @@ function Base.to_indices(
   return @interface interface(a) to_indices(a, inds, I)
 end
 
+# a[mortar([Block(1)[[1, 2]], Block(2)[[1, 3]]])]
+function Base.to_indices(
+  a::AnyAbstractBlockSparseArray,
+  inds,
+  I::Tuple{BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexVector{1}}},Vararg{Any}},
+)
+  return @interface interface(a) to_indices(a, inds, I)
+end
+function Base.to_indices(
+  a::AnyAbstractBlockSparseArray,
+  inds,
+  I::Tuple{BlockVector{<:GenericBlockIndex{1},<:Vector{<:BlockIndexVector{1}}},Vararg{Any}},
+)
+  return @interface interface(a) to_indices(a, inds, I)
+end
+
 # a[[Block(1)[1:2], Block(2)[1:2]], [Block(1)[1:2], Block(2)[1:2]]]
 function Base.to_indices(
   a::AnyAbstractBlockSparseArray, inds, I::Tuple{Vector{<:BlockIndexRange{1}},Vararg{Any}}
 )
   return to_indices(a, inds, (mortar(I[1]), Base.tail(I)...))
 end
 
+# a[[Block(1)[[1, 2]], Block(2)[[1, 2]]], [Block(1)[[1, 2]], Block(2)[[1, 2]]]]
+function Base.to_indices(
+  a::AnyAbstractBlockSparseArray, inds, I::Tuple{Vector{<:BlockIndexVector{1}},Vararg{Any}}
+)
+  return to_indices(a, inds, (mortar(I[1]), Base.tail(I)...))
+end
+
 # BlockArrays `AbstractBlockArray` interface
 function BlockArrays.blocks(a::AnyAbstractBlockSparseArray)
   @interface interface(a) blocks(a)

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -229,6 +229,23 @@ end
   return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...)
 end
 
+@interface ::AbstractBlockSparseArrayInterface function Base.to_indices(
+  a,
+  inds,
+  I::Tuple{BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexVector{1}}},Vararg{Any}},
+)
+  I1 = BlockIndices(I[1], blockedunitrange_getindices(inds[1], I[1]))
+  return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...)
+end
+@interface ::AbstractBlockSparseArrayInterface function Base.to_indices(
+  a,
+  inds,
+  I::Tuple{BlockVector{<:GenericBlockIndex{1},<:Vector{<:BlockIndexVector{1}}},Vararg{Any}},
+)
+  I1 = BlockIndices(I[1], blockedunitrange_getindices(inds[1], I[1]))
+  return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...)
+end
+
 # a[BlockVector([Block(2), Block(1)], [2]), BlockVector([Block(2), Block(1)], [2])]
 # Permute and merge blocks.
 # TODO: This isn't merging blocks yet, that needs to be implemented that.