More matrix factorizations

Author: mtfishman

src/MatrixAlgebra.jl

@@ -0,0 +1,110 @@
+module MatrixAlgebra
+
+using LinearAlgebra: LinearAlgebra
+using MatrixAlgebraKit:
+  eig_full,
+  eig_full!,
+  eig_trunc,
+  eig_trunc!,
+  eig_vals,
+  eig_vals!,
+  eigh_full,
+  eigh_full!,
+  eigh_trunc,
+  eigh_trunc!,
+  eigh_vals,
+  eigh_vals!,
+  left_orth,
+  left_orth!,
+  lq_full,
+  lq_full!,
+  lq_compact,
+  lq_compact!,
+  qr_full,
+  qr_full!,
+  qr_compact,
+  qr_compact!,
+  right_orth,
+  right_orth!,
+  svd_full,
+  svd_full!,
+  svd_compact,
+  svd_compact!,
+  svd_trunc,
+  svd_trunc!
+
+for (f, f_full, f_compact) in (
+  (:qr, :qr_full, :qr_compact),
+  (:qr!, :qr_full!, :qr_compact!),
+  (:lq, :lq_full, :lq_compact),
+  (:lq!, :lq_full!, :lq_compact!),
+)
+  @eval begin
+    function $f(A::AbstractMatrix; full::Bool=false, kwargs...)
+      f = full ? $f_full : $f_compact
+      return f(A; kwargs...)
+    end
+  end
+end
+
+for (eigen, eigh_full, eig_full, eigh_trunc, eig_trunc) in (
+  (:eigen, :eigh_full, :eig_full, :eigh_trunc, :eig_trunc),
+  (:eigen!, :eigh_full!, :eig_full!, :eigh_trunc!, :eig_trunc!),
+)
+  @eval begin
+    function $eigen(A::AbstractMatrix; trunc=nothing, ishermitian=nothing, kwargs...)
+      ishermitian = @something ishermitian LinearAlgebra.ishermitian(A)
+      f = if !isnothing(trunc)
+        ishermitian ? $eigh_trunc : $eig_trunc
+      else
+        ishermitian ? $eigh_full : $eig_full
+      end
+      return f(A; kwargs...)
+    end
+  end
+end
+
+for (eigvals, eigh_vals, eig_vals) in
+    ((:eigvals, :eigh_vals, :eig_vals), (:eigvals!, :eigh_vals!, :eig_vals!))
+  @eval begin
+    function $eigvals(A::AbstractMatrix; ishermitian=nothing, kwargs...)
+      ishermitian = @something ishermitian LinearAlgebra.ishermitian(A)
+      f = (ishermitian ? $eigh_vals : $eig_vals)
+      return f(A; kwargs...)
+    end
+  end
+end
+
+for (svd, svd_trunc, svd_full, svd_compact) in (
+  (:svd, :svd_trunc, :svd_full, :svd_compact),
+  (:svd!, :svd_trunc!, :svd_full!, :svd_compact!),
+)
+  @eval begin
+    function $svd(A::AbstractMatrix; full::Bool=false, trunc=nothing, kwargs...)
+      return if !isnothing(trunc)
+        @assert !full "Specified both full and truncation, currently not supported"
+        $svd_trunc(A; trunc, kwargs...)
+      else
+        (full ? $svd_full : $svd_compact)(A; kwargs...)
+      end
+    end
+  end
+end
+
+for (factorize, left_orth, right_orth) in
+    ((:factorize, :left_orth, :right_orth), (:factorize!, :left_orth!, :right_orth!))
+  @eval begin
+    function $factorize(A::AbstractMatrix; orth=:left, kwargs...)
+      f = if orth == :left
+        $left_orth
+      elseif orth == :right
+        $right_orth
+      else
+        throw(ArgumentError("`orth=$orth` not supported."))
+      end
+      return f(A; kwargs...)
+    end
+  end
+end
+
+end

src/TensorAlgebra.jl

@@ -2,6 +2,8 @@ module TensorAlgebra
 
 export contract, contract!, eigen, eigvals, lq, left_null, qr, right_null, svd, svdvals
 
+include("MatrixAlgebra.jl")
+using .MatrixAlgebra: MatrixAlgebra
 include("blockedtuple.jl")
 include("blockedpermutation.jl")
 include("BaseExtensions/BaseExtensions.jl")

src/factorizations.jl

@@ -1,25 +1,6 @@
 using LinearAlgebra: LinearAlgebra
