Balanced gram_eigh_full convention, add one and operator-construction primitives

Project.toml

@@ -1,6 +1,6 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.9.4"
+version = "0.9.5"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

src/MatrixAlgebra.jl

@@ -206,11 +206,11 @@ for (gram, gram_with_pinv, eigh_full) in (
     @eval begin
         function $gram(A::AbstractMatrix; alg = nothing, kwargs...)
             D, V = MAK.$eigh_full(A, MAK.select_algorithm(MAK.$eigh_full, A, alg))
-            return sqrth_safe(D; kwargs...) * V'
+            return V * sqrth_safe(D; kwargs...)
         end
         function $gram_with_pinv(A::AbstractMatrix; alg = nothing, kwargs...)
             D, V = MAK.$eigh_full(A, MAK.select_algorithm(MAK.$eigh_full, A, alg))
-            return sqrth_safe(D; kwargs...) * V', V * invsqrth_safe(D; kwargs...)
+            return V * sqrth_safe(D; kwargs...), invsqrth_safe(D; kwargs...) * V'
         end
     end
 end
@@ -220,10 +220,11 @@ end
     gram_eigh_full!!(A::AbstractMatrix; alg=nothing, atol=0, rtol=eps(real(eltype(A)))^(3//4)) -> X
 
 Gram factorization of a Hermitian positive semi-definite matrix via its
-eigendecomposition: returns `X = sqrth_safe(D; atol, rtol) * V'` such
-that `A ≈ X' * X`, where `A = V * D * V'`. Eigenvalues below `tol` (see
-[`pow_diag_safe`](@ref)) are clamped to zero. The `!!` variant may
-destroy `A`.
+eigendecomposition (balanced eigh): returns `X = V * sqrth_safe(D; atol, rtol)`
+such that `A ≈ X * X'`, where `A = V * D * V'`. The square-root of `D` is
+absorbed symmetrically into the two factors of the eigendecomposition.
+Eigenvalues below `tol` (see [`pow_diag_safe`](@ref)) are clamped to zero.
+The `!!` variant may destroy `A`.
 
 ## Keyword arguments
 
@@ -242,7 +243,7 @@ julia> A = B' * B;
 
 julia> X = gram_eigh_full(A);
 
-julia> X' * X ≈ A
+julia> X * X' ≈ A
 true
 ```
 
@@ -255,9 +256,9 @@ gram_eigh_full, gram_eigh_full!!
     gram_eigh_full_with_pinv!!(A::AbstractMatrix; alg=nothing, atol=0, rtol=eps(real(eltype(A)))^(3//4)) -> X, Y
 
 Like [`gram_eigh_full`](@ref), but additionally returns
-`Y = V * invsqrth_safe(D; atol, rtol) ≈ pinv(X)` so that `X * Y ≈ I` on
-the rank subspace. Eigenvalues below `tol` are clamped to zero in both
-factors. The `!!` variant may destroy `A`.
+`Y = invsqrth_safe(D; atol, rtol) * V' ≈ pinv(X)`, a left inverse of `X`
+on the rank subspace: `Y * X ≈ I`. Eigenvalues below `tol` are clamped to
+zero in both factors. The `!!` variant may destroy `A`.
 
 ## Keyword arguments
 
@@ -278,10 +279,10 @@ julia> A = B' * B;
 
 julia> X, Y = gram_eigh_full_with_pinv(A);
 
-julia> X' * X ≈ A
+julia> X * X' ≈ A
 true
 
-julia> X * Y ≈ I
+julia> Y * X ≈ I
 true
 ```
 """

src/TensorAlgebra.jl

@@ -22,6 +22,8 @@ include("contract/allocate_output.jl")
 include("contract/contract_matricize.jl")
 include("factorizations.jl")
 include("matrixfunctions.jl")
+include("similar_map.jl")
+include("projectto.jl")
 include("linearbroadcasted.jl")
 
 end

src/factorizations.jl

@@ -31,7 +31,7 @@ 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,
-        :gram_eigh_full, :gram_eigh_full_with_pinv,
+        :gram_eigh_full, :gram_eigh_full_with_pinv, :one,
     )
     @eval begin
         function $f(
@@ -312,8 +312,11 @@ function svd!!(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs
     biperm = trivialbiperm(ndims_codomain, Val(ndims(A)))
     axes_codomain, axes_domain = blocks(axes(A)[biperm])
     axes_U = tuplemortar((axes_codomain, (axes(U, 2),)))
+    axes_S = tuplemortar(((axes(U, 2),), (axes(Vᴴ, 1),)))
     axes_Vᴴ = tuplemortar(((axes(Vᴴ, 1),), axes_domain))
-    return unmatricize(style, U, axes_U), S, unmatricize(style, Vᴴ, axes_Vᴴ)
+    return unmatricize(style, U, axes_U),
+        unmatricize(style, S, axes_S),
+        unmatricize(style, Vᴴ, axes_Vᴴ)
 end
 function svd!!(A::AbstractArray, ndims_codomain::Val; kwargs...)
     return svd!!(FusionStyle(A), A, ndims_codomain; kwargs...)
@@ -443,7 +446,9 @@ end
 
 Gram factorization of a generic N-dimensional array, interpreting it as a
 Hermitian positive semi-definite linear map from the domain to the codomain
-dimensions. Returns `X` such that `A ≈ X' * X` (contracted on the rank leg).
+dimensions. Returns `X` such that `A ≈ X * X'` (contracted on the rank leg),
+i.e. the codomain axes of `X` match the codomain axes of `A` and `X` has a
+single trailing rank axis.
 
 ## Keyword arguments
 
@@ -462,7 +467,7 @@ julia> A = contract((:a, :b, :c, :d), conj(B), (:r, :a, :b), B, (:r, :c, :d));
 
 julia> X = gram_eigh_full(A, (:a, :b, :c, :d), (:a, :b), (:c, :d));
 
-julia> A ≈ contract((:a, :b, :c, :d), conj(X), (:r, :a, :b), X, (:r, :c, :d))
+julia> A ≈ contract((:a, :b, :c, :d), X, (:a, :b, :r), conj(X), (:c, :d, :r))
 true
 ```
 
@@ -478,7 +483,7 @@ function gram_eigh_full!!(
     X = MatrixAlgebra.gram_eigh_full!!(A_mat; kwargs...)
     biperm = trivialbiperm(ndims_codomain, Val(ndims(A)))
     axes_codomain = first(blocks(axes(A)[biperm]))
-    axes_X = tuplemortar(((axes(X, 1),), axes_codomain))
+    axes_X = tuplemortar((axes_codomain, (axes(X, 2),)))
     return unmatricize(style, X, axes_X)
 end
 function gram_eigh_full!!(A::AbstractArray, ndims_codomain::Val; kwargs...)
@@ -501,7 +506,9 @@ end
     gram_eigh_full_with_pinv(A::AbstractArray, biperm::AbstractBlockPermutation{2}; kwargs...) -> X, Y
 
 Like [`gram_eigh_full`](@ref), but additionally returns `Y ≈ pinv(X)` such
-that `X * Y ≈ I` on the rank subspace.
+that `Y * X ≈ I` on the rank subspace (a left inverse). The codomain axes
+of `X` match the codomain axes of `A`; `Y` has a leading rank axis followed
+by the codomain axes.
 
 ## Keyword arguments
 
@@ -522,10 +529,10 @@ julia> A = contract((:a, :b, :c, :d), conj(B), (:r, :a, :b), B, (:r, :c, :d));
 
 julia> X, Y = gram_eigh_full_with_pinv(A, (:a, :b, :c, :d), (:a, :b), (:c, :d));
 
-julia> A ≈ contract((:a, :b, :c, :d), conj(X), (:r, :a, :b), X, (:r, :c, :d))
+julia> A ≈ contract((:a, :b, :c, :d), X, (:a, :b, :r), conj(X), (:c, :d, :r))
 true
 
-julia> contract((:r, :s), X, (:r, :a, :b), Y, (:a, :b, :s)) ≈ I
+julia> contract((:r, :s), Y, (:r, :a, :b), X, (:a, :b, :s)) ≈ I
 true
 ```
 
@@ -540,8 +547,8 @@ function gram_eigh_full_with_pinv!!(
     X, Y = MatrixAlgebra.gram_eigh_full_with_pinv!!(A_mat; kwargs...)
     biperm = trivialbiperm(ndims_codomain, Val(ndims(A)))
     axes_codomain = first(blocks(axes(A)[biperm]))
-    axes_X = tuplemortar(((axes(X, 1),), axes_codomain))
-    axes_Y = tuplemortar((axes_codomain, (axes(Y, 2),)))
+    axes_X = tuplemortar((axes_codomain, (axes(X, 2),)))
+    axes_Y = tuplemortar(((axes(Y, 1),), axes_codomain))
     return unmatricize(style, X, axes_X), unmatricize(style, Y, axes_Y)
 end
 function gram_eigh_full_with_pinv!!(A::AbstractArray, ndims_codomain::Val; kwargs...)
@@ -556,3 +563,55 @@ end
 function gram_eigh_full_with_pinv(A::AbstractArray, ndims_codomain::Val; kwargs...)
     return gram_eigh_full_with_pinv!!(copy(A), ndims_codomain; kwargs...)
 end
+
+"""
+    TensorAlgebra.one(A::AbstractArray, labels_A, labels_codomain, labels_domain) -> Id
+    TensorAlgebra.one(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}) -> Id
+    TensorAlgebra.one(A::AbstractArray, ndims_codomain::Val) -> Id
+    TensorAlgebra.one(A::AbstractArray, biperm::AbstractBlockPermutation{2}) -> Id
+
+Construct the identity operator tensor whose shape mirrors `A`, interpreted as a
+linear map from the domain to the codomain dimensions. The codomain and domain
+partition is specified either via labels or directly through a bi-permutation;
+fused codomain and domain sizes must match. `A` is treated as a shape prototype
+and is not mutated.
+
+Not exported, since exporting would clash with the implicit `Base.one`. Qualify
+as `TensorAlgebra.one(A, ...)`.
+
+See also `MatrixAlgebraKit.one!`.
+
+# Examples
+
+```jldoctest
+julia> using LinearAlgebra: I
+
+julia> using TensorAlgebra: TensorAlgebra, matricize
+
+julia> A = randn(2, 3, 2, 3);
+
+julia> Id = TensorAlgebra.one(A, (:a, :b, :c, :d), (:a, :b), (:c, :d));
+
+julia> matricize(Id, Val(2)) ≈ I
+true
+```
+"""
+one
+
+function one!!(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...)
+    A_mat = matricize(style, A, ndims_codomain)
+    MatrixAlgebraKit.one!(A_mat)
+    biperm = trivialbiperm(ndims_codomain, Val(ndims(A)))
+    axes_codomain, axes_domain = blocks(axes(A)[biperm])
+    return unmatricize(style, A_mat, axes_codomain, axes_domain)
+end
+function one!!(A::AbstractArray, ndims_codomain::Val; kwargs...)
+    return one!!(FusionStyle(A), A, ndims_codomain; kwargs...)
+end
+
+function one(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...)
+    return one!!(style, copy(A), ndims_codomain; kwargs...)
+end
+function one(A::AbstractArray, ndims_codomain::Val; kwargs...)
+    return one!!(copy(A), ndims_codomain; kwargs...)
+end

src/projectto.jl

@@ -0,0 +1,59 @@
+"""
+    projectto!(dest, src) -> dest
+
+Orthogonally project `src` onto the representable subspace of `dest`, in place.
+Per-backend primitive. Tolerance-free and magnitude-blind: drops the
+non-representable component regardless of its size. Pair with
+[`checked_projectto!`](@ref) (or `isapprox`-after) when the discarded weight
+needs to be verified.
+
+The default falls through to `copyto!`, which is the right behavior for dense
+backends where the representable subspace is everything. Block-structured
+backends (e.g. `AbelianGradedArray`, `FusionTensor`) overload this to copy only
+the symmetry-allowed entries.
+"""
+projectto!(dest, src) = copyto!(dest, src)
+
+"""
+    checked_projectto!(dest, src; atol=0, rtol=…) -> dest
+
+Project `src` into `dest` via [`projectto!`](@ref), then verify that the
+discarded component is within the requested tolerance via `isapprox(src, dest; atol, rtol)`. Throws `InexactError` on failure. Backends may specialize this
+verb for a fused, cheaper check (e.g. a one-pass norm comparison).
+"""
+function checked_projectto!(
+        dest, src;
+        atol::Real = 0,
+        rtol::Real = Base.rtoldefault(real(eltype(src)))
+    )
+    projectto!(dest, src)
+    isapprox(src, dest; atol, rtol) ||
+        throw(InexactError(:checked_projectto!, typeof(dest), src))
+    return dest
+end
+
+"""
+    project_map(raw, codomain_axes, domain_axes) -> dest
+
+Allocate a map-shaped array via [`similar_map`](@ref) and project `raw` into it with
+[`projectto!`](@ref). Unchecked: any non-representable component of `raw` is dropped
+silently. The data-bearing member of the `_map` allocator family
+(`similar_map` / `zeros_map` / `project_map`); for the checked variant see
+[`checked_project_map`](@ref).
+"""
+function project_map(raw, codomain_axes, domain_axes)
+    return projectto!(similar_map(raw, eltype(raw), codomain_axes, domain_axes), raw)
+end
+
+"""
+    checked_project_map(raw, codomain_axes, domain_axes; atol=0, rtol=…) -> dest
+
+Allocate via [`similar_map`](@ref) and project `raw` into it with
+[`checked_projectto!`](@ref), throwing `InexactError` if the discarded component
+exceeds tolerance. Default kwargs match `checked_projectto!`.
+"""
+function checked_project_map(raw, codomain_axes, domain_axes; kwargs...)
+    return checked_projectto!(
+        similar_map(raw, eltype(raw), codomain_axes, domain_axes), raw; kwargs...
+    )
+end

src/similar_map.jl

@@ -0,0 +1,32 @@
+"""
+    similar_map(prototype, [T,] codomain_axes, domain_axes) -> O
+
+Allocate an array shaped as a linear map from `domain_axes` to
+`codomain_axes`, with axes `(codomain_axes..., conj.(domain_axes)...)` and
+element type `T` (defaulting to `eltype(prototype)`). The domain axes are
+conjugated so callers can pass both sides in the same (codomain) direction.
+For plain axes `conj` is a no-op. For graded axes it flips the sector arrows.
+`prototype` provides the array backend and is not mutated.
+
+# Examples
+
+```jldoctest
+julia> using TensorAlgebra: similar_map
+
+julia> O = similar_map(
+           randn(3),
+           Float32,
+           (Base.OneTo(2), Base.OneTo(3)),
+           (Base.OneTo(4), Base.OneTo(5))
+       );
+
+julia> eltype(O), size(O)
+(Float32, (2, 3, 4, 5))
+```
+"""
+function similar_map(prototype, ::Type{T}, codomain_axes, domain_axes) where {T}
+    return similar(prototype, T, (codomain_axes..., conj.(domain_axes)...))
+end
+function similar_map(prototype, codomain_axes, domain_axes)
+    return similar_map(prototype, eltype(prototype), codomain_axes, domain_axes)
+end