Initial version of block merging

ITensor · Jun 27, 2024 · 1a22fd3 · 1a22fd3
1 parent 79bf2d4
commit 1a22fd3
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 1 deletion.
diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl
@@ -36,7 +36,9 @@ function Base.getindex(a::SingleBlockView{<:Any,N}, index::Vararg{Int,N}) where
 end
 
 reblock(a) = a
-function reblock(a::SubArray{<:Any,<:Any,<:Any,<:Tuple{Vararg{BlockSlice}}})
+function reblock(
+  a::SubArray{<:Any,<:Any,<:AbstractBlockSparseArray,<:Tuple{Vararg{BlockSlice}}}
+)
   return @view a.parent[map(i -> i.indices, a.indices)...]
 end
 

diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl
@@ -29,6 +29,8 @@ function BlockArrays.viewblock(
 ) where {N}
   return viewblock(a, Tuple(block)...)
 end
+
+# TODO: Define `blocksparse_viewblock`.
 function BlockArrays.viewblock(
   a::AbstractBlockSparseArray{<:Any,N}, block::Vararg{Block{1},N}
 ) where {N}
@@ -38,3 +40,45 @@ function BlockArrays.viewblock(
   end
   return BlockView(a, block)
 end
+
+function Base.view(
+  a::SubArray{
+    T,
+    N,
+    <:AbstractBlockSparseArray{T,N},
+    <:Tuple{Vararg{BlockSlice{<:BlockRange{1,<:Tuple{<:AbstractUnitRange{<:Integer}}}},N}},
+  },
+  block::Block{N},
+) where {T,N}
+  return viewblock(a, block)
+end
+function Base.view(
+  a::SubArray{
+    T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockSlice{<:BlockRange{1}},N}}
+  },
+  block::Vararg{Block{1},N},
+) where {T,N}
+  return viewblock(a, block...)
+end
+function BlockArrays.viewblock(
+  a::SubArray{
+    T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockSlice{<:BlockRange{1}}}}
+  },
+  block::Block{N},
+) where {T,N}
+  return viewblock(a, Tuple(block)...)
+end
+
+# TODO: Define `blocksparse_viewblock`.
+function BlockArrays.viewblock(
+  a::SubArray{
+    T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockSlice{<:BlockRange{1}}}}
+  },
+  block::Vararg{Block{1},N},
+) where {T,N}
+  I = CartesianIndex(Int.(block))
+  if I ∈ stored_indices(blocks(a))
+    return blocks(a)[I]
+  end
+  return BlockView(a, block)
+end
diff --git a/...sors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/...sors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl
@@ -280,6 +280,17 @@ function blocksparse_blocks(a::SubArray)
   return SparseSubArrayBlocks(a)
 end
 
+function blocksparse_blocks(
+  a::SubArray{
+    T,
+    N,
+    <:AbstractBlockSparseArray{T,N},
+    <:Tuple{Vararg{BlockSlice{<:BlockRange{1,<:Tuple{<:AbstractUnitRange{<:Integer}}}},N}},
+  },
+) where {T,N}
+  return @view blocks(parent(a))[map(i -> Int.(i.block), parentindices(a))...]
+end
+
 using BlockArrays: BlocksView
 # TODO: Is this correct in general?
 SparseArrayInterface.nstored(a::BlocksView) = 1