Add trivial bipermutation interface to matricized factorizations and functions (#98)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.5.1"
+version = "0.5.2"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/MatrixAlgebra.jl

@@ -75,14 +75,34 @@ for (svd, svd_trunc, svd_full, svd_compact) in (
         (:svd, :svd_trunc, :svd_full, :svd_compact),
         (:svd!, :svd_trunc!, :svd_full!, :svd_compact!),
     )
+    _svd = Symbol(:_, svd)
     @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
+        function $svd(
+                A::AbstractMatrix;
+                full::Union{Bool, Val} = Val(false),
+                trunc = nothing,
+                kwargs...,
+            )
+            return $_svd(full, trunc, A; kwargs...)
+        end
+        function $_svd(full::Bool, trunc, A::AbstractMatrix; kwargs...)
+            return $_svd(Val(full), trunc, A; kwargs...)
+        end
+        function $_svd(full::Val{false}, trunc::Nothing, A::AbstractMatrix; kwargs...)
+            return $svd_compact(A; kwargs...)
+        end
+        function $_svd(full::Val{false}, trunc, A::AbstractMatrix; kwargs...)
+            return $svd_trunc(A; trunc, kwargs...)
+        end
+        function $_svd(full::Val{true}, trunc::Nothing, A::AbstractMatrix; kwargs...)
+            return $svd_full(A; kwargs...)
+        end
+        function $_svd(full::Val{true}, trunc, A::AbstractMatrix; kwargs...)
+            return throw(
+                ArgumentError(
+                    "Specified both full and truncation, currently not supported"
+                )
+            )
         end
     end
 end

src/factorizations.jl

@@ -7,22 +7,38 @@ for f in (
     @eval begin
         function $f(
                 A::AbstractArray,
-                codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}};
+                codomain_length::Val, domain_length::Val;
                 kwargs...,
             )
             # tensor to matrix
-            A_mat = matricize(A, codomain_perm, domain_perm)
+            A_mat = matricize(A, codomain_length, domain_length)
 
             # factorization
             X, Y = MatrixAlgebra.$f(A_mat; kwargs...)
 
             # matrix to tensor
-            biperm = permmortar((codomain_perm, domain_perm))
+            biperm = blockedtrivialperm((codomain_length, domain_length))
             axes_codomain, axes_domain = blocks(axes(A)[biperm])
             axes_X = tuplemortar((axes_codomain, (axes(X, 2),)))
             axes_Y = tuplemortar(((axes(Y, 1),), axes_domain))
             return unmatricize(X, axes_X), unmatricize(Y, axes_Y)
         end
+    end
+end
+
+for f in (
+        :qr, :lq, :left_polar, :right_polar, :polar, :left_orth, :right_orth, :orth, :factorize,
+        :eigen, :eigvals, :svd, :svdvals, :left_null, :right_null,
+    )
+    @eval begin
+        function $f(
+                A::AbstractArray,
+                codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}};
+                kwargs...,
+            )
+            A_perm = bipermutedims(A, codomain_perm, domain_perm)
+            return $f(A_perm, Val(length(codomain_perm)), Val(length(domain_perm)); kwargs...)
+        end
         function $f(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...)
             biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...)
             return $f(A, blocks(biperm)...; kwargs...)
@@ -36,6 +52,7 @@ end
 """
     qr(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> Q, R
     qr(A::AbstractArray, codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}}; kwargs...) -> Q, R
+    qr(A::AbstractArray, codomain_length::Val, domain_length::Val; kwargs...) -> Q, R
     qr(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> Q, R
 
 Compute the QR decomposition of a generic N-dimensional array, by interpreting it as
@@ -55,6 +72,7 @@ qr
 """
     lq(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> L, Q
     lq(A::AbstractArray, codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}}; kwargs...) -> L, Q
+    lq(A::AbstractArray, codomain_length::Val, domain_length::Val; kwargs...) -> L, Q
     lq(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> L, Q
 
 Compute the LQ decomposition of a generic N-dimensional array, by interpreting it as
@@ -74,6 +92,7 @@ lq
 """
     left_polar(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> W, P
     left_polar(A::AbstractArray, codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}}; kwargs...) -> W, P
+    left_polar(A::AbstractArray, codomain_length::Val, domain_length::Val; kwargs...) -> W, P
     left_polar(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> W, P
 
 Compute the left polar decomposition of a generic N-dimensional array, by interpreting it as
@@ -91,6 +110,7 @@ left_polar
 """
     right_polar(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> P, W
     right_polar(A::AbstractArray, codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}}; kwargs...) -> P, W
+    right_polar(A::AbstractArray, codomain_length::Val, domain_length::Val; kwargs...) -> P, W
     right_polar(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> P, W
 
 Compute the right polar decomposition of a generic N-dimensional array, by interpreting it as
@@ -108,6 +128,7 @@ right_polar
 """
     left_orth(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> V, C
     left_orth(A::AbstractArray, codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}}; kwargs...) -> V, C
+    left_orth(A::AbstractArray, codomain_length::Val, domain_length::Val; kwargs...) -> V, C
     left_orth(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> V, C
 
 Compute the left orthogonal decomposition of a generic N-dimensional array, by interpreting it as
@@ -125,6 +146,7 @@ left_orth
 """
     right_orth(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> C, V
     right_orth(A::AbstractArray, codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}}; kwargs...) -> C, V
