[BlockSparseArrays] Permute and merge blocks

[mtfishman]

This PR adds this functionality for merging blocks while preserving the block sparse structure:

julia> using BlockArrays: Block, BlockedVector

julia> using NDTensors.BlockSparseArrays: BlockSparseArray

julia> a = BlockSparseArray{Float64}([2, 2, 2, 2], [2, 2, 2, 2])
typeof(axes) = Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}

Warning: To temporarily circumvent a bug in printing BlockSparseArrays with mixtures of dual and non-dual axes, the types of the dual axes printed below might not be accurate. The types printed above this message are the correct ones.

4×4-blocked 8×8 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}}, Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}:
 0.0  0.0  │  0.0  0.0  │  0.0  0.0  │  0.0  0.0
 0.0  0.0  │  0.0  0.0  │  0.0  0.0  │  0.0  0.0
 ──────────┼────────────┼────────────┼──────────
 0.0  0.0  │  0.0  0.0  │  0.0  0.0  │  0.0  0.0
 0.0  0.0  │  0.0  0.0  │  0.0  0.0  │  0.0  0.0
 ──────────┼────────────┼────────────┼──────────
 0.0  0.0  │  0.0  0.0  │  0.0  0.0  │  0.0  0.0
 0.0  0.0  │  0.0  0.0  │  0.0  0.0  │  0.0  0.0
 ──────────┼────────────┼────────────┼──────────
 0.0  0.0  │  0.0  0.0  │  0.0  0.0  │  0.0  0.0
 0.0  0.0  │  0.0  0.0  │  0.0  0.0  │  0.0  0.0

julia> @views for I in [Block(1, 1), Block(2, 2), Block(3, 3), Block(4, 4)]
         a[I] = randn(size(a[I]))
       end

julia> a
typeof(axes) = Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}

Warning: To temporarily circumvent a bug in printing BlockSparseArrays with mixtures of dual and non-dual axes, the types of the dual axes printed below might not be accurate. The types printed above this message are the correct ones.

4×4-blocked 8×8 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}}, Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}:
 0.635764  -0.347639  │   0.0      0.0       │   0.0        0.0       │   0.0        0.0     
 0.675381   1.19467   │   0.0      0.0       │   0.0        0.0       │   0.0        0.0     
 ─────────────────────┼──────────────────────┼────────────────────────┼──────────────────────
 0.0        0.0       │   1.83942  0.484242  │   0.0        0.0       │   0.0        0.0     
 0.0        0.0       │  -1.08625  0.541457  │   0.0        0.0       │   0.0        0.0     
 ─────────────────────┼──────────────────────┼────────────────────────┼──────────────────────
 0.0        0.0       │   0.0      0.0       │   0.883359  -2.29911   │   0.0        0.0     
 0.0        0.0       │   0.0      0.0       │  -2.80531    0.267481  │   0.0        0.0     
 ─────────────────────┼──────────────────────┼────────────────────────┼──────────────────────
 0.0        0.0       │   0.0      0.0       │   0.0        0.0       │  -0.397101   0.506121
 0.0        0.0       │   0.0      0.0       │   0.0        0.0       │  -0.255275  -0.36543 

julia> I = BlockedVector(Block.(1:4), [2, 2])
2-blocked 4-element BlockedVector{Block{1, Int64}, BlockArrays.BlockRange{1, Tuple{UnitRange{Int64}}}, Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}:
 Block(1)
 Block(2)
 ────────
 Block(3)
 Block(4)

julia> a[I, I]
typeof(axes) = Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}

Warning: To temporarily circumvent a bug in printing BlockSparseArrays with mixtures of dual and non-dual axes, the types of the dual axes printed below might not be accurate. The types printed above this message are the correct ones.

2×2-blocked 8×8 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}}, Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}:
 0.635764  -0.347639   0.0      0.0       │   0.0        0.0        0.0        0.0     
 0.675381   1.19467    0.0      0.0       │   0.0        0.0        0.0        0.0     
 0.0        0.0        1.83942  0.484242  │   0.0        0.0        0.0        0.0     
 0.0        0.0       -1.08625  0.541457  │   0.0        0.0        0.0        0.0     
 ─────────────────────────────────────────┼────────────────────────────────────────────
 0.0        0.0        0.0      0.0       │   0.883359  -2.29911    0.0        0.0     
 0.0        0.0        0.0      0.0       │  -2.80531    0.267481   0.0        0.0     
 0.0        0.0        0.0      0.0       │   0.0        0.0       -0.397101   0.506121
 0.0        0.0        0.0      0.0       │   0.0        0.0       -0.255275  -0.36543 

