Add support for QR decomposition (#37)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "GradedArrays"
 uuid = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.4.4"
+version = "0.4.5"
 
 [deps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
@@ -25,7 +25,7 @@ GradedArraysTensorAlgebraExt = "TensorAlgebra"
 
 [compat]
 BlockArrays = "1.6.0"
-BlockSparseArrays = "0.6.1"
+BlockSparseArrays = "0.6.2"
 Compat = "4.16.0"
 DerivableInterfaces = "0.4.4"
 FillArrays = "1.13.0"

src/factorizations.jl

@@ -8,7 +8,8 @@ using BlockSparseArrays:
   eachblockaxis,
   mortar_axis
 using LinearAlgebra: Diagonal
-using MatrixAlgebraKit: MatrixAlgebraKit, svd_compact!, svd_full!, svd_trunc!
+using MatrixAlgebraKit:
+  MatrixAlgebraKit, qr_compact!, qr_full!, svd_compact!, svd_full!, svd_trunc!
 
 function BlockSparseArrays.similar_output(
   ::typeof(svd_compact!), A::GradedMatrix, S_axes, alg::BlockPermutedDiagonalAlgorithm
@@ -56,3 +57,19 @@ function BlockSparseArrays.similar_truncate(
   Ṽᴴ = similar(Vᴴ, dual(v_axis), axes(Vᴴ, 2))
   return Ũ, S̃, Ṽᴴ
 end
+
+function BlockSparseArrays.similar_output(
+  ::typeof(qr_compact!), A::GradedMatrix, R_axis, alg::BlockPermutedDiagonalAlgorithm
+)
+  Q = similar(A, axes(A, 1), dual(R_axis))
+  R = similar(A, R_axis, axes(A, 2))
+  return Q, R
+end
+
+function BlockSparseArrays.similar_output(
+  ::typeof(qr_full!), A::GradedMatrix, R_axis, alg::BlockPermutedDiagonalAlgorithm
+)
+  Q = similar(A, axes(A, 1), dual(R_axis))
+  R = similar(A, R_axis, axes(A, 2))
+  return Q, R
+end

test/test_factorizations.jl

@@ -1,7 +1,7 @@
 using BlockArrays: Block, blocksizes
-using GradedArrays: U1, dual, flux, gradedrange
+using GradedArrays: U1, dual, flux, gradedrange, trivial
 using LinearAlgebra: I, diag, svdvals
-using MatrixAlgebraKit: svd_compact, svd_full, svd_trunc
+using MatrixAlgebraKit: qr_compact, qr_full, svd_compact, svd_full, svd_trunc
 using Test: @test, @testset
 
 const elts = (Float32, Float64, ComplexF32, ComplexF64)
@@ -17,9 +17,9 @@ const elts = (Float32, Float64, ComplexF32, ComplexF64)
     @test u * s * vᴴ ≈ a
     @test Array(u'u) ≈ I
     @test Array(vᴴ * vᴴ') ≈ I
-    @test flux(u) == U1(0)
+    @test flux(u) == trivial(flux(a))
     @test flux(s) == flux(a)
-    @test flux(vᴴ) == U1(0)
+    @test flux(vᴴ) == trivial(flux(a))
 
     r1 = gradedrange([U1(0) => i, U1(1) => j])
     r2 = gradedrange([U1(0) => k, U1(1) => l])
@@ -31,9 +31,9 @@ const elts = (Float32, Float64, ComplexF32, ComplexF64)
     @test u * s * vᴴ ≈ a
     @test Array(u'u) ≈ I
     @test Array(vᴴ * vᴴ') ≈ I
-    @test flux(u) == U1(0)
+    @test flux(u) == trivial(flux(a))
     @test flux(s) == flux(a)
-    @test flux(vᴴ) == U1(0)
+    @test flux(vᴴ) == trivial(flux(a))
   end
 end
 
@@ -50,9 +50,9 @@ end
     @test Array(u * u') ≈ I
     @test Array(vᴴ * vᴴ') ≈ I
     @test Array(vᴴ'vᴴ) ≈ I
-    @test flux(u) == U1(0)
+    @test flux(u) == trivial(flux(a))
     @test flux(s) == flux(a)
-    @test flux(vᴴ) == U1(0)
+    @test flux(vᴴ) == trivial(flux(a))
 
     r1 = gradedrange([U1(0) => i, U1(1) => j])
     r2 = gradedrange([U1(0) => k, U1(1) => l])
@@ -65,9 +65,9 @@ end
     @test Array(u * u') ≈ I
     @test Array(vᴴ * vᴴ') ≈ I
     @test Array(vᴴ'vᴴ) ≈ I
-    @test flux(u) == U1(0)
+    @test flux(u) == trivial(flux(a))
     @test flux(s) == flux(a)
-    @test flux(vᴴ) == U1(0)
+    @test flux(vᴴ) == trivial(flux(a))
   end
 end
 
@@ -85,9 +85,9 @@ end
     @test size(vᴴ) == (1, size(a, 2))
     @test Array(u'u) ≈ I
     @test Array(vᴴ * vᴴ') ≈ I
-    @test flux(u) == U1(0)
+    @test flux(u) == trivial(flux(a))
     @test flux(s) == flux(a)
-    @test flux(vᴴ) == U1(0)
+    @test flux(vᴴ) == trivial(flux(a))
 
     r1 = gradedrange([U1(0) => i, U1(1) => j])
     r2 = gradedrange([U1(0) => k, U1(1) => l])
@@ -101,8 +101,62 @@ end
     @test size(vᴴ) == (1, size(a, 2))
     @test Array(u'u) ≈ I
     @test Array(vᴴ * vᴴ') ≈ I
-    @test flux(u) == U1(0)
+    @test flux(u) == trivial(flux(a))
     @test flux(s) == flux(a)
-    @test flux(vᴴ) == U1(0)
+    @test flux(vᴴ) == trivial(flux(a))
+  end
+end
+
+@testset "qr_compact (eltype=$elt)" for elt in elts
+  for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3]
+    r1 = gradedrange([U1(0) => i, U1(1) => j])
+    r2 = gradedrange([U1(0) => k, U1(1) => l])
+    a = zeros(elt, r1, dual(r2))
+    a[Block(2, 2)] = randn(elt, blocksizes(a)[2, 2])
+    @test flux(a) == U1(0)
+    q, r = qr_compact(a)
+    @test q * r ≈ a
+    @test Array(q'q) ≈ I
+    @test flux(q) == trivial(flux(a))
+    @test flux(r) == flux(a)
+
+    r1 = gradedrange([U1(0) => i, U1(1) => j])
+    r2 = gradedrange([U1(0) => k, U1(1) => l])
+    a = zeros(elt, r1, dual(r2))
+    a[Block(1, 2)] = randn(elt, blocksizes(a)[1, 2])
+    @test flux(a) == U1(-1)
+    q, r = qr_compact(a)
+    @test q * r ≈ a
+    @test Array(q'q) ≈ I
+    @test flux(q) == trivial(flux(a))
+    @test flux(r) == flux(a)
+  end
+end
+
+@testset "qr_full (eltype=$elt)" for elt in elts
+  for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3]
+    r1 = gradedrange([U1(0) => i, U1(1) => j])
+    r2 = gradedrange([U1(0) => k, U1(1) => l])
+    a = zeros(elt, r1, dual(r2))
+    a[Block(2, 2)] = randn(elt, blocksizes(a)[2, 2])
+    @test flux(a) == U1(0)
+    q, r = qr_full(a)
+    @test q * r ≈ a
+    @test Array(q'q) ≈ I
+    @test Array(q * q') ≈ I
+    @test flux(q) == trivial(flux(a))
+    @test flux(r) == flux(a)
+
+    r1 = gradedrange([U1(0) => i, U1(1) => j])
+    r2 = gradedrange([U1(0) => k, U1(1) => l])
+    a = zeros(elt, r1, dual(r2))
+    a[Block(1, 2)] = randn(elt, blocksizes(a)[1, 2])
+    @test flux(a) == U1(-1)
+    q, r = qr_full(a)
+    @test q * r ≈ a
+    @test Array(q'q) ≈ I
+    @test Array(q * q') ≈ I
+    @test flux(q) == trivial(flux(a))
+    @test flux(r) == flux(a)
   end
 end