also update svd_full

Author: lkdvos

src/factorizations/svd.jl

@@ -112,55 +112,51 @@ function MatrixAlgebraKit.initialize_output(
   bm, bn = blocksize(A)
 
   brows = eachblockaxis(axes(A, 1))
-  u_axes = similar(brows)
+  bcols = eachblockaxis(axes(A, 2))
+  u_axes = similar(brows, bm)
+  v_axes = similar(bcols, bn)
 
   # fill in values for blocks that are present
-  bIs = collect(eachblockstoredindex(A))
+  bIs = sort!(collect(eachblockstoredindex(A)), by=Int ∘ last ∘ Tuple)
   browIs = Int.(first.(Tuple.(bIs)))
   bcolIs = Int.(last.(Tuple.(bIs)))
-  for bI in eachblockstoredindex(A)
+  for (I, bI) in enumerate(bIs)
     row, col = Int.(Tuple(bI))
-    u_axes[col] = brows[row]
+    u_axes[I] = brows[row]
+    v_axes[I] = bcols[col]
   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)
+  u_axes[length(bIs) .+ (1:length(emptyrows))] .= brows[emptyrows]
   emptycols = setdiff(1:bn, bcolIs)
-  for (row, col) in zip(emptyrows, emptycols)
-    u_axes[col] = brows[row]
-  end
-  for (i, k) in enumerate((length(emptycols) + 1):length(emptyrows))
-    u_axes[bn + i] = brows[emptyrows[k]]
-  end
-
+  v_axes[length(bIs) .+ (1:length(emptycols))] .= bcols[emptycols]
+  
   u_axis = mortar_axis(u_axes)
-  S_axes = (u_axis, axes(A, 2))
+  v_axis = mortar_axis(@show v_axes)
+  S_axes = (u_axis, v_axis)
   U, S, Vt = similar_output(svd_full!, A, S_axes, alg)
 
   # allocate output
-  for bI in eachblockstoredindex(A)
+  for (I, bI) in enumerate(bIs)
     brow, bcol = Tuple(bI)
+    bcol′ = Block(I)
     block = @view!(A[bI])
     block_alg = block_algorithm(alg, block)
