[BlockSparseArrays] Sub-slices of multiple blocks

[mtfishman]

This adds support for taking slices of multiple blocks at once (as proposed in ITensor/BlockSparseArrays.jl#2 and JuliaArrays/BlockArrays.jl#358), for example for a block sparse array:

using BlockArrays: Block
using NDTensors.BlockSparseArrays: BlockSparseArray
a = BlockSparseArray{Float64}([3, 3], [3, 3])
a[Block(1, 1)] = randn(3, 3)
a[Block(2, 2)] = randn(3, 3)

we can now do this:

julia> a
typeof(axes) = Tuple{BlockArrays.BlockedUnitRange{Vector{Int64}}, BlockArrays.BlockedUnitRange{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 6×6 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockArrays.BlockedUnitRange{Vector{Int64}}, BlockArrays.BlockedUnitRange{Vector{Int64}}}}}, Tuple{BlockArrays.BlockedUnitRange{Vector{Int64}}, BlockArrays.BlockedUnitRange{Vector{Int64}}}}:
 -0.945395   -0.116158  -0.470018  │   0.0        0.0        0.0     
  0.0648115   1.33954   -0.53128   │   0.0        0.0        0.0     
  2.21892     1.07313   -1.39657   │   0.0        0.0        0.0     
 ──────────────────────────────────┼─────────────────────────────────
  0.0         0.0        0.0       │   0.961448   0.667434   0.490545
  0.0         0.0        0.0       │  -0.282894  -0.614402  -1.1862  
  0.0         0.0        0.0       │   1.12401    0.924081   0.397953

julia> I = [Block(1)[2:3], Block(2)[1:2]];

julia> a[I, I]
typeof(axes) = Tuple{BlockArrays.BlockedUnitRange{Vector{Int64}}, BlockArrays.BlockedUnitRange{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 4×4 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockArrays.BlockedUnitRange{Vector{Int64}}, BlockArrays.BlockedUnitRange{Vector{Int64}}}}}, Tuple{BlockArrays.BlockedUnitRange{Vector{Int64}}, BlockArrays.BlockedUnitRange{Vector{Int64}}}}:
 1.33954  -0.53128  │   0.0        0.0     
 1.07313  -1.39657  │   0.0        0.0     
 ───────────────────┼──────────────────────
 0.0       0.0      │   0.961448   0.667434
 0.0       0.0      │  -0.282894  -0.614402

One thing you can see is that the implementation is pretty minimal, because it can make use of the infrastructure built for other slicing operations in BlockSparseArrays and the BlockArrays.jl package.

Additionally, I think this kind of slicing operation will be very useful for implementing low-rank block-wise matrix factorizations. For example, we could reduce the rank of an SVD like this:

# Contrived non-truncated block-wise
# SVD of a block diagonal matrix.
using LinearAlgebra: qr

u = BlockSparseArray{Float64}([3, 3], [3, 3])
u[Block(1, 1)] = Matrix(qr(randn(3, 3)).Q)
u[Block(2, 2)] = Matrix(qr(randn(3, 3)).Q)

s = BlockSparseArray{Float64}([3, 3], [3, 3])
s[Block(1, 1)] = Diagonal(normalize([0.9, 0.3, 0.005]))
s[Block(2, 2)] = Diagonal(normalize([0.8, 0.6, 0.001]))

v = BlockSparseArray{Float64}([3, 3], [3, 3])
v[Block(1, 1)] = Matrix(qr(randn(3, 3)).Q)
v[Block(2, 2)] = Matrix(qr(randn(3, 3)).Q)

block_ranks = [2, 2]
I = [Block(1)[1:block_ranks[1]], Block(2)[1:block_ranks[2]]]
u_truncated = u[:, I]
s_truncated = s[I, I]
v_truncated = v[I, :]

which truncates the singular values and associated singular vectors of each block down to the specified block ranks:

julia> u
typeof(axes) = Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{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 6×6 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}}, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
 -0.201584   0.975938  0.0831202  │   0.0        0.0       0.0      
 -0.137013  -0.112125  0.984203   │   0.0        0.0       0.0      
  0.969841   0.187011  0.156319   │   0.0        0.0       0.0      
 ─────────────────────────────────┼─────────────────────────────────
  0.0        0.0       0.0        │  -0.79183    0.607681  0.0610696
  0.0        0.0       0.0        │  -0.155642  -0.29747   0.941959 
  0.0        0.0       0.0        │   0.590577   0.736366  0.330126 

julia> u_truncated
typeof(axes) = Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{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 6×4 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}}, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
 -0.201584   0.975938  │   0.0        0.0     
 -0.137013  -0.112125  │   0.0        0.0     
  0.969841   0.187011  │   0.0        0.0     
 ──────────────────────┼──────────────────────
  0.0        0.0       │  -0.79183    0.607681
  0.0        0.0       │  -0.155642  -0.29747 
  0.0        0.0       │   0.590577   0.736366

julia> s
typeof(axes) = Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{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 6×6 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}}, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
 0.94867  0.0       0.0         │  0.0  0.0  0.0  
 0.0      0.316223  0.0         │  0.0  0.0  0.0  
 0.0      0.0       0.00527039  │  0.0  0.0  0.0  
 ───────────────────────────────┼─────────────────
 0.0      0.0       0.0         │  0.8  0.0  0.0  
 0.0      0.0       0.0         │  0.0  0.6  0.0  
 0.0      0.0       0.0         │  0.0  0.0  0.001

julia> s_truncated
typeof(axes) = Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{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 4×4 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}}, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}}:
 0.94867  0.0       │  0.0  0.0
 0.0      0.316223  │  0.0  0.0
 ───────────────────┼──────────
 0.0      0.0       │  0.8  0.0
 0.0      0.0       │  0.0  0.6

julia> norm(u_truncated * s_truncated * v_truncated - u * s * v)
0.005364420303655687

This kind of slicing can be performed either with the syntax I = [Block(1)[1:2], Block(2)[1:2]] to define the slice of the blocks in each dimension as shown above or using the syntax I = mortar([Block(1)[1:2], Block(2)[1:2]]). mortar is one of the BlockArrays.jl interfaces for making a BlockVector, so it acts like a vector over [Block(1)[1], Block(1)[2], Block(2)[1], Block(2)[2]] but keeps track of the block structure as well.

To-do:

• Add support for slicing with I = [Block(1)[1:2], Block(2)[1:2]], in addition to the current implementation which is based around I = mortar([Block(1)[1:2], Block(2)[1:2]]).

@ogauthe, @emstoudenmire