[WIP] Make SVD more general to accommodate graded arrays

Author: mtfishman

src/BlockArraysExtensions/blockedunitrange.jl

@@ -9,11 +9,22 @@ using BlockArrays:
   BlockSlice,
   BlockVector,
   block,
+  blockedrange,
   blockindex,
+  blocklengths,
   findblock,
   findblockindex,
   mortar
 
+function blockaxeses(a::AbstractUnitRange)
+  Base.require_one_based_indexing(a)
+  return map(Base.OneTo, blocklengths(a))
+end
+
+function mortar_axis(axeses)
+  return blockedrange(length.(axeses))
+end
+
 # Custom `BlockedUnitRange` constructor that takes a unit range
 # and a set of block lengths, similar to `BlockArray(::AbstractArray, blocklengths...)`.
 function blockedunitrange(a::AbstractUnitRange, blocklengths)

src/factorizations/svd.jl

@@ -25,7 +25,16 @@ end
 # the blocktype and element type - something like S = similar(A, BlockType(...))
 function _similar_S(A::AbstractBlockSparseMatrix, s_axis)
   T = real(eltype(A))
-  return BlockSparseArray{T,2,Diagonal{T,Vector{T}}}(undef, (s_axis, s_axis))
+  return BlockSparseMatrix{T,Diagonal{T,Vector{T}}}(undef, (s_axis, s_axis))
+end
+
+function similar_output(
+  ::typeof(svd_compact!), A, s_axis::AbstractUnitRange, alg::MatrixAlgebraKit.Algorithm
+)
+  U = similar(A, axes(A, 1), s_axis)
+  S = _similar_S(A, s_axis)
+  Vt = similar(A, s_axis, axes(A, 2))
+  return U, S, Vt
 end
 
 function MatrixAlgebraKit.initialize_output(
@@ -34,9 +43,9 @@ function MatrixAlgebraKit.initialize_output(
   bm, bn = blocksize(A)
   bmn = min(bm, bn)
 
-  brows = blocklengths(axes(A, 1))
-  bcols = blocklengths(axes(A, 2))
-  slengths = Vector{Int}(undef, bmn)
+  brows = blockaxeses(axes(A, 1))
+  bcols = blockaxeses(axes(A, 2))
+  s_axeses = Vector{eltype(brows)}(undef, bmn)
 
   # fill in values for blocks that are present
   bIs = collect(eachblockstoredindex(A))
@@ -46,21 +55,19 @@ function MatrixAlgebraKit.initialize_output(
     row, col = Int.(Tuple(bI))
     nrows = brows[row]
     ncols = bcols[col]
-    slengths[col] = min(nrows, ncols)
+    s_axeses[col] = min(nrows, ncols)
   end
 
   # 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 = setdiff(1:bm, browIs)
   emptycols = setdiff(1:bn, bcolIs)
   for (row, col) in zip(emptyrows, emptycols)
-    slengths[col] = min(brows[row], bcols[col])
+    s_axeses[col] = min(brows[row], bcols[col])
   end
 
-  s_axis = blockedrange(slengths)
-  U = similar(A, axes(A, 1), s_axis)
-  S = _similar_S(A, s_axis)
-  Vt = similar(A, s_axis, axes(A, 2))
+  s_axis = mortar_axis(s_axeses)
+  U, S, Vt = similar_output(svd_compact!, A, s_axis, alg)
 
   # allocate output
   for bI in eachblockstoredindex(A)