-    U[brow, bcol], S[bcol, bcol], Vt[bcol, bcol] = MatrixAlgebraKit.initialize_output(
+    U[brow, bcol′], S[bcol′, bcol′], Vt[bcol′, bcol] = MatrixAlgebraKit.initialize_output(
       svd_full!, block, block_alg
     )
   end
 
   # allocate output for blocks that aren't present -- do we also fill identities here?
-  for (row, col) in zip(emptyrows, emptycols)
-    @view!(U[Block(row, col)])
-    @view!(Vt[Block(col, col)])
+  for (I, row) in enumerate(emptyrows)
+    @view!(U[Block(row, I)])
   end
-  # also handle extra rows/cols
-  for i in (length(emptyrows) + 1):length(emptycols)
-    @view!(Vt[Block(emptycols[i], emptycols[i])])
+  for (I, col) in enumerate(emptycols)
+    @view!(Vt[Block(I, col)])
   end
-  for (i, k) in enumerate((length(emptycols) + 1):length(emptyrows))
-    @view!(U[Block(emptyrows[k], bn + i)])
-  end
-
+  
   return U, S, Vt
 end
 
@@ -185,8 +181,7 @@ function MatrixAlgebraKit.check_input(
     isa(Vᴴ, AbstractBlockSparseMatrix)
   @assert eltype(A) == eltype(U) == eltype(Vᴴ)
   @assert real(eltype(A)) == eltype(S)
-  @assert axes(A, 1) == axes(U, 1) && axes(A, 2) == axes(Vᴴ, 1) == axes(Vᴴ, 2)
-  @assert axes(S, 2) == axes(A, 2)
+  @assert axes(A, 1) == axes(U, 1) && axes(A, 2) == axes(Vᴴ, 2)
   return nothing
 end
 
@@ -208,7 +203,6 @@ function MatrixAlgebraKit.svd_compact!(
   end
 
   # fill in identities for blocks that aren't present
-  bIs = collect(eachblockstoredindex(A))
   browIs = Int.(first.(Tuple.(bIs)))
   bcolIs = Int.(last.(Tuple.(bIs)))
   emptyrows = setdiff(1:blocksize(A, 1), browIs)
@@ -230,36 +224,30 @@ function MatrixAlgebraKit.svd_full!(
   check_input(svd_full!, A, (U, S, Vᴴ))
 
   # do decomposition on each block
-  for bI in eachblockstoredindex(A)
+  bIs = sort!(collect(eachblockstoredindex(A)); by=Int ∘ last ∘ Tuple)
+  for (I, bI) in enumerate(bIs)
     brow, bcol = Tuple(bI)
-    usvᴴ = (@view!(U[brow, bcol]), @view!(S[bcol, bcol]), @view!(Vᴴ[bcol, bcol]))
+    bcol′ = Block(I)
+    usvᴴ = (@view!(U[brow, bcol′]), @view!(S[bcol′, bcol′]), @view!(Vᴴ[bcol′, bcol]))
     block = @view!(A[bI])
     block_alg = block_algorithm(alg, block)
     usvᴴ′ = svd_full!(block, usvᴴ, block_alg)
     @assert usvᴴ === usvᴴ′ "svd_full! might not be in-place"
   end
 
   # fill in identities for blocks that aren't present
-  bIs = collect(eachblockstoredindex(A))
   browIs = Int.(first.(Tuple.(bIs)))
   bcolIs = Int.(last.(Tuple.(bIs)))
   emptyrows = setdiff(1:blocksize(A, 1), browIs)
   emptycols = setdiff(1:blocksize(A, 2), bcolIs)
   # needs copyto! instead because size(::LinearAlgebra.I) doesn't work
   # U[Block(row, col)] = LinearAlgebra.I
   # Vt[Block(col, col)] = LinearAlgebra.I
-  for (row, col) in zip(emptyrows, emptycols)
-    copyto!(@view!(U[Block(row, col)]), LinearAlgebra.I)
-    copyto!(@view!(Vᴴ[Block(col, col)]), LinearAlgebra.I)
-  end
-
-  # also handle extra rows/cols
-  for i in (length(emptyrows) + 1):length(emptycols)
-    copyto!(@view!(Vᴴ[Block(emptycols[i], emptycols[i])]), LinearAlgebra.I)
+  for (I, row) in enumerate(emptyrows)
+    copyto!(@view!(U[Block(row, length(bIs) + I)]), LinearAlgebra.I)
   end
-  bn = blocksize(A, 2)
-  for (i, k) in enumerate((length(emptycols) + 1):length(emptyrows))
-    copyto!(@view!(U[Block(emptyrows[k], bn + i)]), LinearAlgebra.I)
+  for (I, col) in enumerate(emptycols)
+    copyto!(@view!(Vᴴ[Block(length(bIs) + I, col)]), LinearAlgebra.I)
   end
 
   return (U, S, Vᴴ)

test/test_issues.jl

@@ -14,4 +14,12 @@ using Test: @test, @testset
   @test U * S * Vᴴ ≈ a
   @test U' * U ≈ LinearAlgebra.I
   @test Vᴴ * Vᴴ' ≈ LinearAlgebra.I
+  
+  U, S, Vᴴ = svd_full(a);
+
+  @test U * S * Vᴴ ≈ a
+  @test U' * U ≈ LinearAlgebra.I
+  @test U * U' ≈ LinearAlgebra.I
+  @test Vᴴ * Vᴴ' ≈ LinearAlgebra.I
+  @test Vᴴ' * Vᴴ ≈ LinearAlgebra.I
 end