[WIP] Logical indexing

Author: mtfishman

Project.toml

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

src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -67,12 +67,19 @@ for f in (:axes, :unsafe_indices, :axes1, :first, :last, :size, :length, :unsafe
   @eval Base.$f(S::BlockIndices) = Base.$f(S.indices)
 end
 Base.getindex(S::BlockIndices, i::Integer) = getindex(S.indices, i)
+
+function _blockslice(x, y::AbstractUnitRange{<:Integer})
+  return BlockSlice(x, y)
+end
+function _blockslice(x, y::AbstractVector{<:Integer})
+  return BlockIndices(x, y)
+end
 function Base.getindex(S::BlockIndices, i::BlockSlice{<:Block{1}})
   # TODO: Check that `i.indices` is consistent with `S.indices`.
   # It seems like this isn't handling the case where `i` is a
   # subslice of a block correctly (i.e. it ignores `i.indices`).
   @assert length(S.indices[Block(i)]) == length(i.indices)
-  return BlockSlice(S.blocks[Int(Block(i))], S.indices[Block(i)])
+  return _blockslice(S.blocks[Int(Block(i))], S.indices[Block(i)])
 end
 
 # This is used in slicing like:
@@ -347,7 +354,7 @@ function blockrange(
   axis::AbstractUnitRange,
   r::BlockVector{<:BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}},
 )
-  return map(b -> Block(b), blocks(r))
+  return map(Block, blocks(r))
 end
 
 # This handles slicing with `:`/`Colon()`.
@@ -365,6 +372,14 @@ function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:Integer})
   return Block.(Base.OneTo(1))
 end
 
+function blockrange(axis::AbstractUnitRange, r::BlockIndexVector)
+  return Block(r):Block(r)
+end
+
+function blockrange(axis::AbstractUnitRange, r::Vector{<:BlockIndexVector})
+  return map(Block, r)
+end
+
 function blockrange(axis::AbstractUnitRange, r)
   return error("Slicing not implemented for range of type `$(typeof(r))`.")
 end

src/BlockArraysExtensions/blockedunitrange.jl

@@ -10,6 +10,7 @@ using BlockArrays:
   BlockVector,
   block,
   blockedrange,
+  blockfirsts,
   blockindex,
   blocklengths,
   findblock,
@@ -169,6 +170,19 @@ function blockedunitrange_getindices(
   return mortar(map(b -> a[b], blocks(indices)))
 end
 
+function blockedunitrange_getindices(
+  a::AbstractBlockedUnitRange, indices::AbstractArray{Bool}
+)
+  blocked_indices = BlockedVector(indices, axes(a))
+  return mortar(
+    map(Base.OneTo(blocklength(blocked_indices))) do b
+      binds = blocked_indices[Block(b)]
+      bstart = blockfirsts(only(axes(blocked_indices)))[b]
+      return findall(binds) .+ (bstart - 1)
+    end,
+  )
+end
+
 # TODO: Move this to a `BlockArraysExtensions` library.
 function blockedunitrange_getindices(a::AbstractBlockedUnitRange, indices)
   return error("Not implemented.")
@@ -197,6 +211,27 @@ function to_blockindices(a::AbstractBlockedUnitRange{<:Integer}, I::UnitRange{<:
   )
 end
 
+struct BlockIndexVector{T<:Integer,I<:AbstractVector{T},TB<:Integer} <:
+       AbstractVector{BlockIndex{1,Tuple{TB},Tuple{T}}}
+  block::Block{1,TB}
+  indices::I
+end
+Base.length(a::BlockIndexVector) = length(a.indices)
+Base.size(a::BlockIndexVector) = (length(a),)
+BlockArrays.Block(a::BlockIndexVector) = a.block
+Base.getindex(a::BlockIndexVector, I::Integer) = Block(a)[a.indices[I]]
+Base.copy(a::BlockIndexVector) = BlockIndexVector(a.block, copy(a.indices))
+
+function to_blockindices(a::AbstractBlockedUnitRange{<:Integer}, I::AbstractArray{Bool})
+  I_blocks = blocks(BlockedVector(I, blocklengths(a)))
+  return mortar(
+    map(eachindex(I_blocks)) do b
+      I_b = findall(I_blocks[b])
+      BlockIndexVector(Block(b), I_b)
+    end,
+  )
+end
+
 # This handles non-blocked slices.
 # For example:
 # a = BlockSparseArray{Float64}([2, 2, 2, 2])

src/abstractblocksparsearray/unblockedsubarray.jl

@@ -4,7 +4,10 @@ using BlockArrays: BlockArrays, Block, BlockIndexRange, BlockSlice
 using TypeParameterAccessors: TypeParameterAccessors, parenttype, similartype
 
 const UnblockedIndices = Union{
-  Vector{<:Integer},BlockSlice{<:Block{1}},BlockSlice{<:BlockIndexRange{1}}
+  Vector{<:Integer},
+  BlockSlice{<:Block{1}},
+  BlockSlice{<:BlockIndexRange{1}},
+  BlockSlice{<:BlockIndexVector},
 }
 
 const UnblockedSubArray{T,N} = SubArray{

src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -40,6 +40,12 @@ function Base.to_indices(
   return @interface BlockSparseArrayInterface() to_indices(a, inds, I)
 end
 
+function Base.to_indices(
+  a::AnyAbstractBlockSparseArray, inds, I::Tuple{AbstractArray{Bool},Vararg{Any}}
+)
+  return @interface BlockSparseArrayInterface() to_indices(a, inds, I)
+end
+
 # a[[Block(2), Block(1)], [Block(2), Block(1)]]
 function Base.to_indices(
   a::AnyAbstractBlockSparseArray, inds, I::Tuple{Vector{<:Block{1}},Vararg{Any}}

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -146,6 +146,14 @@ end
   return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...)
 end
 
+@interface ::AbstractBlockSparseArrayInterface function Base.to_indices(
+  a, inds, I::Tuple{AbstractArray{Bool},Vararg{Any}}
+)
+  bs1 = to_blockindices(inds[1], I[1])
+  I1 = BlockIndices(bs1, blockedunitrange_getindices(inds[1], I[1]))
+  return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...)
+end
+
 # Special case when there is no blocking.
 @interface ::AbstractBlockSparseArrayInterface function Base.to_indices(
   a,