[WIP] [BlockSparseArrays] (Truncated) SVD for Abstract(Block)(Sparse)Array

[lkdvos]

This is a first draft of functionality to extend singular value decompositions to blockarrays. It features quite a few components that probably have to be discussed further.

svd(!)

The core functionality is centered around the typical LinearAlgebra.svd.

usv = svd(A) # returns factorization object
U, S, V = usv # returns components

For generic block-arrays, the idea would be that the rows of U, and the columns of Vt have the same structure as the rows and columns of A, such that U * Diagonal(S) * Vt = A. For the columns of U and the rows of Vt however, there is no structure that can be inferred, such that this would generically be a single block.
Additionally, implementation-wise it is probably the most efficient to first cast to a single dense array, then do the SVD, and then restore the blocked structure. For that reason, the return type of the matrices is also BlockedArray instead of BlockArray. This is implemented and works:

julia> a = mortar([rand(i, j) for i in [2, 2], j in [2, 2]])
2×2-blocked 4×4 BlockMatrix{Float64}:
 0.899443  0.24551    │  0.735501  0.367721
 0.60512   0.354016   │  0.854856  0.807514
 ─────────────────────┼────────────────────
 0.379235  0.371387   │  0.540997  0.597405
 0.983355  0.0194433  │  0.187833  0.497739

julia> svd(a)
NDTensors.BlockSparseArrays.SVD{Float64, Float64, BlockedMatrix{Float64, Matrix{Float64}, Tuple{BlockedOneTo{Int64, Vector{Int64}}, BlockedOneTo{Int64, Vector{Int64}}}}, BlockedVector{Float64, Vector{Float64}, Tuple{BlockedOneTo{Int64, Vector{Int64}}}}, BlockedMatrix{Float64, Matrix{Float64}, Tuple{BlockedOneTo{Int64, Vector{Int64}}, BlockedOneTo{Int64, Vector{Int64}}}}}
U factor:
2×1-blocked 4×4 BlockedMatrix{Float64}:
 -0.538016   0.185128  -0.809959   0.14224  
 -0.593322  -0.454957   0.177765  -0.63983  
 ───────────────────────────────────────────
 -0.408608  -0.407648   0.310857   0.755136 
 -0.43766    0.769782   0.464471  -0.0124688
singular values:
1-blocked 4-element BlockedVector{Float64}:
 2.2459630332217984 
 0.6638504762523649 
 0.3266654288265734 
 0.09879357698918856
Vt factor:
1×2-blocked 4×4 BlockedMatrix{Float64}:
 -0.635932  -0.223688   │  -0.537043  -0.507088
  0.743515  -0.379664   │  -0.495153  -0.24055 
 -0.14178   -0.0350279  │  -0.576573   0.803888
  0.15058    0.896991   │  -0.366022  -0.19688

If we want to keep that behavior of the blockstructure, the LinearAlgebra.SVD struct is however insufficient, as this requires the types of U an Vt to be the same. For BlockArray objects, the type of axes is captured in the type parameters, such that if the rows and colums of A are not described by axes of the same type, this fails.

SVD struct

In order to circumvent this problem, I copied most of the LinearAlgebra.svd files, and created a new SVD struct that does not have this restriction, but otherwise functions identically. With this, I have new functions svd(!) that by default fall back to the LinearAlgebra implementation, but rewrap them in the new struct.
Some questions here:

• Do we want to overload LinearAlgebra.svd! instead of keeping them separate instead, only returning the custom SVD struct whenever required?
• Should I rename the SVD struct in order to avoid confusion with the LinearAlgebra-defined one?

BlockSparse (Diagonal) SVD

Expanding this to BlockSparseArrays, I first added an implementation for BlockDiagonal, which is just a BlockSparseMatrix with a Diagonal{T,<:AbstractMatrix{T}} storage type.
In this case, the SVD happens block-by-block, and as a result we have extra information about the columns of U and the rows of Vt, as they should retain this blocked structure.

julia> a = BlockDiagonal([rand(T, i, j) for (i, j) in zip(m, n)])
typeof(axes) = Tuple{BlockedOneTo{Int64, Vector{Int64}}, BlockedOneTo{Int64, Vector{Int64}}}

Warning: To temporarily circumvent ...

2×2-blocked 4×4 BlockSparseArray{Float64, 2, Matrix{Float64}, Diagonal{Matrix{Float64}, Vector{Matrix{Float64}}}, Tuple{BlockedOneTo{Int64, Vector{Int64}}, BlockedOneTo{Int64, Vector{Int64}}}}:
 0.583632  0.168925  │   ⋅         ⋅      
 0.911104  0.842136  │   ⋅         ⋅      
 ────────────────────┼────────────────────
  ⋅         ⋅        │  0.324859  0.471626
  ⋅         ⋅        │  0.188073  0.720221

