Allow two int-tuple inputs to represent biperm in factorizations and matrix functions (#97)

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.0"
+version = "0.5.1"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/factorizations.jl

@@ -5,33 +5,42 @@ for f in (
         :qr, :lq, :left_polar, :right_polar, :polar, :left_orth, :right_orth, :orth, :factorize,
     )
     @eval begin
-        function $f(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...)
-            biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...)
-            return $f(A, biperm; kwargs...)
-        end
-        function $f(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...)
+        function $f(
+                A::AbstractArray,
+                codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}};
+                kwargs...,
+            )
             # tensor to matrix
-            A_mat = matricize(A, biperm)
+            A_mat = matricize(A, codomain_perm, domain_perm)
 
             # factorization
             X, Y = MatrixAlgebra.$f(A_mat; kwargs...)
 
             # matrix to tensor
+            biperm = permmortar((codomain_perm, domain_perm))
             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
+        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...)
+        end
+        function $f(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...)
+            return $f(A, blocks(biperm)...; kwargs...)
+        end
     end
 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, biperm::AbstractBlockPermutation{2}; kwargs...) -> Q, R
 
 Compute the QR decomposition of a generic N-dimensional array, by interpreting it as
 a linear map from the domain to the codomain indices. These can be specified either via
-their labels, or directly through a `biperm`.
+their labels or directly through a bi-permutation.
 
 ## Keyword arguments
 
@@ -45,11 +54,12 @@ 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, biperm::AbstractBlockPermutation{2}; kwargs...) -> L, Q
 
 Compute the LQ decomposition of a generic N-dimensional array, by interpreting it as
 a linear map from the domain to the codomain indices. These can be specified either via
-their labels, or directly through a `biperm`.
+their labels or directly through a bi-permutation.
 
 ## Keyword arguments
 
@@ -63,11 +73,12 @@ 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, biperm::AbstractBlockPermutation{2}; kwargs...) -> W, P
 
 Compute the left polar decomposition of a generic N-dimensional array, by interpreting it as
 a linear map from the domain to the codomain indices. These can be specified either via
-their labels, or directly through a `biperm`.
+their labels or directly through a bi-permutation.
 
 ## Keyword arguments
 
@@ -79,11 +90,12 @@ 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, biperm::AbstractBlockPermutation{2}; kwargs...) -> P, W
 
 Compute the right polar decomposition of a generic N-dimensional array, by interpreting it as
 a linear map from the domain to the codomain indices. These can be specified either via
-their labels, or directly through a `biperm`.
+their labels or directly through a bi-permutation.
 
 ## Keyword arguments
 
@@ -95,11 +107,12 @@ 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, biperm::AbstractBlockPermutation{2}; kwargs...) -> V, C
 
 Compute the left orthogonal decomposition of a generic N-dimensional array, by interpreting it as
 a linear map from the domain to the codomain indices. These can be specified either via
-their labels, or directly through a `biperm`.
+their labels or directly through a bi-permutation.
 
 ## Keyword arguments
 
@@ -111,11 +124,12 @@ 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, biperm::AbstractBlockPermutation{2}; kwargs...) -> C, V
 
 Compute the right orthogonal decomposition of a generic N-dimensional array, by interpreting it as
 a linear map from the domain to the codomain indices. These can be specified either via
-their labels, or directly through a `biperm`.
+their labels or directly through a bi-permutation.
 
 ## Keyword arguments
 
@@ -127,11 +141,12 @@ 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, biperm::AbstractBlockPermutation{2}; kwargs...) -> X, Y
 
 Compute the decomposition of a generic N-dimensional array, by interpreting it as
 a linear map from the domain to the codomain indices. These can be specified either via
-their labels, or directly through a `biperm`.
+their labels or directly through a bi-permutation.
 
 ## Keyword arguments
 
