catch non-full sparse matrices and simplify diagonal implementation

Author: lkdvos

Project.toml

@@ -8,6 +8,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 DerivableInterfaces = "6c5e35bf-e59e-4898-b73c-732dcc4ba65f"
+DiagonalArrays = "74fd4be6-21e2-4f6f-823a-4360d37c7a77"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 GradedUnitRanges = "e2de450a-8a67-46c7-b59c-01d5a3d041c5"
@@ -32,6 +33,7 @@ Aqua = "0.8.9"
 ArrayLayouts = "1.10.4"
 BlockArrays = "1.2.0"
 DerivableInterfaces = "0.3.7"
+DiagonalArrays = "0.2.2"
 Dictionaries = "0.4.3"
 FillArrays = "1.13.0"
 GPUArraysCore = "0.1.0, 0.2"

src/abstractblocksparsearray/abstractblocksparsematrix.jl

@@ -43,8 +43,8 @@ function _allocate_svd_output(A::AbstractBlockSparseMatrix, full::Bool, ::Algori
 
   # 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)
+  emptyrows = findall(∉(browIs), 1:bm)
+  emptycols = findall(∉(bcolIs), 1:bn)
   for (row, col) in zip(emptyrows, emptycols)
     slengths[col] = min(brows[row], bcols[col])
   end
@@ -73,5 +73,20 @@ function svd!(
     S[bcol] = bUSV.S
     Vt[bcol, bcol] = bUSV.Vt
   end
+
+  # fill in values for blocks that aren't present, pairing them in order of occurence
+  block_inds_S = eachblockstoredindex(S)
+  i = findfirst(∉(block_inds_S), blockaxes(S, 1))
+  bIs = collect(eachblockstoredindex(U))
+  browIs = Int.(first.(Tuple.(bIs)))
+  emptyrows = findall(∉(browIs), 1:blocksize(U, 1))
+  j = 0
+  while !isnothing(i)
+    copyto!(@view!(Vt[Block(i), Block(i)]), LinearAlgebra.I)
+    j += 1
+    copyto!(@view!(U[Block(emptyrows[j]), Block(i)]), LinearAlgebra.I)
+    i = findnext(∉(block_inds_S), blockaxes(S, 1), i + 1)
+  end
+
   return SVD(U, S, Vt)
 end

src/blocksparsearray/blockdiagonalarray.jl

@@ -1,5 +1,6 @@
 # type alias for block-diagonal
 using LinearAlgebra: Diagonal
+using DiagonalArrays: DiagonalArrays, diagonal
 
 const BlockDiagonal{T,A,Axes,V<:AbstractVector{A}} = BlockSparseMatrix{
   T,A,Diagonal{A,V},Axes
@@ -15,3 +16,11 @@ function BlockDiagonal(blocks::AbstractVector{<:AbstractMatrix})
     Diagonal(blocks), (blockedrange(size.(blocks, 1)), blockedrange(size.(blocks, 2)))
   )
 end
+
+function DiagonalArrays.diagonal(S::BlockSparseVector)
+  D = similar(S, (axes(S, 1), axes(S, 1)))
+  for bI in eachblockstoredindex(S)
+    D[bI, bI] = diagonal(@view!(S[bI]))
+  end
+  return D
+end

test/Project.toml

@@ -4,6 +4,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 BlockSparseArrays = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
+DiagonalArrays = "74fd4be6-21e2-4f6f-823a-4360d37c7a77"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 GradedUnitRanges = "e2de450a-8a67-46c7-b59c-01d5a3d041c5"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"

test/basics/test_svd.jl

@@ -1,14 +1,15 @@
 using Test
 using BlockSparseArrays
-using BlockSparseArrays: BlockSparseArray, svd, BlockDiagonal
+using BlockSparseArrays: BlockSparseArray, svd, BlockDiagonal, eachblockstoredindex
 using BlockArrays
-using LinearAlgebra: LinearAlgebra, Diagonal, svdvals
 using Random
+using DiagonalArrays: diagonal
+using LinearAlgebra: LinearAlgebra
 
 function test_svd(a, usv)
   U, S, V = usv
 
-  @test U * Diagonal(S) * V' ≈ a
+  @test U * diagonal(S) * V' ≈ a
   @test U' * U ≈ LinearAlgebra.I
   @test V' * V ≈ LinearAlgebra.I
 end
@@ -43,33 +44,35 @@ end
   usv = svd(a)
   # TODO: `BlockDiagonal * Adjoint` errors
   test_svd(a, usv)
-  @test usv.U isa BlockDiagonal
-  @test usv.Vt isa BlockDiagonal
-  @test usv.S isa BlockVector
 end
 
-a = mortar([rand(2, 2) for i in 1:2, j in 1:3])
-usv = svd(a)
-test_svd(a, usv)
-
-a = mortar([rand(2, 2) for i in 1:3, j in 1:2])
-usv = svd(a)
-test_svd(a, usv)
-
 # blocksparse 
 # -----------
 @testset "($m, $n) BlockSparseMatrix{$T}" for ((m, n), T) in
                                               Iterators.product(blockszs, eltypes)
   a = BlockSparseArray{T}(m, n)
+
+  # test empty matrix
+  usv_empty = svd(a)
+  test_svd(a, usv_empty)
+
+  # test blockdiagonal
   for i in LinearAlgebra.diagind(blocks(a))
     I = CartesianIndices(blocks(a))[i]
     a[Block(I.I...)] = rand(T, size(blocks(a)[i]))
   end
-  perm = Random.randperm(length(m))
-  a = a[Block.(perm), Block.(1:length(n))]
-
-  # errors because `blocks(a)[CartesianIndex.(...)]` is not implemented
   usv = svd(a)
-  # TODO: `BlockDiagonal * Adjoint` errors
   test_svd(a, usv)
+
+  perm = Random.randperm(length(m))
+  b = a[Block.(perm), Block.(1:length(n))]
+  usv = svd(b)
+  test_svd(b, usv)
+
+  # test permuted blockdiagonal with missing row/col
+  I_removed = rand(eachblockstoredindex(b))
+  c = copy(b)
+  delete!(blocks(c).storage, CartesianIndex(Int.(Tuple(I_removed))))
+  usv = svd(c)
+  test_svd(c, usv)
 end