Generalize block sparse matricize (#163)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "BlockSparseArrays"
 uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
-version = "0.8.1"
+version = "0.9.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -22,9 +22,11 @@ TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 
 [weakdeps]
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
+TensorProducts = "decf83d6-1968-43f4-96dc-fdb3fe15fc6d"
 
 [extensions]
 BlockSparseArraysTensorAlgebraExt = "TensorAlgebra"
+BlockSparseArraysTensorProductsExt = "TensorProducts"
 
 [compat]
 Adapt = "4.1.1"
@@ -43,6 +45,7 @@ MatrixAlgebraKit = "0.2.2"
 SparseArraysBase = "0.7.1"
 SplitApplyCombine = "1.2.3"
 TensorAlgebra = "0.3.2"
+TensorProducts = "0.1.7"
 Test = "1.10"
 TypeParameterAccessors = "0.4.1"
 julia = "1.10"

docs/Project.toml

@@ -6,6 +6,6 @@ Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 
 [compat]
 BlockArrays = "1"
-BlockSparseArrays = "0.8"
+BlockSparseArrays = "0.9"
 Documenter = "1"
 Literate = "2"

examples/Project.toml

@@ -5,5 +5,5 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 BlockArrays = "1"
-BlockSparseArrays = "0.8"
+BlockSparseArrays = "0.9"
 Test = "1"

ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl

@@ -15,19 +15,51 @@ function TensorAlgebra.FusionStyle(::Type{<:AbstractBlockSparseArray})
   return BlockReshapeFusion()
 end
 
+using BlockArrays: Block, blocklength, blocks
+using BlockSparseArrays: blocksparse
+using SparseArraysBase: eachstoredindex
+using TensorAlgebra: TensorAlgebra, matricize, unmatricize
 function TensorAlgebra.matricize(
   ::BlockReshapeFusion, a::AbstractArray, biperm::BlockedTrivialPermutation{2}
 )
-  new_axes = fuseaxes(axes(a), biperm)
-  return blockreshape(a, new_axes)
+  ax = fuseaxes(axes(a), biperm)
+  reshaped_blocks_a = reshape(blocks(a), map(blocklength, ax))
+  key(I) = Block(Tuple(I))
+  value(I) = matricize(reshaped_blocks_a[I], biperm)
+  Is = eachstoredindex(reshaped_blocks_a)
+  bs = if isempty(Is)
+    # Catch empty case and make sure the type is constrained properly.
+    # This seems to only be necessary in Julia versions below v1.11,
+    # try removing it when we drop support for those versions.
+    keytype = Base.promote_op(key, eltype(Is))
+    valtype = Base.promote_op(value, eltype(Is))
+    valtype′ = !isconcretetype(valtype) ? AbstractMatrix{eltype(a)} : valtype
+    Dict{keytype,valtype′}()
+  else
+    Dict(key(I) => value(I) for I in Is)
+  end
+  return blocksparse(bs, ax)
 end
 
+using BlockArrays: blocklengths
 function TensorAlgebra.unmatricize(
   ::BlockReshapeFusion,
   m::AbstractMatrix,
-  blocked_axes::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractUnitRange}}},
+  blocked_ax::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractUnitRange}}},
 )
-  return blockreshape(m, Tuple(blocked_axes)...)
+  ax = Tuple(blocked_ax)
+  reshaped_blocks_m = reshape(blocks(m), map(blocklength, ax))
+  function f(I)
+    block_axes_I = BlockedTuple(
+      map(ntuple(identity, length(ax))) do i
+        return Base.axes1(ax[i][Block(I[i])])
+      end,
+      blocklengths(blocked_ax),
+    )
+    return unmatricize(reshaped_blocks_m[I], block_axes_I)
+  end
+  bs = Dict(Block(Tuple(I)) => f(I) for I in eachstoredindex(reshaped_blocks_m))
+  return blocksparse(bs, ax)
 end
 
 end

ext/BlockSparseArraysTensorProductsExt/BlockSparseArraysTensorProductsExt.jl

@@ -0,0 +1,14 @@
+module BlockSparseArraysTensorProductsExt
+
+using BlockSparseArrays: BlockUnitRange, blockrange, eachblockaxis
+using TensorProducts: TensorProducts, tensor_product
+# TODO: Dispatch on `FusionStyle` to allow different kinds of products,
+# for example to allow merging common symmetry sectors.
+function TensorProducts.tensor_product(a1::BlockUnitRange, a2::BlockUnitRange)
+  new_blockaxes = vec(
+    map(splat(tensor_product), Iterators.product(eachblockaxis(a1), eachblockaxis(a2)))
+  )
+  return blockrange(new_blockaxes)
+end
+
+end

src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -180,7 +180,10 @@ end
 # ```
 # but includes `BlockIndices`, where the blocks aren't contiguous.
 const BlockSliceCollection = Union{
-  Base.Slice,BlockSlice{<:BlockRange{1}},BlockIndices{<:Vector{<:Block{1}}}
+  Base.Slice,
+  BlockSlice{<:Block{1}},
+  BlockSlice{<:BlockRange{1}},
+  BlockIndices{<:Vector{<:Block{1}}},
 }
 const BlockIndexRangeSlice = BlockSlice{
   <:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}}
