More general truncation and 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.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,12 @@ function Base.getindex(r::BlockUnitRange, I::Block{1})
   return eachblockaxis(r)[Int(I)] .+ (first(r.r[I]) - 1)
 end
 
+using BlockArrays: BlockedOneTo
+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

@@ -189,6 +189,18 @@ function Base.view(
 ) where {T,N}
   return viewblock(a, block...)
 end
+# Fix ambiguity error with BlockArrays.jl.
+function Base.view(
+  a::SubArray{
+    T,
+    N,
+    <:AbstractBlockSparseArray{T,N},
+    <:Tuple{Vararg{Union{Base.Slice,BlockSlice{Union{},<:Integer}},N}},
+  },
+  block::Block{N},
+) where {T,N}
+  return viewblock(a, block)
+end
 
 # XXX: TODO: Distinguish if a sub-view of the block needs to be taken!
 # Define a new `SubBlockSlice` which is used in:
@@ -302,6 +314,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(
@@ -319,6 +339,11 @@ function BlockArrays.viewblock(
   end
   return @view parent(a)[brs...]
 end
+
+## function BlockArrays.viewblock
+##   a::SubArray{
+##     T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{Union{Base.Slice, BlockArrays.BlockSlice{Union{}, T} where T<:Integer, BlockSparseArrays.BlockIndices{Union{}, T} where T<:Integer},N}}, ::Vararg{Block{1}, N}) where {T, N}
+
 # TODO: Define `@interface interface(a) viewblock`.
 function BlockArrays.viewblock(
   a::SubArray{