julia> svd(a)
NDTensors.BlockSparseArrays.SVD{Float64, Float64, BlockSparseArray{Float64, 2, Matrix{Float64}, Diagonal{Matrix{Float64}, Vector{Matrix{Float64}}}, Tuple{BlockedOneTo{Int64, Vector{Int64}}, BlockedOneTo{Int64, Vector{Int64}}}}, BlockVector{Float64, Vector{Vector{Float64}}, Tuple{BlockedOneTo{Int64, Vector{Int64}}}}, BlockSparseArray{Float64, 2, Matrix{Float64}, Diagonal{Matrix{Float64}, Vector{Matrix{Float64}}}, Tuple{BlockedOneTo{Int64, Vector{Int64}}, BlockedOneTo{Int64, Vector{Int64}}}}}
U factor:
typeof(axes) = Tuple{BlockedOneTo{Int64, Vector{Int64}}, BlockedOneTo{Int64, Vector{Int64}}}

Warning: To temporarily circumvent ...

2×2-blocked 4×4 BlockSparseArray{Float64, 2, Matrix{Float64}, Diagonal{Matrix{Float64}, Vector{Matrix{Float64}}}, Tuple{BlockedOneTo{Int64, Vector{Int64}}, BlockedOneTo{Int64, Vector{Int64}}}}:
 -0.415014  -0.909815  │   ⋅          ⋅      
 -0.909815   0.415014  │   ⋅          ⋅      
 ──────────────────────┼─────────────────────
   ⋅          ⋅        │  0.603525   0.797344
   ⋅          ⋅        │  0.797344  -0.603525
singular values:
2-blocked 4-element BlockVector{Float64}:
 1.3589527927192564 
 0.24841857379494714
 ───────────────────
 0.9259816883073599 
 0.15688225526219635
Vt factor:
typeof(axes) = Tuple{BlockedOneTo{Int64, Vector{Int64}}, BlockedOneTo{Int64, Vector{Int64}}}

Warning: To temporarily circumvent ...

2×2-blocked 4×4 BlockSparseArray{Float64, 2, Matrix{Float64}, Diagonal{Matrix{Float64}, Vector{Matrix{Float64}}}, Tuple{BlockedOneTo{Int64, Vector{Int64}}, BlockedOneTo{Int64, Vector{Int64}}}}:
 -0.788218  -0.615396  │   ⋅          ⋅      
 -0.615396   0.788218  │   ⋅          ⋅      
 ──────────────────────┼─────────────────────
   ⋅          ⋅        │  0.373678   0.927558
   ⋅          ⋅        │  0.927558  -0.373678

Here, it might be beneficial to have S be either a BlockSparseVector or a BlockVector, depending on whether or not the full SVD is asked (not implemented yet)

BlockSparse SVD

