[BlockSparseArrays] Towards abelian block fusion

[mtfishman]

This PR makes progress towards abelian symmetric block fusion.

• Define splitdims(::SectorFusion, ::BlockSparseArray, ...), the inverse operation of fusedims(::SectorFusion, ::BlockSparseArray, ...). This will fix the failing symmetric contraction tests.

This is the current functionality:

using BlockArrays: Block, blocksize
using NDTensors.BlockSparseArrays: BlockSparseArray
using NDTensors.GradedAxes: gradedrange
using NDTensors.Sectors: U1
using NDTensors.TensorAlgebra: fusedims, splitdims
using Random: randn!

d1 = gradedrange([U1(0) => 1, U1(1) => 1])
d2 = gradedrange([U1(1) => 1, U1(0) => 1])
a = BlockSparseArray{Float64}(d1, d2, d1, d2)
for i in 1:minimum(blocksize(a))
  b = Block(i, i, i, i)
  a[b] = randn!(a[b])
end
m = fusedims(a, (1, 2), (3, 4))
axes(m, 1)
axes(m, 2)
a′ = splitdims(m, (d1, d2), (d1, d2))
a == a′

which outputs:

Output

julia> d1 = gradedrange([U1(0) => 1, U1(1) => 1])
2-element Vector{U1}:
 U1(0)
 U1(1)
isdual = false
2-blocked 2-element BlockArrays.BlockedUnitRange{Vector{Int64}}:
 1
 ─
 2

julia> d2 = gradedrange([U1(1) => 1, U1(0) => 1])
2-element Vector{U1}:
 U1(1)
 U1(0)
isdual = false
2-blocked 2-element BlockArrays.BlockedUnitRange{Vector{Int64}}:
 1
 ─
 2

julia> a = BlockSparseArray{Float64}(d1, d2, d1, d2)
2×2×2×2-blocked 2×2×2×2 BlockSparseArray{Float64, 4, Array{Float64, 4}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Array{Float64, 4}, 4, NDTensors.BlockSparseArrays.BlockZero{NTuple{4, NDTensors.GradedAxes.GradedUnitRange{Vector{Int64}, U1}}}}, NTuple{4, NDTensors.GradedAxes.GradedUnitRange{Vector{Int64}, U1}}}:
[:, :, 1, 1] =
 0.0  0.0
 0.0  0.0

[:, :, 2, 1] =
 0.0  0.0
 0.0  0.0

[:, :, 1, 2] =
 0.0  0.0
 0.0  0.0

[:, :, 2, 2] =
 0.0  0.0
 0.0  0.0

julia> for i in 1:minimum(blocksize(a))
         b = Block(i, i, i, i)
         a[b] = randn!(a[b])
       end

julia> m = fusedims(a, (1, 2), (3, 4))
4×4-blocked 4×4 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{NDTensors.GradedAxes.GradedUnitRange{Vector{Int64}, U1}, NDTensors.GradedAxes.GradedUnitRange{Vector{Int64}, U1}}}}, Tuple{NDTensors.GradedAxes.GradedUnitRange{Vector{Int64}, U1}, NDTensors.GradedAxes.GradedUnitRange{Vector{Int64}, U1}}}:
 0.0   0.0        0.0       0.0
 0.0  -0.172995   0.0       0.0
 0.0   0.0       -0.919495  0.0
 0.0   0.0        0.0       0.0

julia> axes(m, 1)
4-element Vector{U1}:
 U1(0)
 U1(1)
 U1(1)
 U1(2)
isdual = false
4-blocked 4-element BlockArrays.BlockedUnitRange{Vector{Int64}}:
 1
 ─
 2
 ─
 3
 ─
 4

julia> axes(m, 2)
4-element Vector{U1}:
 U1(0)
 U1(1)
 U1(1)
 U1(2)
isdual = false
4-blocked 4-element BlockArrays.BlockedUnitRange{Vector{Int64}}:
 1
 ─
 2
 ─
 3
 ─
 4

julia> a′ = splitdims(m, (d1, d2), (d1, d2))
2×2×2×2-blocked 2×2×2×2 BlockSparseArray{Float64, 4, Array{Float64, 4}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Array{Float64, 4}, 4, NDTensors.BlockSparseArrays.BlockZero{NTuple{4, NDTensors.GradedAxes.GradedUnitRange{Vector{Int64}, U1}}}}, NTuple{4, NDTensors.GradedAxes.GradedUnitRange{Vector{Int64}, U1}}}:
[:, :, 1, 1] =
 -0.172995  0.0
  0.0       0.0

[:, :, 2, 1] =
 0.0  0.0
 0.0  0.0

[:, :, 1, 2] =
 0.0  0.0
 0.0  0.0

[:, :, 2, 2] =
 0.0   0.0
 0.0  -0.919495

julia> a == a′
true

So you can see that it sorts the blocks by sector, which is the first step towards fusion, but doesn't merge the blocks yet. I will cover block merging in a future PR.

@emstoudenmire @ogauthe

[mtfishman]