@@ -143,11 +158,12 @@ 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, biperm::AbstractBlockPermutation{2}; kwargs...) -> D, V
 
 Compute the eigenvalue decomposition of a generic N-dimensional array, by interpreting it as
 a linear map from the domain to the codomain indices. These can be specified either via
-their labels, or directly through a `biperm`.
+their labels or directly through a bi-permutation.
 
 ## Keyword arguments
 
@@ -161,28 +177,37 @@ See also `MatrixAlgebraKit.eig_full!`, `MatrixAlgebraKit.eig_trunc!`, `MatrixAlg
 """
 function eigen(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...)
     biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...)
-    return eigen(A, biperm; kwargs...)
+    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}};
+        kwargs...,
+    )
     # tensor to matrix
-    A_mat = matricize(A, biperm)
+    A_mat = matricize(A, codomain_perm, domain_perm)
 
     # factorization
     D, V = MatrixAlgebra.eigen!(A_mat; kwargs...)
 
     # matrix to tensor
+    biperm = permmortar((codomain_perm, domain_perm))
     axes_codomain, = blocks(axes(A)[biperm])
     axes_V = tuplemortar((axes_codomain, (axes(V, ndims(V)),)))
     return D, unmatricize(V, axes_V)
 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, biperm::AbstractBlockPermutation{2}; kwargs...) -> D
 
 Compute the eigenvalues of a generic N-dimensional array, by interpreting it as
 a linear map from the domain to the codomain indices. These can be specified either via
-their labels, or directly through a `biperm`. The output is a vector of eigenvalues.
+their labels or directly through a bi-permutation. The output is a vector of eigenvalues.
 
 ## Keyword arguments
 
@@ -194,20 +219,28 @@ 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, biperm; kwargs...)
+    return eigvals(A, blocks(biperm)...; kwargs...)
 end
 function eigvals(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...)
-    A_mat = matricize(A, biperm)
+    return eigvals(A, blocks(biperm)...; kwargs...)
+end
+function eigvals(
+        A::AbstractArray,
+        codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}};
+        kwargs...,
+    )
+    A_mat = matricize(A, codomain_perm, domain_perm)
     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, biperm::AbstractBlockPermutation{2}; kwargs...) -> U, S, Vᴴ
 
 Compute the SVD decomposition of a generic N-dimensional array, by interpreting it as
 a linear map from the domain to the codomain indices. These can be specified either via
-their labels, or directly through a `biperm`.
+their labels or directly through a bi-permutation.
 
 ## Keyword arguments
 
@@ -220,16 +253,24 @@ See also `MatrixAlgebraKit.svd_full!`, `MatrixAlgebraKit.svd_compact!`, and `Mat
 """
 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...)
+    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}};
+        kwargs...,
+    )
     # tensor to matrix
-    A_mat = matricize(A, biperm)
+    A_mat = matricize(A, codomain_perm, domain_perm)
 
     # factorization
     U, S, Vᴴ = MatrixAlgebra.svd!(A_mat; kwargs...)
 
     # matrix to tensor
+    biperm = permmortar((codomain_perm, domain_perm))
     axes_codomain, axes_domain = blocks(axes(A)[biperm])
     axes_U = tuplemortar((axes_codomain, (axes(U, 2),)))
     axes_Vᴴ = tuplemortar(((axes(Vᴴ, 1),), axes_domain))
@@ -238,30 +279,38 @@ 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, biperm::AbstractBlockPermutation{2}) -> S
 
 Compute the singular values of a generic N-dimensional array, by interpreting it as
 a linear map from the domain to the codomain indices. These can be specified either via
-their labels, or directly through a `biperm`. The output is a vector of singular values.
+their labels or directly through a bi-permutation. The output is a vector of singular values.
 
 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, biperm)
+    return svdvals(A, blocks(biperm)...)
 end
 function svdvals(A::AbstractArray, biperm::AbstractBlockPermutation{2})
-    A_mat = matricize(A, biperm)
+    return svdvals(A, blocks(biperm)...)
+end
+function svdvals(
+        A::AbstractArray,
+        codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}}
+    )
+    A_mat = matricize(A, codomain_perm, domain_perm)
     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, biperm::AbstractBlockPermutation{2}; kwargs...) -> N
 
 Compute the left nullspace of a generic N-dimensional array, by interpreting it as
 a linear map from the domain to the codomain indices. These can be specified either via
-their labels, or directly through a `biperm`.
+their labels or directly through a bi-permutation.
 The output satisfies `N' * A ≈ 0` and `N' * N ≈ I`.
 
 ## Keyword arguments
@@ -274,23 +323,32 @@ The output satisfies `N' * A ≈ 0` and `N' * N ≈ I`.
 """
 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, biperm; kwargs...)