julia> using NDTensors.BlockSparseArrays: block_stored_indices

julia> block_stored_indices(a[I, I])
2-element Dictionaries.Dictionary{CartesianIndex{2}, Block{2, Int64}}
 CartesianIndex(1, 1) │ Block(1, 1)
 CartesianIndex(2, 2) │ Block(2, 2)

It also adds support for setting a new blocking at a more fine-grained level while preserving the block sparse structure:

julia> using BlockArrays: blockedrange

julia> I = blockedrange([1, 3, 3, 1])
4-blocked 8-element BlockArrays.BlockedOneTo{Int64, Vector{Int64}}:
 1
 ─
 2
 3
 4
 ─
 5
 6
 7
 ─
 8

julia> a[I, I]
typeof(axes) = Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}

Warning: To temporarily circumvent a bug in printing BlockSparseArrays with mixtures of dual and non-dual axes, the types of the dual axes printed below might not be accurate. The types printed above this message are the correct ones.

4×4-blocked 8×8 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}}, Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}:
 0.635764  │  -0.347639   0.0      0.0       │   0.0        0.0        0.0       │   0.0     
 ──────────┼─────────────────────────────────┼───────────────────────────────────┼───────────
 0.675381  │   1.19467    0.0      0.0       │   0.0        0.0        0.0       │   0.0     
 0.0       │   0.0        1.83942  0.484242  │   0.0        0.0        0.0       │   0.0     
 0.0       │   0.0       -1.08625  0.541457  │   0.0        0.0        0.0       │   0.0     
 ──────────┼─────────────────────────────────┼───────────────────────────────────┼───────────
 0.0       │   0.0        0.0      0.0       │   0.883359  -2.29911    0.0       │   0.0     
 0.0       │   0.0        0.0      0.0       │  -2.80531    0.267481   0.0       │   0.0     
 0.0       │   0.0        0.0      0.0       │   0.0        0.0       -0.397101  │   0.506121
 ──────────┼─────────────────────────────────┼───────────────────────────────────┼───────────
 0.0       │   0.0        0.0      0.0       │   0.0        0.0       -0.255275  │  -0.36543 

julia> block_stored_indices(a[I, I])
8-element Dictionaries.Dictionary{CartesianIndex{2}, Block{2, Int64}}
 CartesianIndex(1, 1) │ Block(1, 1)
 CartesianIndex(2, 1) │ Block(2, 1)
 CartesianIndex(1, 2) │ Block(1, 2)
 CartesianIndex(2, 2) │ Block(2, 2)
 CartesianIndex(3, 3) │ Block(3, 3)
 CartesianIndex(4, 3) │ Block(4, 3)
 CartesianIndex(3, 4) │ Block(3, 4)
 CartesianIndex(4, 4) │ Block(4, 4)

[mtfishman]

This PR now supports permuting and merging blocks in a single slicing operation:

julia> using BlockArrays: Block, BlockedVector

julia> using NDTensors.BlockSparseArrays: BlockSparseArray, block_stored_indices

julia> a = BlockSparseArray{Float64}([2, 2, 2, 2], [2, 2, 2, 2]);

julia> @views for I in [Block(1, 1), Block(2, 2), Block(3, 3), Block(4, 4)]
         a[I] = randn(size(a[I]))
       end

julia> a
typeof(axes) = Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}

Warning: To temporarily circumvent a bug in printing BlockSparseArrays with mixtures of dual and non-dual axes, the types of the dual axes printed below might not be accurate. The types printed above this message are the correct ones.

4×4-blocked 8×8 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}}, Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}:
  0.35991  -0.363998  │   0.0       0.0      │   0.0       0.0        │   0.0       0.0     
 -1.0964    0.478964  │   0.0       0.0      │   0.0       0.0        │   0.0       0.0     
 ─────────────────────┼──────────────────────┼────────────────────────┼─────────────────────
  0.0       0.0       │   1.36241   1.66477  │   0.0       0.0        │   0.0       0.0     
  0.0       0.0       │  -1.44313  -1.43454  │   0.0       0.0        │   0.0       0.0     
 ─────────────────────┼──────────────────────┼────────────────────────┼─────────────────────
  0.0       0.0       │   0.0       0.0      │   0.253089  1.26423    │   0.0       0.0     
  0.0       0.0       │   0.0       0.0      │  -0.581412  0.0754013  │   0.0       0.0     
 ─────────────────────┼──────────────────────┼────────────────────────┼─────────────────────
  0.0       0.0       │   0.0       0.0      │   0.0       0.0        │   0.459771  0.227871
  0.0       0.0       │   0.0       0.0      │   0.0       0.0        │  -0.966115  0.223543

