[BlockSparseArrays] Towards block merging

[mtfishman]

This adds some functionality that will be needed for slicing operations that merge blocks, one of the feature requests listed in ITensor/BlockSparseArrays.jl#2. These operations will be useful for fusion operations of symmetric tensors, removing symmetries from symmetric tensors, among other operations.

Given this block sparse array:

using BlockArrays: Block, BlockedVector, blockedrange
using NDTensors.BlockSparseArrays: BlockSparseArray
a = BlockSparseArray{Float64}([2, 2, 2, 2], [2, 2, 2, 2])
@views for I in [
  Block(1, 1),
  Block(2, 2),
  Block(3, 3),
  Block(4, 4),
]
  a[I] = randn(size(a[I]))
end

this PR enables the following operations:

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.765348  -1.15712    │   0.0        0.0       │   0.0       0.0       │   0.0        0.0     
 0.176078  -0.0421353  │   0.0        0.0       │   0.0       0.0       │   0.0        0.0     
 ──────────────────────┼────────────────────────┼───────────────────────┼──────────────────────
 0.0        0.0        │  -1.48538   -1.40037   │   0.0       0.0       │   0.0        0.0     
 0.0        0.0        │  -0.718883  -0.086352  │   0.0       0.0       │   0.0        0.0     
 ──────────────────────┼────────────────────────┼───────────────────────┼──────────────────────
 0.0        0.0        │   0.0        0.0       │  -0.671303  0.773009  │   0.0        0.0     
 0.0        0.0        │   0.0        0.0       │   0.367588  0.379923  │   0.0        0.0     
 ──────────────────────┼────────────────────────┼───────────────────────┼──────────────────────
 0.0        0.0        │   0.0        0.0       │   0.0       0.0       │  -0.610198  -0.397977
 0.0        0.0        │   0.0        0.0       │   0.0       0.0       │   0.887972  -0.850992

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

julia> @view a[I, I]
8×8 view(::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}}}}, BlockArrays.BlockedOneTo([4, 8]), BlockArrays.BlockedOneTo([4, 8])) with eltype Float64 with indices BlockArrays.BlockedOneTo([4, 8])×BlockArrays.BlockedOneTo([4, 8]):
 0.765348  -1.15712     0.0        0.0       │   0.0       0.0        0.0        0.0     
 0.176078  -0.0421353   0.0        0.0       │   0.0       0.0        0.0        0.0     
 0.0        0.0        -1.48538   -1.40037   │   0.0       0.0        0.0        0.0     
 0.0        0.0        -0.718883  -0.086352  │   0.0       0.0        0.0        0.0     
 ────────────────────────────────────────────┼───────────────────────────────────────────
 0.0        0.0         0.0        0.0       │  -0.671303  0.773009   0.0        0.0     
 0.0        0.0         0.0        0.0       │   0.367588  0.379923   0.0        0.0     
 0.0        0.0         0.0        0.0       │   0.0       0.0       -0.610198  -0.397977
 0.0        0.0         0.0        0.0       │   0.0       0.0        0.887972  -0.850992

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> @view a[I, I]
8×8 view(::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}}}}, BlockArrays.BlockedOneTo([4, 8]), BlockArrays.BlockedOneTo([4, 8])) with eltype Float64 with indices BlockArrays.BlockedOneTo([4, 8])×BlockArrays.BlockedOneTo([4, 8]):
 0.765348  -1.15712     0.0        0.0       │   0.0       0.0        0.0        0.0     
 0.176078  -0.0421353   0.0        0.0       │   0.0       0.0        0.0        0.0     
 0.0        0.0        -1.48538   -1.40037   │   0.0       0.0        0.0        0.0     
 0.0        0.0        -0.718883  -0.086352  │   0.0       0.0        0.0        0.0     
 ────────────────────────────────────────────┼───────────────────────────────────────────
 0.0        0.0         0.0        0.0       │  -0.671303  0.773009   0.0        0.0     
 0.0        0.0         0.0        0.0       │   0.367588  0.379923   0.0        0.0     
 0.0        0.0         0.0        0.0       │   0.0       0.0       -0.610198  -0.397977
 0.0        0.0         0.0        0.0       │   0.0       0.0        0.887972  -0.850992

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

julia> @view a[I, I]
8×8 view(::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}}}}, [7, 8, 5, 6, 3, 4, 1, 2], [7, 8, 5, 6, 3, 4, 1, 2]) with eltype Float64 with indices BlockArrays.BlockedOneTo([4, 8])×BlockArrays.BlockedOneTo([4, 8]):
 -0.610198  -0.397977   0.0       0.0       │   0.0        0.0       0.0        0.0      
  0.887972  -0.850992   0.0       0.0       │   0.0        0.0       0.0        0.0      
  0.0        0.0       -0.671303  0.773009  │   0.0        0.0       0.0        0.0      
  0.0        0.0        0.367588  0.379923  │   0.0        0.0       0.0        0.0      
 ───────────────────────────────────────────┼────────────────────────────────────────────
  0.0        0.0        0.0       0.0       │  -1.48538   -1.40037   0.0        0.0      
  0.0        0.0        0.0       0.0       │  -0.718883  -0.086352  0.0        0.0      
  0.0        0.0        0.0       0.0       │   0.0        0.0       0.765348  -1.15712  
  0.0        0.0        0.0       0.0       │   0.0        0.0       0.176078  -0.0421353

However, the implementation isn't complete yet, and basic operations like copy(b) and b[Block(1, 1)] don't work right now. I'll leave that for future PRs.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 43.51%. Comparing base (82cfd76) to head (8da6a14).
Report is 22 commits behind head on main.

❗ Current head 8da6a14 differs from pull request most recent head 398ce48

Please upload reports for the commit 398ce48 to get more accurate results.

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

❗ There is a different number of reports uploaded between BASE (82cfd76) and HEAD (8da6a14). Click for more details.

HEAD has 2 uploads less than BASE

| Flag | BASE (82cfd76) | HEAD (8da6a14) | |------|------|------| ||3|1|

Additional details and impacted files

@@             Coverage Diff             @@
##             main    ITensor/ITensors.jl#1512       +/-   ##
===========================================
- Coverage   78.05%   43.51%   -34.54%     
===========================================
  Files         148      136       -12     
  Lines        9679     8783      -896     
===========================================
- Hits         7555     3822     -3733     
- Misses       2124     4961     +2837     

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