Start fixing slicing with unit ranges

ITensor · Jun 21, 2024 · 951f495 · 951f495
1 parent 93ca493
commit 951f495
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/...src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/...src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl
@@ -25,6 +25,13 @@ function Base.to_indices(
   return blocksparse_to_indices(a, inds, I)
 end
 
+# a[1:2, 1:2]
+function Base.to_indices(
+  a::BlockSparseArrayLike, inds, I::Tuple{UnitRange{<:Integer},Vararg{Any}}
+)
+  return blocksparse_to_indices(a, inds, I)
+end
+
 # BlockArrays `AbstractBlockArray` interface
 BlockArrays.blocks(a::BlockSparseArrayLike) = blocksparse_blocks(a)
 

diff --git a/...sors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/...sors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl
@@ -57,11 +57,19 @@ function blocksparse_getindex(
   return a_merged
 end
 
+# a[[Block(2), Block(1)], [Block(2), Block(1)]]
 function blocksparse_to_indices(a, inds, I::Tuple{Vector{<:Block{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[1:2, 1:2]
+function blocksparse_to_indices(a, inds, I::Tuple{UnitRange{<:Integer},Vararg{Any}})
+  bs1 = blockrange(inds[1], I[1])
+  I1 = BlockSlice(bs1, blockedunitrange_getindices(inds[1], I[1]))
+  return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...)
+end
+
 # TODO: Need to implement this!
 function block_merge end