julia> I = BlockedVector(Block.(reverse(1:4)), [2, 2])
2-blocked 4-element BlockedVector{Block{1, Int64}}:
 Block(4)
 Block(3)
 ────────
 Block(2)
 Block(1)

julia> a[I, I]
typeof(axes) = Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}

Warning: To temporarily circumvent a bug in printing BlockSparseArrays with mixtures of dual and non-dual axes, the types of the dual axes printed below might not be accurate. The types printed above this message are the correct ones.

2×2-blocked 8×8 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}}, Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}:
  0.459771  0.227871   0.0       0.0        │   0.0       0.0       0.0       0.0     
 -0.966115  0.223543   0.0       0.0        │   0.0       0.0       0.0       0.0     
  0.0       0.0        0.253089  1.26423    │   0.0       0.0       0.0       0.0     
  0.0       0.0       -0.581412  0.0754013  │   0.0       0.0       0.0       0.0     
 ───────────────────────────────────────────┼─────────────────────────────────────────
  0.0       0.0        0.0       0.0        │   1.36241   1.66477   0.0       0.0     
  0.0       0.0        0.0       0.0        │  -1.44313  -1.43454   0.0       0.0     
  0.0       0.0        0.0       0.0        │   0.0       0.0       0.35991  -0.363998
  0.0       0.0        0.0       0.0        │   0.0       0.0      -1.0964    0.478964

This is useful for abelian fusion, it is a key step in fusing multiple dimensions/axes/indices into a single one. I'll demonstrate that in the next post.

[mtfishman]

Here is a demonstration of abelian fusion, specifically matricizing an n-dimensional abelian symmetric tensor, as of this PR (before it was permuting the blocks, but not merging adjacent blocks that are in the same symmetry sectors):

julia> using BlockArrays: Block

julia> using NDTensors.BlockSparseArrays: BlockSparseArray

julia> using NDTensors.GradedAxes: GradedAxes, gradedrange

julia> using NDTensors.TensorAlgebra: fusedims, splitdims

julia> struct U1
         n::Int
       end

julia> GradedAxes.dual(c::U1) = U1(-c.n)

julia> GradedAxes.fuse_labels(c1::U1, c2::U1) = U1(c1.n + c2.n)

julia> Base.isless(c1::U1, c2::U1) = isless(c1.n, c2.n)

julia> r = gradedrange([U1(0) => 1, U1(1) => 1])
2-blocked 2-element BlockArrays.BlockedOneTo{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}, Vector{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}}}:
 NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}(1, U1(0))
 ──────────────────────────────────────────────────────────────
 NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}(2, U1(1))

julia> a = BlockSparseArray{Float64}(r, r, r, r);

julia> @views for I in [
         Block(1, 1, 1, 1),
         Block(2, 1, 2, 1),
         Block(2, 1, 1, 2),
         Block(1, 2, 2, 1),
         Block(1, 2, 1, 2),
         Block(2, 2, 2, 2),
       ]
         a[I] = randn(size(a[I]))
       end

julia> a
typeof(axes) = NTuple{4, BlockArrays.BlockedOneTo{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}, Vector{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}}}}

Warning: To temporarily circumvent a bug in printing BlockSparseArrays with mixtures of dual and non-dual axes, the types of the dual axes printed below might not be accurate. The types printed above this message are the correct ones.

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, BlockArrays.BlockedOneTo{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}, Vector{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}}}}}}, NTuple{4, BlockArrays.BlockedOneTo{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}, Vector{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}}}}}:
[:, :, 1, 1] =
 0.0958877  0.0
 0.0        0.0

[:, :, 2, 1] =
 0.0      -0.111154
 1.75521   0.0

[:, :, 1, 2] =
  0.0       -0.440967
 -0.550559   0.0

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

julia> m = fusedims(a, (1, 2), (3, 4))
typeof(axes) = Tuple{BlockArrays.BlockedOneTo{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}, Vector{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}}}, BlockArrays.BlockedOneTo{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}, Vector{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}}}}