+    right_orth(A::AbstractArray, codomain_length::Val, domain_length::Val; kwargs...) -> C, V
     right_orth(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> C, V
 
 Compute the right orthogonal decomposition of a generic N-dimensional array, by interpreting it as
@@ -142,6 +164,7 @@ right_orth
 """
     factorize(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> X, Y
     factorize(A::AbstractArray, codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}}; kwargs...) -> X, Y
+    factorize(A::AbstractArray, codomain_length::Val, domain_length::Val; kwargs...) -> X, Y
     factorize(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> X, Y
 
 Compute the decomposition of a generic N-dimensional array, by interpreting it as
@@ -159,6 +182,7 @@ factorize
 """
     eigen(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> D, V
     eigen(A::AbstractArray, codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}}; kwargs...) -> D, V
+    eigen(A::AbstractArray, codomain_length::Val, domain_length::Val; kwargs...) -> D, V
     eigen(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> D, V
 
 Compute the eigenvalue decomposition of a generic N-dimensional array, by interpreting it as
@@ -175,26 +199,18 @@ their labels or directly through a bi-permutation.
 See also `MatrixAlgebraKit.eig_full!`, `MatrixAlgebraKit.eig_trunc!`, `MatrixAlgebraKit.eig_vals!`,
 `MatrixAlgebraKit.eigh_full!`, `MatrixAlgebraKit.eigh_trunc!`, and `MatrixAlgebraKit.eigh_vals!`.
 """
-function eigen(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...)
-    biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...)
-    return eigen(A, blocks(biperm)...; kwargs...)
-end
-function eigen(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...)
-    return eigen(A, blocks(biperm)...; kwargs...)
-end
 function eigen(
         A::AbstractArray,
-        codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}};
+        codomain_length::Val, domain_length::Val;
         kwargs...,
     )
     # tensor to matrix
-    A_mat = matricize(A, codomain_perm, domain_perm)
-
+    A_mat = matricize(A, codomain_length, domain_length)
     # factorization
     D, V = MatrixAlgebra.eigen!(A_mat; kwargs...)
 
     # matrix to tensor
-    biperm = permmortar((codomain_perm, domain_perm))
+    biperm = blockedtrivialperm((codomain_length, domain_length))
     axes_codomain, = blocks(axes(A)[biperm])
     axes_V = tuplemortar((axes_codomain, (axes(V, ndims(V)),)))
     return D, unmatricize(V, axes_V)
@@ -203,6 +219,7 @@ end
 """
     eigvals(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> D
     eigvals(A::AbstractArray, codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}}; kwargs...) -> D
+    eigvals(A::AbstractArray, codomain_length::Val, domain_length::Val; kwargs...) -> D
     eigvals(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> D
 
 Compute the eigenvalues of a generic N-dimensional array, by interpreting it as
@@ -217,25 +234,19 @@ their labels or directly through a bi-permutation. The output is a vector of eig
 
 See also `MatrixAlgebraKit.eig_vals!` and `MatrixAlgebraKit.eigh_vals!`.
 """
-function eigvals(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...)
-    biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...)
-    return eigvals(A, blocks(biperm)...; kwargs...)
-end
-function eigvals(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...)
-    return eigvals(A, blocks(biperm)...; kwargs...)
-end
 function eigvals(
         A::AbstractArray,
-        codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}};
+        codomain_length::Val, domain_length::Val;
         kwargs...,
     )
-    A_mat = matricize(A, codomain_perm, domain_perm)
+    A_mat = matricize(A, codomain_length, domain_length)
     return MatrixAlgebra.eigvals!(A_mat; kwargs...)
 end
 
 """
     svd(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> U, S, Vᴴ
     svd(A::AbstractArray, codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}}; kwargs...) -> U, S, Vᴴ
+    svd(A::AbstractArray, codomain_length::Val, domain_length::Val; kwargs...) -> U, S, Vᴴ
     svd(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> U, S, Vᴴ
 
 Compute the SVD decomposition of a generic N-dimensional array, by interpreting it as
@@ -251,26 +262,18 @@ their labels or directly through a bi-permutation.
 
 See also `MatrixAlgebraKit.svd_full!`, `MatrixAlgebraKit.svd_compact!`, and `MatrixAlgebraKit.svd_trunc!`.
 """
-function svd(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...)
-    biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...)
-    return svd(A, blocks(biperm)...; kwargs...)
-end
-function svd(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...)
-    return svd(A, blocks(biperm)...; kwargs...)
-end
 function svd(
         A::AbstractArray,
-        codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}};
+        codomain_length::Val, domain_length::Val;
         kwargs...,
     )
     # tensor to matrix
-    A_mat = matricize(A, codomain_perm, domain_perm)
-
+    A_mat = matricize(A, codomain_length, domain_length)
     # factorization
     U, S, Vᴴ = MatrixAlgebra.svd!(A_mat; kwargs...)
 
     # matrix to tensor
-    biperm = permmortar((codomain_perm, domain_perm))
+    biperm = blockedtrivialperm((codomain_length, domain_length))
     axes_codomain, axes_domain = blocks(axes(A)[biperm])
     axes_U = tuplemortar((axes_codomain, (axes(U, 2),)))
     axes_Vᴴ = tuplemortar(((axes(Vᴴ, 1),), axes_domain))
@@ -280,6 +283,7 @@ end
 """
     svdvals(A::AbstractArray, labels_A, labels_codomain, labels_domain) -> S
     svdvals(A::AbstractArray, codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}}) -> S
