Implement blockwise polar decomposition (#125)

Author: mtfishman

Project.toml

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

src/BlockSparseArrays.jl

@@ -48,5 +48,6 @@ include("factorizations/svd.jl")
 include("factorizations/truncation.jl")
 include("factorizations/qr.jl")
 include("factorizations/lq.jl")
+include("factorizations/polar.jl")
 
 end

src/factorizations/polar.jl

@@ -0,0 +1,59 @@
+using MatrixAlgebraKit:
+  MatrixAlgebraKit,
+  PolarViaSVD,
+  check_input,
+  default_algorithm,
+  left_polar!,
+  right_polar!,
+  svd_compact!
+
+function MatrixAlgebraKit.check_input(::typeof(left_polar!), A::AbstractBlockSparseMatrix)
+  @views for I in eachblockstoredindex(A)
+    m, n = size(A[I])
+    m >= n ||
+      throw(ArgumentError("each input matrix block needs at least as many rows as columns"))
+  end
+  return nothing
+end
+function MatrixAlgebraKit.check_input(::typeof(right_polar!), A::AbstractBlockSparseMatrix)
+  @views for I in eachblockstoredindex(A)
+    m, n = size(A[I])
+    m <= n ||
+      throw(ArgumentError("each input matrix block needs at least as many columns as rows"))
+  end
+  return nothing
+end
+
+function MatrixAlgebraKit.left_polar!(A::AbstractBlockSparseMatrix, alg::PolarViaSVD)
+  check_input(left_polar!, A)
+  # TODO: Use more in-place operations here, avoid `copy`.
+  U, S, Vᴴ = svd_compact!(A, alg.svdalg)
+  W = U * Vᴴ
+  # TODO: `copy` is required for now because of:
+  # https://github.com/ITensor/BlockSparseArrays.jl/issues/24
+  # Remove when that is fixed.
+  P = copy(Vᴴ') * S * Vᴴ
+  return (W, P)
+end
+function MatrixAlgebraKit.right_polar!(A::AbstractBlockSparseMatrix, alg::PolarViaSVD)
+  check_input(right_polar!, A)
+  # TODO: Use more in-place operations here, avoid `copy`.
+  U, S, Vᴴ = svd_compact!(A, alg.svdalg)
+  Wᴴ = U * Vᴴ
+  # TODO: `copy` is required for now because of:
+  # https://github.com/ITensor/BlockSparseArrays.jl/issues/24
+  # Remove when that is fixed.
+  P = U * S * copy(U')
+  return (P, Wᴴ)
+end
+
+function MatrixAlgebraKit.default_algorithm(
+  ::typeof(left_polar!), a::AbstractBlockSparseMatrix; kwargs...
+)
+  return PolarViaSVD(default_algorithm(svd_compact!, a; kwargs...))
+end
+function MatrixAlgebraKit.default_algorithm(
+  ::typeof(right_polar!), a::AbstractBlockSparseMatrix; kwargs...
+)
+  return PolarViaSVD(default_algorithm(svd_compact!, a; kwargs...))
+end

test/test_factorizations.jl

@@ -1,10 +1,12 @@
 using BlockArrays: Block, BlockedMatrix, BlockedVector, blocks, mortar
 using BlockSparseArrays: BlockSparseArray, BlockDiagonal, eachblockstoredindex
 using MatrixAlgebraKit:
+  left_polar,
   lq_compact,
   lq_full,
   qr_compact,
   qr_full,
+  right_polar,
   svd_compact,
   svd_full,
   svd_trunc,
@@ -215,3 +217,23 @@ end
     @test A ≈ L * Q
   end
 end
+
+@testset "left_polar (T=$T)" for T in (Float32, Float64, ComplexF32, ComplexF64)
+  A = BlockSparseArray{T}(undef, ([3, 4], [2, 3]))
+  A[Block(1, 1)] = randn(T, 3, 2)
+  A[Block(2, 2)] = randn(T, 4, 3)
+
+  U, C = left_polar(A)
+  @test U * C ≈ A
+  @test Matrix(U'U) ≈ LinearAlgebra.I
+end
+
+@testset "right_polar (T=$T)" for T in (Float32, Float64, ComplexF32, ComplexF64)
+  A = BlockSparseArray{T}(undef, ([2, 3], [3, 4]))
+  A[Block(1, 1)] = randn(T, 2, 3)
+  A[Block(2, 2)] = randn(T, 3, 4)
+
+  C, U = right_polar(A)
+  @test C * U ≈ A
+  @test Matrix(U * U') ≈ LinearAlgebra.I
+end