Fix Issue 162 - SVD with empty rows and cols

Project.toml

@@ -1,7 +1,7 @@
 name = "BlockSparseArrays"
 uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.8.1"
+version = "0.8.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/factorizations/svd.jl

@@ -53,47 +53,44 @@ function MatrixAlgebraKit.initialize_output(
 
   brows = eachblockaxis(axes(A, 1))
   bcols = eachblockaxis(axes(A, 2))
-  u_axes = similar(brows, bmn)
-  v_axes = similar(brows, bmn)
+  s_axes = similar(brows, bmn)
 
   # 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] = infimum(brows[row], bcols[col])
-    v_axes[col] = infimum(bcols[col], brows[row])
+    s_axes[I] = infimum(brows[row], 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)
   emptycols = setdiff(1:bn, bcolIs)
-  for (row, col) in zip(emptyrows, emptycols)
-    u_axes[col] = infimum(brows[row], bcols[col])
-    v_axes[col] = infimum(bcols[col], brows[row])
+  for (I, (row, col)) in enumerate(zip(emptyrows, emptycols))
+    s_axes[I + length(bIs)] = infimum(brows[row], bcols[col])
   end
 
-  u_axis = mortar_axis(u_axes)
-  v_axis = mortar_axis(v_axes)
-  S_axes = (u_axis, v_axis)
+  s_axis = mortar_axis(s_axes)
+  S_axes = (s_axis, s_axis)
   U, S, Vt = similar_output(svd_compact!, 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_compact!, 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, col)) in enumerate(zip(emptyrows, emptycols))
+    @view!(U[Block(row, I + length(bIs))])
+    @view!(Vt[Block(I + length(bIs), col)])
   end
 
   return U, S, Vt
@@ -115,53 +112,49 @@ 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(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)])
-  end
-  # also handle extra rows/cols
-  for i in (length(emptyrows) + 1):length(emptycols)
-    @view!(Vt[Block(emptycols[i], emptycols[i])])
+  for (I, row) in enumerate(emptyrows)
+    @view!(U[Block(row, length(bIs) + I)])
   end
-  for (i, k) in enumerate((length(emptycols) + 1):length(emptyrows))
-    @view!(U[Block(emptyrows[k], bn + i)])
+  for (I, col) in enumerate(emptycols)
+    @view!(Vt[Block(length(bIs) + I, col)])
   end
 
   return U, S, Vt
@@ -188,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
 
@@ -199,27 +191,28 @@ function MatrixAlgebraKit.svd_compact!(
   check_input(svd_compact!, 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_compact!(block, usvᴴ, block_alg)
     @assert usvᴴ === usvᴴ′ "svd_compact! 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
   # Vᴴ[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)
+  for (I, (row, col)) in enumerate(zip(emptyrows, emptycols))
+    copyto!(@view!(U[Block(row, I + length(bIs))]), LinearAlgebra.I)
+    copyto!(@view!(Vᴴ[Block(I + length(bIs), col)]), LinearAlgebra.I)
   end
 
   return (U, S, Vᴴ)
@@ -231,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

@@ -0,0 +1,25 @@
+using BlockArrays
+using BlockSparseArrays
+using BlockSparseArrays: blocksparse
+using MatrixAlgebraKit
+using LinearAlgebra: LinearAlgebra
+using Test: @test, @testset
+
+@testset "Issue 162" begin
+  ax = (blockedrange([2, 2]), blockedrange([2, 2, 2]))
+  bs = Dict(Block(1, 1) => randn(2, 2), Block(2, 3) => randn(2, 2))
+  a = blocksparse(bs, ax)
+  U, S, Vᴴ = svd_compact(a)
+
+  @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