Better support for slicing with Colon

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.7.22"
+version = "0.7.23"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -185,6 +185,7 @@ const GenericBlockIndexVectorSlices = BlockIndices{
   <:BlockVector{<:GenericBlockIndex{1},<:Vector{<:BlockIndexVector}}
 }
 const SubBlockSliceCollection = Union{
+  Base.Slice,
   BlockIndexRangeSlice,
   BlockIndexRangeSlices,
   BlockIndexVectorSlices,

src/BlockArraysExtensions/blockedunitrange.jl

@@ -349,6 +349,15 @@ BlockArrays.Block(b::BlockIndexVector) = b.block
 
 Base.copy(a::BlockIndexVector) = BlockIndexVector(a.block, copy.(a.indices))
 
+# Copied from BlockArrays.BlockIndexRange.
+function Base.show(io::IO, B::BlockIndexVector)
+  show(io, Block(B))
+  print(io, "[")
+  print_tuple_elements(io, B.indices)
+  print(io, "]")
+end
+Base.show(io::IO, ::MIME"text/plain", B::BlockIndexVector) = show(io, B)
+
 function Base.getindex(b::AbstractBlockedUnitRange, Kkr::BlockIndexVector{1})
   return b[block(Kkr)][Kkr.indices...]
 end

src/BlockArraysExtensions/blockrange.jl

@@ -16,6 +16,11 @@ function Base.getindex(r::BlockUnitRange, I::Block{1})
   return eachblockaxis(r)[Int(I)] .+ (first(r.r[I]) - 1)
 end
 
+const BlockOneTo{T<:Integer,B,CS,R<:BlockedOneTo{T,CS}} = BlockUnitRange{T,B,CS,R}
+Base.axes(S::Base.Slice{<:BlockOneTo}) = (S.indices,)
+Base.axes1(S::Base.Slice{<:BlockOneTo}) = S.indices
+Base.unsafe_indices(S::Base.Slice{<:BlockOneTo}) = (S.indices,)
+
 function BlockArrays.combine_blockaxes(r1::BlockUnitRange, r2::BlockUnitRange)
   if eachblockaxis(r1) ≠ eachblockaxis(r2)
     return throw(ArgumentError("BlockUnitRanges must have the same block axes"))

src/abstractblocksparsearray/views.jl

@@ -302,6 +302,14 @@ end
 
 blockedslice_blocks(x::BlockSlice) = x.block
 blockedslice_blocks(x::BlockIndices) = x.blocks
+# Reinterpret the slice blockwise.
+function blockedslice_blocks(x::Base.Slice)
+  return mortar(
+    map(BlockRange(x.indices)) do b
+      return BlockIndexRange(b, Base.Slice(Base.axes1(x.indices[b])))
+    end,
+  )
+end
 
 # TODO: Define `@interface interface(a) viewblock`.
 function BlockArrays.viewblock(