@ogauthe you were asking about ways of manipulating blocks of GradedUnitRange, here is a demonstration:

using NDTensors.GradedAxes: blockmergesortperm, fuse, gradedrange

d = gradedrange([U1(0) => 1, U1(1) => 1, U1(0) => 1])
d[[Block(1), Block(3), Block(2)]]
d[BlockVector([Block(1), Block(3), Block(2)], [2, 1])]
blockmergesortperm(d)
d[blockmergesortperm(d)]
fuse(d, d)

which outputs:

julia> using NDTensors.GradedAxes: blockmergesortperm, fuse, gradedrange

julia> d = gradedrange([U1(0) => 1, U1(1) => 1, U1(0) => 1])
3-element Vector{U1}:
 U1(0)
 U1(1)
 U1(0)
isdual = false
3-blocked 3-element BlockedUnitRange{Vector{Int64}}:
 1
 ─
 2
 ─
 3

julia> d[[Block(1), Block(3), Block(2)]]
3-element Vector{U1}:
 U1(0)
 U1(0)
 U1(1)
isdual = false
3-blocked 3-element BlockedUnitRange{Vector{Int64}}:
 1
 ─
 2
 ─
 3

julia> d[BlockVector([Block(1), Block(3), Block(2)], [2, 1])]
2-element Vector{U1}:
 U1(0)
 U1(1)
isdual = false
2-blocked 3-element BlockedUnitRange{Vector{Int64}}:
 1
 2
 ─
 3

julia> blockmergesortperm(d)
2-blocked 3-element BlockVector{Block{1, Int64}}:
 Block(1)
 Block(3)
 ────────
 Block(2)

julia> d[blockmergesortperm(d)]
2-element Vector{U1}:
 U1(0)
 U1(1)
isdual = false
2-blocked 3-element BlockedUnitRange{Vector{Int64}}:
 1
 2
 ─
 3

julia> fuse(d, d)
3-element Vector{U1}:
 U1(0)
 U1(1)
 U1(2)
isdual = false
3-blocked 9-element BlockedUnitRange{Vector{Int64}}:
 1
 2
 3
 4
 ─
 5
 6
 7
 8
 ─
 9

[codecov-commenter]

Codecov Report

Attention: 37 lines in your changes are missing coverage. Please review.

Comparison is base (98b95a2) 84.23% compared to head (21e2db4) 53.14%.
Report is 4 commits behind head on main.

❗ Current head 21e2db4 differs from pull request most recent head 45340c2. Consider uploading reports for the commit 45340c2 to get more accurate results

Files	Patch %	Lines
src/physics/autompo/opsum_to_mpo_generic.jl	0.00%	14 Missing ⚠️
src/mps/abstractprojmpo/projmpo_mps.jl	0.00%	8 Missing ⚠️
src/mps/abstractprojmpo/projmposum.jl	0.00%	7 Missing ⚠️
src/mps/abstractprojmpo/projmps.jl	0.00%	6 Missing ⚠️
src/mps/abstractmps.jl	0.00%	1 Missing ⚠️
src/mps/mpo.jl	0.00%	1 Missing ⚠️

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #1326       +/-   ##
===========================================
- Coverage   84.23%   53.14%   -31.10%     
===========================================
  Files         106       97        -9     
  Lines        8841     8020      -821     
===========================================
- Hits         7447     4262     -3185     
- Misses       1394     3758     +2364     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[emstoudenmire]

Really interesting - I like the interface

[mtfishman]

In the latest I added support for splitdims, which does the reverse of fusedims, specifically it inverses the permutation of the blocks and then reshapes back to a higher-order block sparse tensor.

[emstoudenmire]

Great, so basically does everything needed for Abelian symmetric tensors except some optimizations related to merging blocks (I guess merging them in the sense of permuting and merging the blocking inside the axes)?

[mtfishman]

This PR introduced permuting sectors to put common sectors next to each other. As you can see in the first post, the output of fusedims has U1(1) sectors which are next to each other, while the code before this PR was based on a more naive block reshaping where they would have been apart. The more naive code that doesn't involve sector permutations still exists, and is used if a BlockSparseArray doesn't have symmetries attached to the axes. Those two options are chosen through a FusionStyle trait that is derived by analyzing the types of the axes, which is also introduced in this PR.

The current implementation doesn't merge the adjacent U1(1) sectors yet. That will require some new functionality in BlockSparseArrays to add (or more general, perform mapping/broadcasting operations) of BlockSparseArrays with non-matching block structures. Currently mapping/broadcasting is hardcoded to only work with BlockSparseArrays where the axes have the same block structure.

[emstoudenmire]

All makes sense - thanks!

[ogauthe]

Concerning splitdims, I think the code should check that prod(split_axes) == axis_to_split and return an error if not.
One can be less strict by allowing missing blocks: it would be allowed to split a spin 0 into 2 spins 1/2, implicitly assuming there is a spin 1 block that disappeared. But it should always be forbidden to split U1(0) into U1(1) ⊗ U1(1).

[mtfishman]

Agreed, I was just keeping the code simpler for now but we can add checks like that.