@@ -267,13 +270,15 @@ tuple_oneto(n) = ntuple(identity, n)
 
 function _blockreshape(a::AbstractArray, axes::Tuple{Vararg{AbstractUnitRange}})
   reshaped_blocks_a = reshape(blocks(a), blocklength.(axes))
-  reshaped_a = similar(a, axes)
-  for I in eachstoredindex(reshaped_blocks_a)
-    block_size_I = map(i -> length(axes[i][Block(I[i])]), tuple_oneto(length(axes)))
+  function f(I)
+    block_axes_I = map(ntuple(identity, length(axes))) do i
+      return Base.axes1(axes[i][Block(I[i])])
+    end
     # TODO: Better converter here.
-    reshaped_a[Block(Tuple(I))] = reshape(reshaped_blocks_a[I], block_size_I)
+    return reshape(reshaped_blocks_a[I], block_axes_I)
   end
-  return reshaped_a
+  bs = Dict(Block(Tuple(I)) => f(I) for I in eachstoredindex(reshaped_blocks_a))
+  return blocksparse(bs, axes)
 end
 
 function blockreshape(

src/BlockArraysExtensions/blockrange.jl

@@ -8,13 +8,21 @@ end
 function blockrange(eachblockaxis)
   return BlockUnitRange(blockedrange(length.(eachblockaxis)), eachblockaxis)
 end
+function blockrange(first::Integer, eachblockaxis)
+  return BlockUnitRange(blockedrange(first, length.(eachblockaxis)), eachblockaxis)
+end
 Base.first(r::BlockUnitRange) = first(r.r)
 Base.last(r::BlockUnitRange) = last(r.r)
 BlockArrays.blocklasts(r::BlockUnitRange) = blocklasts(r.r)
 eachblockaxis(r::BlockUnitRange) = r.eachblockaxis
 function Base.getindex(r::BlockUnitRange, I::Block{1})
   return eachblockaxis(r)[Int(I)] .+ (first(r.r[I]) - 1)
 end
+function Base.getindex(r::BlockUnitRange, I::BlockRange{1})
+  return blockrange(first(r), eachblockaxis(r)[Int.(I)])
+end
+Base.axes(r::BlockUnitRange) = (blockrange(eachblockaxis(r)),)
+Base.axes1(r::BlockUnitRange) = blockrange(eachblockaxis(r))
 
 using BlockArrays: BlockedOneTo
 const BlockOneTo{T<:Integer,B,CS,R<:BlockedOneTo{T,CS}} = BlockUnitRange{T,B,CS,R}

src/abstractblocksparsearray/abstractblocksparsearray.jl

@@ -87,12 +87,19 @@ function Base.setindex!(
       ),
     )
   end
-  # Custom `_convert` works around the issue that
-  # `convert(::Type{<:Diagonal}, ::AbstractMatrix)` isnt' defined
-  # in Julia v1.10 (https://github.com/JuliaLang/julia/pull/48895,
-  # https://github.com/JuliaLang/julia/pull/52487).
-  # TODO: Delete once we drop support for Julia v1.10.
-  blocks(a)[Int.(I)...] = _convert(blocktype(a), value)
+  if isstored(a, I...)
+    # This writes into existing blocks, or constructs blocks
+    # using the axes.
+    aI = @view! a[I...]
+    aI .= value
+  else
+    # Custom `_convert` works around the issue that
+    # `convert(::Type{<:Diagonal}, ::AbstractMatrix)` isnt' defined
+    # in Julia v1.10 (https://github.com/JuliaLang/julia/pull/48895,
+    # https://github.com/JuliaLang/julia/pull/52487).
+    # TODO: Delete `_convert` once we drop support for Julia v1.10.
+    blocks(a)[Int.(I)...] = _convert(blocktype(a), value)
+  end
   return a
 end
 

src/blocksparsearray/blocksparsearray.jl

@@ -262,13 +262,18 @@ function blocksparsezeros(::BlockType{A}, axes...) where {A<:AbstractArray}
   # to make a bit more generic.
   return BlockSparseArray{eltype(A),ndims(A),A}(undef, axes...)
 end
-function blocksparse(d::Dict{<:Block,<:AbstractArray}, axes...)
-  a = blocksparsezeros(BlockType(valtype(d)), axes...)
+function blocksparse(d::Dict{<:Block,<:AbstractArray}, ax::Tuple)
+  a = blocksparsezeros(BlockType(valtype(d)), ax...)
   for I in eachindex(d)
     a[I] = d[I]
   end
   return a
 end
+function blocksparse(
+  d::Dict{<:Block,<:AbstractArray}, blocklens::AbstractVector{<:Integer}...
+)
+  return blocksparse(d, map(blockedrange, blocklens))
+end
 
 # Base `AbstractArray` interface
 Base.axes(a::BlockSparseArray) = a.axes

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -537,7 +537,11 @@ function SparseArraysBase.getunstoredindex(
   return error("Not implemented.")
 end
 
+# Convert a blockwise slice on a block sparse array
+# to an elementwise slice on the corresponding sparse array
+# of blocks.
 to_blocks_indices(I::BlockSlice{<:BlockRange{1}}) = Int.(I.block)
+to_blocks_indices(I::BlockSlice{<:Block{1}}) = Int(I.block):Int(I.block)
 to_blocks_indices(I::BlockIndices{<:Vector{<:Block{1}}}) = Int.(I.blocks)
 to_blocks_indices(I::Base.Slice) = Base.OneTo(blocklength(I.indices))