More factorizations (polar, orth, etc.) (#74)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "NamedDimsArrays"
 uuid = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.6.1"
+version = "0.6.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -36,6 +36,6 @@ LinearAlgebra = "1.10"
 MapBroadcast = "0.1.6"
 Random = "1.10"
 SimpleTraits = "0.9.4"
-TensorAlgebra = "0.2"
+TensorAlgebra = "0.2.9"
 TypeParameterAccessors = "0.2, 0.3"
 julia = "1.10"

src/tensoralgebra.jl

@@ -6,12 +6,17 @@ using TensorAlgebra:
   contract!,
   eigen,
   eigvals,
+  factorize,
   fusedims,
   left_null,
+  left_orth,
+  left_polar,
   lq,
   permmortar,
   qr,
   right_null,
+  right_orth,
+  right_polar,
   splitdims,
   svd,
   svdvals
@@ -133,55 +138,57 @@ function TensorAlgebra.splitdims(na::AbstractNamedDimsArray, splitters::Pair...)
   return nameddimsarray(a_split, names_split)
 end
 
-function TensorAlgebra.qr(
-  a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs...
+# Generic interface for forwarding binary factorizations
+# to the corresponding functions in TensorAlgebra.jl.
+function factorize_with(
+  f, a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs...
 )
   codomain = to_nameddimsindices(a, dimnames_codomain)
   domain = to_nameddimsindices(a, dimnames_domain)
-  q_unnamed, r_unnamed = qr(dename(a), nameddimsindices(a), codomain, domain; kwargs...)
-  name_q = randname(dimnames(a, 1))
-  name_r = name_q
-  namedindices_q = named(last(axes(q_unnamed)), name_q)
-  namedindices_r = named(first(axes(r_unnamed)), name_r)
-  nameddimsindices_q = (codomain..., namedindices_q)
-  nameddimsindices_r = (namedindices_r, domain...)
-  q = nameddimsarray(q_unnamed, nameddimsindices_q)
-  r = nameddimsarray(r_unnamed, nameddimsindices_r)
-  return q, r
+  x_unnamed, y_unnamed = f(dename(a), nameddimsindices(a), codomain, domain; kwargs...)
+  name_x = randname(dimnames(a, 1))
+  name_y = name_x
+  namedindices_x = named(last(axes(x_unnamed)), name_x)
+  namedindices_y = named(first(axes(y_unnamed)), name_y)
+  nameddimsindices_x = (codomain..., namedindices_x)
+  nameddimsindices_y = (namedindices_y, domain...)
+  x = nameddimsarray(x_unnamed, nameddimsindices_x)
+  y = nameddimsarray(y_unnamed, nameddimsindices_y)
+  return x, y
 end
-function TensorAlgebra.qr(a::AbstractNamedDimsArray, dimnames_codomain; kwargs...)
+function factorize_with(f, a::AbstractNamedDimsArray, dimnames_codomain; kwargs...)
   codomain = to_nameddimsindices(a, dimnames_codomain)
   domain = setdiff(nameddimsindices(a), codomain)
-  return qr(a, codomain, domain; kwargs...)
-end
-function LinearAlgebra.qr(a::AbstractNamedDimsArray, args...; kwargs...)
-  return TensorAlgebra.qr(a, args...; kwargs...)
+  return factorize_with(f, a, codomain, domain; kwargs...)
 end
 
-function TensorAlgebra.lq(
-  a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs...
-)
-  codomain = to_nameddimsindices(a, dimnames_codomain)
-  domain = to_nameddimsindices(a, dimnames_domain)
-  l_unnamed, q_unnamed = lq(dename(a), nameddimsindices(a), codomain, domain; kwargs...)
-  name_l = randname(dimnames(a, 1))
-  name_q = name_l
-  namedindices_l = named(last(axes(l_unnamed)), name_l)
-  namedindices_q = named(first(axes(q_unnamed)), name_q)
-  nameddimsindices_l = (codomain..., namedindices_l)
-  nameddimsindices_q = (namedindices_q, domain...)
-  l = nameddimsarray(l_unnamed, nameddimsindices_l)
-  q = nameddimsarray(q_unnamed, nameddimsindices_q)
-  return l, q
+for f in [:qr, :lq, :left_polar, :right_polar, :left_orth, :right_orth, :factorize]
+  @eval begin
+    function TensorAlgebra.$f(
+      a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs...
+    )
+      return factorize_with($f, a, dimnames_codomain, dimnames_domain; kwargs...)
+    end
+    function TensorAlgebra.$f(a::AbstractNamedDimsArray, dimnames_codomain; kwargs...)
+      return factorize_with($f, a, dimnames_codomain; kwargs...)
+    end
+  end
 end
-function TensorAlgebra.lq(a::AbstractNamedDimsArray, dimnames_codomain; kwargs...)
-  codomain = to_nameddimsindices(a, dimnames_codomain)
-  domain = setdiff(nameddimsindices(a), codomain)
-  return lq(a, codomain, domain; kwargs...)
+
+# Overload LinearAlgebra functions where relevant.
+function LinearAlgebra.qr(a::AbstractNamedDimsArray, args...; kwargs...)
+  return TensorAlgebra.qr(a, args...; kwargs...)
 end
 function LinearAlgebra.lq(a::AbstractNamedDimsArray, args...; kwargs...)
   return TensorAlgebra.lq(a, args...; kwargs...)
 end
+function LinearAlgebra.factorize(a::AbstractNamedDimsArray, args...; kwargs...)
+  return TensorAlgebra.factorize(a, args...; kwargs...)
+end
+
+#
+# Non-binary factorizations.
+#
 
 function TensorAlgebra.svd(
   a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs...

test/test_tensoralgebra.jl

@@ -1,5 +1,14 @@
-using LinearAlgebra: lq, norm, qr, svd
-using NamedDimsArrays: dename, left_null, nameddimsindices, namedoneto, right_null
+using LinearAlgebra: factorize, lq, norm, qr, svd
+using NamedDimsArrays:
+  dename,
+  left_null,
+  left_orth,
+  left_polar,
+  nameddimsindices,
+  namedoneto,
+  right_null,
+  right_orth,
+  right_polar
 using TensorAlgebra: TensorAlgebra, contract, fusedims, splitdims
 using Test: @test, @testset, @test_broken
 elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
@@ -50,14 +59,14 @@ elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     a = randn(elt, i, j)
     # TODO: Should this be allowed?
     # TODO: Add support for specifying new name.
-    for f in (qr, lq)
+    for f in (qr, lq, left_polar, right_polar, left_orth, right_orth, factorize)
       x, y = f(a, (i,))
       @test x * y ≈ a
     end
 
     a = randn(elt, i, j, k, l)
     # TODO: Add support for specifying new name.
-    for f in (qr, lq)
+    for f in (qr, lq, left_polar, right_polar, left_orth, right_orth, factorize)
       x, y = f(a, (i, k), (j, l))
       @test x * y ≈ a
     end