+    return left_null(A, blocks(biperm)...; kwargs...)
 end
 function left_null(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...)
-    A_mat = matricize(A, biperm)
+    return left_null(A, blocks(biperm)...; kwargs...)
+end
+function left_null(
+        A::AbstractArray,
+        codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}};
+        kwargs...,
+    )
+    A_mat = matricize(A, codomain_perm, domain_perm)
     N = MatrixAlgebraKit.left_null!(A_mat; kwargs...)
+    biperm = permmortar((codomain_perm, domain_perm))
     axes_codomain = first(blocks(axes(A)[biperm]))
     axes_N = tuplemortar((axes_codomain, (axes(N, 2),)))
     return unmatricize(N, axes_N)
 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, biperm::AbstractBlockPermutation{2}; kwargs...) -> Nᴴ
 
 Compute the right nullspace of a generic N-dimensional array, by interpreting it as
 a linear map from the domain to the codomain indices. These can be specified either via
-their labels, or directly through a `biperm`.
+their labels or directly through a bi-permutation.
 The output satisfies `A * Nᴴ' ≈ 0` and `Nᴴ * Nᴴ' ≈ I`.
 
 ## Keyword arguments
@@ -303,11 +361,19 @@ The output satisfies `A * Nᴴ' ≈ 0` and `Nᴴ * Nᴴ' ≈ I`.
 """
 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, biperm; kwargs...)
+    return right_null(A, blocks(biperm)...; kwargs...)
 end
 function right_null(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...)
-    A_mat = matricize(A, biperm)
+    return right_null(A, blocks(biperm)...; kwargs...)
+end
+function right_null(
+        A::AbstractArray,
+        codomain_perm::Tuple{Vararg{Int}}, domain_perm::Tuple{Vararg{Int}};
+        kwargs...,
+    )
+    A_mat = matricize(A, codomain_perm, domain_perm)
     Nᴴ = MatrixAlgebraKit.right_null!(A_mat; kwargs...)
+    biperm = permmortar((codomain_perm, domain_perm))
     axes_domain = last(blocks((axes(A)[biperm])))
     axes_Nᴴ = tuplemortar(((axes(Nᴴ, 1),), axes_domain))
     return unmatricize(Nᴴ, axes_Nᴴ)

src/matricize.jl

@@ -154,16 +154,11 @@ function unmatricize(
     )
 end
 
-function unmatricize(
-        m::AbstractMatrix,
-        blocked_axes::BlockedTuple{2, <:Any, <:Tuple{Vararg{AbstractUnitRange}}},
-    )
+function unmatricize(m::AbstractMatrix, blocked_axes::AbstractBlockTuple{2})
     return unmatricize(FusionStyle(m), m, blocked_axes)
 end
 function unmatricize(
-        style::FusionStyle,
-        m::AbstractMatrix,
-        blocked_axes::BlockedTuple{2, <:Any, <:Tuple{Vararg{AbstractUnitRange}}},
+        style::FusionStyle, m::AbstractMatrix, blocked_axes::AbstractBlockTuple{2}
     )
     return unmatricize(style, m, blocks(blocked_axes)...)
 end