Then, to also implement the generic BlockSparse case, I selected two possible implementations. Either the input is "permuted block-diagonal", as can arise in the context of tensors with nonzero flux, or it is a generic block sparse matrix.
There is a runtime check for this, (based on #1515) which permutes the array to become BlockDiagonal if possible, and otherwise falls back to a dense implementation.

Here, the questions mostly concern whether we want to have this runtime check. I would guess that this is something that can be statically inferred from the type of the axes, ie if the sparsity is due to symmetries, it has to be blockdiagonal in some form, while if not, it is presumably quite rare to have this specific structure. We could consider a trait to enable or disable this implementation.

Truncated SVD

Finally, in order to also include the truncation, I added an interface roughly based on the TensorKit implementation, which features a truncate(F::SVD) to truncate a given decomposition, as well as a tsvd(A) which simply combines the decomposition and truncation step. This second form could be useful in the future for algorithms that cannot decouple these two, such as Krylov methods.

The interface supports truncation by a number of different methods, based around keeping the n first values after sorting according to some criteria, as well as discarding values via a filtering mechanism. There are convenient constructors for the cases that are probably used most, but the system allows quite a lot of freedom. Additionally, because of the specialization on TruncationScheme, it is easily extensible in the future.

julia> tsvd(a)
NDTensors.BlockSparseArrays.SVD{Float64, Float64, BlockedMatrix{Float64, Matrix{Float64}, Tuple{BlockedOneTo{Int64, Vector{Int64}}, BlockedOneTo{Int64, Vector{Int64}}}}, BlockedVector{Float64, Vector{Float64}, Tuple{BlockedOneTo{Int64, Vector{Int64}}}}, BlockedMatrix{Float64, Matrix{Float64}, Tuple{BlockedOneTo{Int64, Vector{Int64}}, BlockedOneTo{Int64, Vector{Int64}}}}}
U factor:
2×1-blocked 4×4 BlockedMatrix{Float64}:
 -0.672324   0.299543   0.533151   0.417138
 -0.265515  -0.723181   0.435688  -0.465496
 ──────────────────────────────────────────
 -0.412697   0.490917  -0.221498  -0.734589
 -0.554223  -0.382471  -0.690554   0.263986
singular values:
1-blocked 4-element BlockedVector{Float64}:
 2.156295710384511  
 0.841059553579276  
 0.4119978358023982 
 0.17695709459403608
Vt factor:
1×2-blocked 4×4 BlockedMatrix{Float64}:
 -0.472017  -0.55809   │  -0.605812  -0.314209
  0.212759  -0.81694   │   0.404206   0.35208 
 -0.505012   0.143246  │  -0.17094    0.8338  
 -0.690577   0.025016  │   0.663618  -0.286513

julia> tsvd(a; trunc=truncbelow(0.2))
NDTensors.BlockSparseArrays.SVD{Float64, Float64, BlockedMatrix{Float64, Matrix{Float64}, Tuple{BlockedOneTo{Int64, Vector{Int64}}, Base.OneTo{Int64}}}, Vector{Float64}, BlockedMatrix{Float64, Matrix{Float64}, Tuple{Base.OneTo{Int64}, BlockedOneTo{Int64, Vector{Int64}}}}}
U factor:
2×1-blocked 4×3 BlockedMatrix{Float64, Matrix{Float64}, Tuple{BlockedOneTo{Int64, Vector{Int64}}, Base.OneTo{Int64}}}:
 -0.672324   0.299543   0.533151
 -0.265515  -0.723181   0.435688
 ───────────────────────────────
 -0.412697   0.490917  -0.221498
 -0.554223  -0.382471  -0.690554
singular values:
3-element Vector{Float64}:
 2.156295710384511
 0.841059553579276
 0.4119978358023982
Vt factor:
1×2-blocked 3×4 BlockedMatrix{Float64, Matrix{Float64}, Tuple{Base.OneTo{Int64}, BlockedOneTo{Int64, Vector{Int64}}}}:
 -0.472017  -0.55809   │  -0.605812  -0.314209
  0.212759  -0.81694   │   0.404206   0.35208 
 -0.505012   0.143246  │  -0.17094    0.8338

In order to make this work with BlockSparseArrays, I am missing the functionality to have U[:,[1,2,4]] slice the matrix but retain as much of the block structure as possible:

julia> a = BlockDiagonal([rand(T, i, j) for (i, j) in zip(m, n)])
typeof(axes) = Tuple{BlockedOneTo{Int64, Vector{Int64}}, 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 4×4 BlockSparseArray{Float64, 2, Matrix{Float64}, Diagonal{Matrix{Float64}, Vector{Matrix{Float64}}}, Tuple{BlockedOneTo{Int64, Vector{Int64}}, BlockedOneTo{Int64, Vector{Int64}}}}:
 0.545446  0.847428  │   ⋅         ⋅      
 0.214882  0.472071  │   ⋅         ⋅      
 ────────────────────┼────────────────────
  ⋅         ⋅        │  0.988776  0.313793
  ⋅         ⋅        │  0.394168  0.616615

julia> tsvd(a; trunc=truncbelow(0.2))
ERROR: Slicing not implemented for range of type `Vector{Int64}`.

As you see, much to discuss. I attached my preliminary test file, and added the code in places that I felt are relevant, but obviously I am completely open to suggestions, reorganisations, etc.

[mtfishman]

Consider changing to try_to_blockdiagonal as discussed above.

[mtfishman]

Also, does this handle rectangular block matrices, i.e. when the block pattern itself is rectangular?

[mtfishman]

I think this is the canonical way to convert Block(1, 1) to (1, 1), as an alternative to using field access:

Int.(Tuple(Block(1, 1)))

[mtfishman]

Though we could consider keeping the Block wrapper, i.e. Tuple(Block(1, 1)), and then the permutations are block permutations. Then below, instead of blocks(A)[tuple.(invperm(I), J)], it could be written as a slicing operation on the block sparse matrix A[...]. I slightly prefer that approach. In that case I would call the function something like try_to_blockdiagonal_blockperm to make it clear it is outputting a block permutation.

[lkdvos]

I played around with this, but within the current constraints I don't think I can make that work. I really need a vector of blocks in this case, in other words linear slicing into the block sparse matrix. A[vec_of_blockinds] just cannot work, for the reasons we discussed previously, because the slicing wants to return a BlockArray, which would then be a vector of matrices that all have different sizes.

My way around this is to perform linear index slicing into the blocks, which I think is fair to assume to return a vector of blocks (not wrapped into a BlockArray). I agree that it would be nicer to not have to do this, but I would also argue that it is not outside the allowed operations of the interface. blocks should return an AbstractArray containing the different blocks, and I should be allowed to index into that to obtain different blocks.

[mtfishman]

Hmm I'm not quite following, let's discuss this.

[lkdvos]

In this second iteration, I implemented most of the comments as you suggested.

There are two main problems that hold this back currently:

• BlockDiagonal * Adjoint and vice-versa don't work
• blocks(::AbstractBlockSparseArray)[vec_of_cartesianinds] does not work

I'll handle these first separately, and then continue working here. I'll also split off the tsvd stuff before this PR is ready to merge.