-using MatrixAlgebraKit:
-  eig_full!,
-  eig_trunc!,
-  eig_vals!,
-  eigh_full!,
-  eigh_trunc!,
-  eigh_vals!,
-  left_null!,
-  left_orth!,
-  left_polar!,
-  lq_full!,
-  lq_compact!,
-  qr_full!,
-  qr_compact!,
-  right_null!,
-  right_orth!,
-  right_polar!,
-  svd_full!,
-  svd_compact!,
-  svd_trunc!,
-  svd_vals!
+using .MatrixAlgebra: MatrixAlgebra
+using MatrixAlgebraKit: MatrixAlgebraKit
 
 """
     qr(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> Q, R
@@ -41,12 +22,12 @@ function qr(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs..
   biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...)
   return qr(A, biperm; kwargs...)
 end
-function qr(A::AbstractArray, biperm::BlockedPermutation{2}; full::Bool=false, kwargs...)
+function qr(A::AbstractArray, biperm::BlockedPermutation{2}; kwargs...)
   # tensor to matrix
   A_mat = fusedims(A, biperm)
 
   # factorization
-  Q, R = full ? qr_full!(A_mat; kwargs...) : qr_compact!(A_mat; kwargs...)
+  Q, R = MatrixAlgebra.qr(A_mat; kwargs...)
 
   # matrix to tensor
   axes_codomain, axes_domain = blockpermute(axes(A), biperm)
@@ -75,12 +56,12 @@ function lq(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs..
   biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...)
   return lq(A, biperm; kwargs...)
 end
-function lq(A::AbstractArray, biperm::BlockedPermutation{2}; full::Bool=false, kwargs...)
+function lq(A::AbstractArray, biperm::BlockedPermutation{2}; kwargs...)
   # tensor to matrix
   A_mat = fusedims(A, biperm)
 
   # factorization
-  L, Q = (full ? lq_full! : lq_compact!)(A_mat; kwargs...)
+  L, Q = MatrixAlgebra.lq(A_mat; kwargs...)
 
   # matrix to tensor
   axes_codomain, axes_domain = blockpermute(axes(A), biperm)
@@ -111,25 +92,12 @@ function eigen(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwarg
   biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...)
   return eigen(A, biperm; kwargs...)
 end
-function eigen(
-  A::AbstractArray,
-  biperm::BlockedPermutation{2};
-  trunc=nothing,
-  ishermitian=nothing,
-  kwargs...,
-)
+function eigen(A::AbstractArray, biperm::BlockedPermutation{2}; kwargs...)
   # tensor to matrix
   A_mat = fusedims(A, biperm)
 
-  ishermitian = @something ishermitian LinearAlgebra.ishermitian(A_mat)
-
   # factorization
-  f! = if !isnothing(trunc)
-    ishermitian ? eigh_trunc! : eig_trunc!
-  else
-    ishermitian ? eigh_full! : eig_full!
-  end
-  D, V = f!(A_mat; kwargs...)
+  D, V = MatrixAlgebra.eigen!(A_mat; kwargs...)
 
   # matrix to tensor
   axes_codomain, = blockpermute(axes(A), biperm)
@@ -161,11 +129,9 @@ function eigvals(
   A::AbstractArray, biperm::BlockedPermutation{2}; ishermitian=nothing, kwargs...
 )
   A_mat = fusedims(A, biperm)
-  ishermitian = @something ishermitian LinearAlgebra.ishermitian(A_mat)
-  return (ishermitian ? eigh_vals! : eig_vals!)(A_mat; kwargs...)
+  return MatrixAlgebra.eigvals!(A_mat; kwargs...)
 end
 
-# TODO: separate out the algorithm selection step from the implementation
 """
     svd(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> U, S, Vᴴ
     svd(A::AbstractArray, biperm::BlockedPermutation{2}; kwargs...) -> U, S, Vᴴ
@@ -187,23 +153,12 @@ function svd(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs.
   biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...)
   return svd(A, biperm; kwargs...)
 end
-function svd(
-  A::AbstractArray,
-  biperm::BlockedPermutation{2};
-  full::Bool=false,
-  trunc=nothing,
-  kwargs...,
-)
+function svd(A::AbstractArray, biperm::BlockedPermutation{2}; kwargs...)
   # tensor to matrix
   A_mat = fusedims(A, biperm)
 
   # factorization
-  if !isnothing(trunc)
-    @assert !full "Specified both full and truncation, currently not supported"
-    U, S, Vᴴ = svd_trunc!(A_mat; trunc, kwargs...)
-  else
-    U, S, Vᴴ = full ? svd_full!(A_mat; kwargs...) : svd_compact!(A_mat; kwargs...)
-  end
+  U, S, Vᴴ = MatrixAlgebra.svd!(A_mat; kwargs...)
 
   # matrix to tensor
   axes_codomain, axes_domain = blockpermute(axes(A), biperm)
@@ -228,7 +183,7 @@ function svdvals(A::AbstractArray, labels_A, labels_codomain, labels_domain)
 end
 function svdvals(A::AbstractArray, biperm::BlockedPermutation{2})
   A_mat = fusedims(A, biperm)
-  return svd_vals!(A_mat)
+  return MatrixAlgebraKit.svd_vals!(A_mat)
 end
 
 """
@@ -254,7 +209,7 @@ function left_null(A::AbstractArray, labels_A, labels_codomain, labels_domain; k
 end
 function left_null(A::AbstractArray, biperm::BlockedPermutation{2}; kwargs...)
   A_mat = fusedims(A, biperm)
-  N = left_null!(A_mat; kwargs...)
+  N = MatrixAlgebraKit.left_null!(A_mat; kwargs...)
   axes_codomain, _ = blockpermute(axes(A), biperm)
   axes_N = (axes_codomain..., axes(N, 2))
   N_tensor = splitdims(N, axes_N)
@@ -284,7 +239,7 @@ function right_null(A::AbstractArray, labels_A, labels_codomain, labels_domain;
 end
 function right_null(A::AbstractArray, biperm::BlockedPermutation{2}; kwargs...)
   A_mat = fusedims(A, biperm)
-  Nᴴ = right_null!(A_mat; kwargs...)
+  Nᴴ = MatrixAlgebraKit.right_null!(A_mat; kwargs...)
   _, axes_domain = blockpermute(axes(A), biperm)
   axes_Nᴴ = (axes(Nᴴ, 1), axes_domain...)
   return splitdims(Nᴴ, axes_Nᴴ)
@@ -313,7 +268,7 @@ function left_polar(A::AbstractArray, biperm::BlockedPermutation{2}; kwargs...)
   A_mat = fusedims(A, biperm)
 
   # factorization
-  W, P = left_polar!(A_mat; kwargs...)
+  W, P = MatrixAlgebraKit.left_polar!(A_mat; kwargs...)
 
   # matrix to tensor
   axes_codomain, axes_domain = blockpermute(axes(A), biperm)
@@ -345,7 +300,7 @@ function right_polar(A::AbstractArray, biperm::BlockedPermutation{2}; kwargs...)
   A_mat = fusedims(A, biperm)
 
   # factorization
-  P, W = right_polar!(A_mat; kwargs...)
+  P, W = MatrixAlgebraKit.right_polar!(A_mat; kwargs...)
 
   # matrix to tensor
   axes_codomain, axes_domain = blockpermute(axes(A), biperm)
@@ -377,7 +332,7 @@ function left_orth(A::AbstractArray, biperm::BlockedPermutation{2}; kwargs...)
   A_mat = fusedims(A, biperm)
 
   # factorization
-  V, C = left_orth!(A_mat; kwargs...)
+  V, C = MatrixAlgebraKit.left_orth!(A_mat; kwargs...)
 
   # matrix to tensor
   axes_codomain, axes_domain = blockpermute(axes(A), biperm)
@@ -409,7 +364,7 @@ function right_orth(A::AbstractArray, biperm::BlockedPermutation{2}; kwargs...)
   A_mat = fusedims(A, biperm)
 
   # factorization
-  P, W = right_orth!(A_mat; kwargs...)
+  P, W = MatrixAlgebraKit.right_orth!(A_mat; kwargs...)
 
   # matrix to tensor
   axes_codomain, axes_domain = blockpermute(axes(A), biperm)
@@ -441,7 +396,7 @@ function factorize(A::AbstractArray, biperm::BlockedPermutation{2}; orth=:left,
   A_mat = fusedims(A, biperm)
 
   # factorization
-  X, Y = (orth == :left ? left_orth! : right_orth!)(A_mat; kwargs...)
+  X, Y = MatrixAlgebra.factorize!(A_mat; kwargs...)
 
   # matrix to tensor
   axes_codomain, axes_domain = blockpermute(axes(A), biperm)