Warning: To temporarily circumvent a bug in printing BlockSparseArrays with mixtures of dual and non-dual axes, the types of the dual axes printed below might not be accurate. The types printed above this message are the correct ones.

3×3-blocked 4×4 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockArrays.BlockedOneTo{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}, Vector{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}}}, BlockArrays.BlockedOneTo{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}, Vector{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}}}}}}, Tuple{BlockArrays.BlockedOneTo{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}, Vector{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}}}, BlockArrays.BlockedOneTo{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}, Vector{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}}}}}:
 0.0958877  │   0.0        0.0       │   0.0    
 ───────────┼────────────────────────┼──────────
 0.0        │   1.75521   -0.550559  │   0.0    
 0.0        │  -0.111154  -0.440967  │   0.0    
 ───────────┼────────────────────────┼──────────
 0.0        │   0.0        0.0       │  -1.13633

We can also split the dimensions to recover the original n-dimensional tensor with splitdims:

julia> splitdims(m, (r, r), (r, r))
typeof(axes) = NTuple{4, BlockArrays.BlockedOneTo{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}, Vector{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}}}}

Warning: To temporarily circumvent a bug in printing BlockSparseArrays with mixtures of dual and non-dual axes, the types of the dual axes printed below might not be accurate. The types printed above this message are the correct ones.

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, BlockArrays.BlockedOneTo{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}, Vector{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}}}}}}, NTuple{4, BlockArrays.BlockedOneTo{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}, Vector{NDTensors.LabelledNumbers.LabelledInteger{Int64, U1}}}}}:
[:, :, 1, 1] =
 0.0958877  0.0
 0.0        0.0

[:, :, 2, 1] =
 0.0      -0.111154
 1.75521   0.0

[:, :, 1, 2] =
  0.0       -0.440967
 -0.550559   0.0

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

@emstoudenmire to translate this back to the existing NDTensors.j/ITensors.jl code, this is performing two combiner operations (one combining the first two indices, and one combining the last two indices) in a single operation. The code implementing fusedims is quite simple in the end. Basically the steps are:

1. call permutedims if needed depending on the fusion call,
2. perform a naive reshape into a block sparse matrix that doesn't involve any block permutations/block merging, and
3. use the block permutation and merging operation shown in my previous comment which are determined by sorting and grouping the blocks based on the symmetry sector.

This PR brings us one step closer to replacing the current TensorStorage-based system of NDTensors.jl with the new system based off of the new submodules NamedDimsArrays, BlockSparseArrays, DiagonalArrays, GradedAxes, Sectors, etc. This is a key operation needed for matricizing n-dimension abelian symmetric tensors in order to perform blockwise factorizations like QR, eigendecompositions, and SVD. Performing the actualy blockwise matrix factorizations still need to be implemented, that will be left for a future PR, but once those are implemented we should be feature-complete with the TensorStorage-based code, besides fixing some known issues with running the new code on GPUs, and also a few missing features like direct summing/concatenating block sparse arrays.

[mtfishman]

@ogauthe this involves a lot of changes to the design of slicing of BlockSparseArrays, but doesn't change the interface (it just adds support for new slicing operations). Some of the output types of slicing operations are different after this PR, but in general they should be simpler and easier to work with. Once tests pass I'll merge and register this, let me know if you have any issues with it in your code.

Note that fusedims can be used to convert "dense" abelian symmetric tensors to matricized/fusion tensors, so is an alternative to your implementation when the tensors are abelian. It would be good to confirm it outputs the same thing as your code for abelian tensors, for example that the two follow the same conventions in terms of where blocks end up in the data matrix. Also my implementation is quite different from yours, it would be interesting to compare the steps of the two and see if they can share some operations or be designed in similar ways. I realize a lot of the steps will be different to handle non-abelian symmetries but I think it would still be an illuminating exercise and may inspire a new perspective on how to design the dense to fusion tensor conversion for non-abelian tensors.

[ogauthe]

Nice! I did not observe any change in BlockSparseArrays behavior in FusionTensors, the changes are transparent.

I confirm that I get the same data matrix as obtained from fusedims, I just need to take dual axes in the domain. This is already my convention; if I change to dual in the codomain, then the matrix blocks would be in opposite order. I don't know if fine control of dual is a feature you wish for a BlockSparseArray, it is not strictly required in the abelian case.

ft = FusionTensor(a, GradedAxes.dual.((r,r)), (r,r))
data_matrix(ft) ≈ m  # true