+    svdvals(A::AbstractArray, codomain_length::Val, domain_length::Val) -> S
     svdvals(A::AbstractArray, biperm::AbstractBlockPermutation{2}) -> S
 
 Compute the singular values of a generic N-dimensional array, by interpreting it as
@@ -288,24 +292,18 @@ their labels or directly through a bi-permutation. The output is a vector of sin
 
 See also `MatrixAlgebraKit.svd_vals!`.
 """
-function svdvals(A::AbstractArray, labels_A, labels_codomain, labels_domain)
-    biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...)
-    return svdvals(A, blocks(biperm)...)
-end
-function svdvals(A::AbstractArray, biperm::AbstractBlockPermutation{2})
-    return svdvals(A, blocks(biperm)...)
-end
 function svdvals(
         A::AbstractArray,
-        codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}}
+        codomain_length::Val, domain_length::Val
     )
-    A_mat = matricize(A, codomain_perm, domain_perm)
+    A_mat = matricize(A, codomain_length, domain_length)
     return MatrixAlgebra.svdvals!(A_mat)
 end
 
 """
     left_null(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> N
     left_null(A::AbstractArray, codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}}; kwargs...) -> N
+    left_null(A::AbstractArray, codomain_length::Val, domain_length::Val; kwargs...) -> N
     left_null(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> N
 
 Compute the left nullspace of a generic N-dimensional array, by interpreting it as
@@ -321,21 +319,14 @@ The output satisfies `N' * A ≈ 0` and `N' * N ≈ I`.
   The options are `:qr`, `:qrpos` and `:svd`. The former two require `0 == atol == rtol`.
   The default is `:qrpos` if `atol == rtol == 0`, and `:svd` otherwise.
 """
-function left_null(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...)
-    biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...)
-    return left_null(A, blocks(biperm)...; kwargs...)
-end
-function left_null(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...)
-    return left_null(A, blocks(biperm)...; kwargs...)
-end
 function left_null(
         A::AbstractArray,
-        codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}};
+        codomain_length::Val, domain_length::Val;
         kwargs...,
     )
-    A_mat = matricize(A, codomain_perm, domain_perm)
+    A_mat = matricize(A, codomain_length, domain_length)
     N = MatrixAlgebraKit.left_null!(A_mat; kwargs...)
-    biperm = permmortar((codomain_perm, domain_perm))
+    biperm = blockedtrivialperm((codomain_length, domain_length))
     axes_codomain = first(blocks(axes(A)[biperm]))
     axes_N = tuplemortar((axes_codomain, (axes(N, 2),)))
     return unmatricize(N, axes_N)
@@ -344,6 +335,7 @@ end
 """
     right_null(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> Nᴴ
     right_null(A::AbstractArray, codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}}; kwargs...) -> Nᴴ
+    right_null(A::AbstractArray, codomain_length::Val, domain_length::Val; kwargs...) -> Nᴴ
     right_null(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> Nᴴ
 
 Compute the right nullspace of a generic N-dimensional array, by interpreting it as
@@ -359,21 +351,14 @@ The output satisfies `A * Nᴴ' ≈ 0` and `Nᴴ * Nᴴ' ≈ I`.
   The options are `:lq`, `:lqpos` and `:svd`. The former two require `0 == atol == rtol`.
   The default is `:lqpos` if `atol == rtol == 0`, and `:svd` otherwise.
 """
-function right_null(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...)
-    biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...)
-    return right_null(A, blocks(biperm)...; kwargs...)
-end
-function right_null(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...)
-    return right_null(A, blocks(biperm)...; kwargs...)
-end
 function right_null(
         A::AbstractArray,
-        codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}};
+        codomain_length::Val, domain_length::Val;
         kwargs...,
     )
-    A_mat = matricize(A, codomain_perm, domain_perm)
+    A_mat = matricize(A, codomain_length, domain_length)
     Nᴴ = MatrixAlgebraKit.right_null!(A_mat; kwargs...)
-    biperm = permmortar((codomain_perm, domain_perm))
+    biperm = blockedtrivialperm((codomain_length, domain_length))
     axes_domain = last(blocks((axes(A)[biperm])))
     axes_Nᴴ = tuplemortar(((axes(Nᴴ, 1),), axes_domain))
     return unmatricize(Nᴴ, axes_Nᴴ)