Rewrite svd to work with rectangular matrices

Author: lkdvos

src/abstractblocksparsearray/abstractblocksparsematrix.jl

@@ -3,20 +3,74 @@ const AbstractBlockSparseMatrix{T} = AbstractBlockSparseArray{T,2}
 # SVD is implemented by trying to
 # 1. Attempt to find a block-diagonal implementation by permuting
 # 2. Fallback to AbstractBlockArray implementation via BlockedArray
-function svd(
-  A::AbstractBlockSparseMatrix; full::Bool=false, alg::Algorithm=default_svd_alg(A)
-)
-  T = LinearAlgebra.eigtype(eltype(A))
-  A′_and_blockperms = try_to_blockdiagonal(A)
 
-  if isnothing(A′_and_blockperms)
-    # not block-diagonal, fall back to dense case
-    Adense = eigencopy_oftype(A, T)
-    return svd!(Adense; full, alg)
+function eigencopy_oftype(A::AbstractBlockSparseMatrix, T)
+  if is_block_permutation_matrix(A)
+    Acopy = similar(A, T)
+    for bI in eachblockstoredindex(A)
+      Acopy[bI] = eigencopy_oftype(A[bI], T)
+    end
+    return Acopy
+  else
+    return BlockedMatrix{T}(A)
+  end
+end
+
+function is_block_permutation_matrix(a::AbstractBlockSparseMatrix)
+  return allunique(first ∘ Tuple, eachblockstoredindex(a)) &&
+         allunique(last ∘ Tuple, eachblockstoredindex(a))
+end
+
+function _allocate_svd_output(A::AbstractBlockSparseMatrix, full::Bool, ::Algorithm)
+  @assert !full "TODO"
+  bm, bn = blocksize(A)
+  bmn = min(bm, bn)
+
+  brows = blocklengths(axes(A, 1))
+  bcols = blocklengths(axes(A, 2))
+  slengths = Vector{Int}(undef, bmn)
+
+  # fill in values for blocks that are present
+  bIs = collect(eachblockstoredindex(A))
+  browIs = Int.(first.(Tuple.(bIs)))
+  bcolIs = Int.(last.(Tuple.(bIs)))
+  for bI in eachblockstoredindex(A)
+    row, col = Int.(Tuple(bI))
+    nrows = brows[row]
+    ncols = bcols[col]
+    slengths[col] = min(nrows, ncols)
   end
 
-  # compute block-by-block and permute back
-  A″, (I, J) = A′
-  F = svd!(eigencopy_oftype(A″, T); full, alg)
-  return SVD(F.U[I, J], F.S, F.Vt)
+  # fill in values for blocks that aren't present, pairing them in order of occurence
+  # this is a convention, which at least gives the expected results for blockdiagonal
+  emptyrows = findall(∉(browIs), 1:bmn)
+  emptycols = findall(∉(bcolIs), 1:bmn)
+  for (row, col) in zip(emptyrows, emptycols)
+    slengths[col] = min(brows[row], bcols[col])
+  end
+
+  U = similar(A, axes(A, 1), blockedrange(slengths))
+  S = similar(A, real(eltype(A)), blockedrange(slengths))
+  Vt = similar(A, blockedrange(slengths), axes(A, 2))
+
+  return U, S, Vt
+end
+
+function svd(A::AbstractBlockSparseMatrix; kwargs...)
+  return svd!(eigencopy_oftype(A, LinearAlgebra.eigtype(eltype(A))); kwargs...)
+end
+
+function svd!(
+  A::AbstractBlockSparseMatrix; full::Bool=false, alg::Algorithm=default_svd_alg(A)
+)
+  @assert is_block_permutation_matrix(A) "Cannot keep sparsity: use `svd` to convert to `BlockedMatrix"
+  U, S, Vt = _allocate_svd_output(A, full, alg)
+  for bI in eachblockstoredindex(A)
+    bUSV = svd!(A[bI]; full, alg)
+    brow, bcol = Int.(Tuple(bI))
+    U[Block(brow, bcol)] = bUSV.U
+    S[Block(bcol)] = bUSV.S
+    Vt[Block(bcol, bcol)] = bUSV.Vt
+  end
+  return SVD(U, S, Vt)
 end

src/blocksparsearray/blockdiagonalarray.jl

@@ -11,38 +11,9 @@ function BlockDiagonal(blocks::AbstractVector{<:AbstractMatrix})
   )
 end
 
-# Cast to block-diagonal implementation if permuted-blockdiagonal
-function try_to_blockdiagonal_perm(A)
-  inds = map(x -> Int.(Tuple(x)), vec(collect(block_stored_indices(A))))
-  I = first.(inds)
-  allunique(I) || return nothing
-  J = last.(inds)
-  p = sortperm(J)
-  Jsorted = J[p]
-  allunique(Jsorted) || return nothing
-  return Block.(I[p], Jsorted)
-end
-
-"""
-    try_to_blockdiagonal(A)
-
-Attempt to find a permutation of blocks that makes `A` blockdiagonal. If unsuccesful,
-returns nothing, otherwise returns both the blockdiagonal `B` as well as the permutation `I, J`.
-"""
-function try_to_blockdiagonal(A::AbstractBlockSparseMatrix)
-  perm = try_to_blockdiagonal_perm(A)
-  isnothing(perm) && return perm
-  I = first.(Tuple.(perm))
-  J = last.(Tuple.(perm))
-  diagblocks = map(invperm(I), J) do i, j
-    return A[Block(i, j)]
-  end
-  return BlockDiagonal(diagblocks), perm
-end
-
 # SVD implementation
-function eigencopy_oftype(A::BlockDiagonal, S)
-  diag = map(Base.Fix2(eigencopy_oftype, S), A.blocks.diag)
+function eigencopy_oftype(A::BlockDiagonal, T)
+  diag = map(Base.Fix2(eigencopy_oftype, T), A.blocks.diag)
   return BlockDiagonal(diag)
 end