Generalize matricize

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

ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl

@@ -15,19 +15,41 @@ 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))
+  bs = Dict(
+    Block(Tuple(I)) => matricize(reshaped_blocks_a[I], biperm) for
+    I in eachstoredindex(reshaped_blocks_a)
+  )
+  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

src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -267,13 +267,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/abstractblocksparsearray/abstractblocksparsearray.jl

@@ -92,7 +92,8 @@ function Base.setindex!(
   # 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)
+  aI = @view! a[I...]
+  copyto!(aI, value)
   return a
 end
 

src/blocksparsearray/blocksparsearray.jl

@@ -262,13 +262,19 @@ 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)
+  blockaxtype = Tuple{map(eltype ∘ eachblockaxis, ax)...}
+  # TODO: Catch if inference fails and use `valtype(d)` instead.
+  blockt = Base.promote_op(similar, Type{valtype(d)}, blockaxtype)
+  a = blocksparsezeros(BlockType(blockt), ax)
   for I in eachindex(d)
     a[I] = d[I]
   end
   return a
 end
+function blocksparse(d::Dict{<:Block,<:AbstractArray}, ax::AbstractUnitRange...)
+  return blocksparse(d, ax)
+end
 
 # Base `AbstractArray` interface
 Base.axes(a::BlockSparseArray) = a.axes

src/blocksparsearrayinterface/getunstoredblock.jl

@@ -20,7 +20,7 @@ end
 @inline function Base.getindex(a::ZeroBlocks{N,A}, I::Vararg{Int,N}) where {N,A}
   # TODO: Use `BlockArrays.eachblockaxes`.
   ax = ntuple(N) do d
-    return only(axes(a.parentaxes[d][Block(I[d])]))
+    return eachblockaxis(a.parentaxes[d])[I[d]]
   end
   !isconcretetype(A) && return zero!(similar(Array{eltype(A),N}, ax))
   return zero!(similar(A, ax))

src/factorizations/svd.jl

@@ -49,21 +49,22 @@ function MatrixAlgebraKit.initialize_output(
   ::typeof(svd_compact!), A::AbstractBlockSparseMatrix, alg::BlockPermutedDiagonalAlgorithm
 )
   bm, bn = blocksize(A)
-  bmn = min(bm, bn)
+  (bmn, mindim) = findmin((bm, bn))
 
   brows = eachblockaxis(axes(A, 1))
   bcols = eachblockaxis(axes(A, 2))
   u_axes = similar(brows, bmn)
-  v_axes = similar(brows, bmn)
+  v_axes = similar(bcols, bmn)
 
   # fill in values for blocks that are present
   bIs = collect(eachblockstoredindex(A))
   browIs = Int.(first.(Tuple.(bIs)))
   bcolIs = Int.(last.(Tuple.(bIs)))
   for bI in eachblockstoredindex(A)
     row, col = Int.(Tuple(bI))
-    u_axes[col] = infimum(brows[row], bcols[col])
-    v_axes[col] = infimum(bcols[col], brows[row])
+    dim = (row, col)[mindim]
+    u_axes[dim] = infimum(brows[row], bcols[col])
+    v_axes[dim] = infimum(bcols[col], brows[row])
   end
 
   # fill in values for blocks that aren't present, pairing them in order of occurence
@@ -83,9 +84,10 @@ function MatrixAlgebraKit.initialize_output(
   # allocate output
   for bI in eachblockstoredindex(A)
     brow, bcol = Tuple(bI)
+    bdim = (brow, bcol)[mindim]
     block = @view!(A[bI])
     block_alg = block_algorithm(alg, block)
-    U[brow, bcol], S[bcol, bcol], Vt[bcol, bcol] = MatrixAlgebraKit.initialize_output(
+    U[brow, bdim], S[bdim, bdim], Vt[bdim, bcol] = MatrixAlgebraKit.initialize_output(
       svd_compact!, block, block_alg
     )
   end