Balanced gram_eigh_full convention, add one and operator-construction primitives (#177)

## Summary

- `gram_eigh_full` and `gram_eigh_full_with_pinv` flipped to the
balanced `A ≈ X * X'` convention (was right-Gram `A ≈ X' * X`).
Breaking, but lands as a patch bump since `gram_eigh_full` is new in
v0.9.4 with no known downstream users.
- `TensorAlgebra.one(A, codomain, domain)` for identity operator
tensors. Not exported (clashes with `Base.one`).
- `similar_map(prototype, [T,] codomain_axes, domain_axes)` for
allocating linear-map-shaped arrays. `NamedDimsArrays.similar_operator`
routes through this
(ITensor/NamedDimsArrays.jl#229).
- `projectto!` and `checked_projectto!` for projecting into a restricted
subspace.
- `project_map` and `checked_project_map` as `similar_map` +
`projectto!` allocators.
- `trivialrange(::Type{<:AbstractUnitRange})` for the identity range
under `tensor_product_axis`. Defaults to `Base.OneTo(1)`, overloaded by
downstream packages (for example a charge-0 graded range).
- The `!!` (may-destroy-input) factorization variants in `MatrixAlgebra`
are no longer exported and no longer appear in the `gram_eigh_full`
docstrings, treating them as internal. They remain callable as qualified
`MatrixAlgebra.svd!!`-style functions. Nothing outside TensorAlgebra
used them.

---------

Co-authored-by: Claude <noreply@anthropic.com>

Author: mtfishman

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

@@ -1,36 +1,25 @@
 module MatrixAlgebra
 
 export eigen,
-    eigen!!,
     eigvals,
-    eigvals!!,
     factorize,
-    factorize!!,
     gram_eigh_full,
-    gram_eigh_full!!,
     gram_eigh_full_with_pinv,
-    gram_eigh_full_with_pinv!!,
     invsqrt_diag_safe,
     invsqrth_safe,
     lq,
-    lq!!,
     orth,
-    orth!!,
     polar,
-    polar!!,
     pow_diag_safe,
     powh_safe,
     qr,
-    qr!!,
     sqrt_diag_safe,
     sqrth_safe,
     svd,
-    svd!!,
-    svdvals,
-    svdvals!!
+    svdvals
 
-import MatrixAlgebraKit as MAK
 using LinearAlgebra: LinearAlgebra, Diagonal, isdiag, norm
+using MatrixAlgebraKit: MatrixAlgebraKit as MAK
 
 for (f, f_full, f_compact) in (
         (:qr, :qr_full, :qr_compact),
@@ -206,24 +195,23 @@ 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
 
 """
     gram_eigh_full(A::AbstractMatrix; alg=nothing, atol=0, rtol=eps(real(eltype(A)))^(3//4)) -> X
-    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.
 
 ## Keyword arguments
 
@@ -242,22 +230,21 @@ julia> A = B' * B;
 
 julia> X = gram_eigh_full(A);
 
-julia> X' * X ≈ A
+julia> X * X' ≈ A
 true
 ```
 
 See also [`gram_eigh_full_with_pinv`](@ref).
 """
-gram_eigh_full, 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
-    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.
 
 ## Keyword arguments
 
@@ -278,14 +265,14 @@ 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
 ```
 """
-gram_eigh_full_with_pinv, gram_eigh_full_with_pinv!!
+gram_eigh_full_with_pinv
 
 for (svd, svd_trunc, svd_full, svd_compact) in (
         (:svd, :svd_trunc, :svd_full, :svd_compact),

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/blockedtuple.jl

@@ -110,6 +110,10 @@ function Base.map(f, bt::AbstractBlockTuple)
     return widened_constructorof(typeof(bt))(t, Val(blocklengths(bt)))
 end
 
+function Base.invperm(bt::AbstractBlockTuple)
+    return widened_constructorof(typeof(bt))(invperm(Tuple(bt)), Val(blocklengths(bt)))
+end
+
 function Base.show(io::IO, bt::AbstractBlockTuple)
     return print(io, nameof(typeof(bt)), blocks(bt))
 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(S, 1),), (axes(S, 2),)))
     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),), conj.(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/linearbroadcasted.jl

@@ -1,5 +1,5 @@
-import Base.Broadcast as BC
-import LinearAlgebra as LA
+using Base.Broadcast: Broadcast as BC
+using LinearAlgebra: LinearAlgebra as LA
 
 # TermInterface-like interface.
 iscall(x) = false

src/matricize.jl

@@ -77,6 +77,19 @@ function tensor_product_fusionstyle(r1::AbstractUnitRange, r2::AbstractUnitRange
     return FusionStyle(FusionStyle(r1), FusionStyle(r2))
 end
 
+"""
+    TensorAlgebra.trivialrange(R::Type{<:AbstractUnitRange})
+    TensorAlgebra.trivialrange(r::AbstractUnitRange)
+
+Return the identity range for `tensor_product_axis` on ranges of type `R`,
+i.e. a one-dimensional range `t` for which fusing `t` with any other range
+of the same family leaves that range unchanged. Defaults to `Base.OneTo(1)`;
+downstream packages overload the type-level method to return their own
+identity (for example, a charge-0 one-dimensional sector for a graded range).
+"""
+trivialrange(r::AbstractUnitRange) = trivialrange(typeof(r))
+trivialrange(::Type{<:AbstractUnitRange}) = Base.OneTo(1)
+
 function fused_axis(
         style::FusionStyle, side::Val{:codomain}, a::AbstractArray,
         axes_codomain::Tuple{Vararg{AbstractUnitRange}},

src/permutedimsadd.jl

@@ -1,6 +1,6 @@
-import StridedViews as SV
 using FunctionImplementations: permuteddims
 using Strided: Strided
+using StridedViews: StridedViews as SV
 
 # Specify if an array is on CPU. This is helpful for backends that don't support
 # operations on GPU, such as Strided.jl.

src/projectto.jl

@@ -0,0 +1,48 @@
+"""
+    projectto!(dest, src) -> dest
+
+Project `src` into the restricted space of `dest` without checking which
+components may have been projected out. Defaults to `copyto!`. See
+[`checked_projectto!`](@ref) for a checked version.
+"""
+projectto!(dest, src) = copyto!(dest, src)
+
+"""
+    checked_projectto!(dest, src; kwargs...) -> dest
+
+Project `src` into the restricted space of `dest` via [`projectto!`](@ref)
+and verify via `isapprox(src, dest; kwargs...)` that the discarded
+component is within tolerance. Keyword arguments are forwarded to
+`isapprox`. The default tolerances are subject to change in future
+versions.
+"""
+function checked_projectto!(dest, src; kwargs...)
+    projectto!(dest, src)
+    isapprox(src, dest; kwargs...) ||
+        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). See [`checked_project_map`](@ref) for
+a checked version.
+"""
+function project_map(raw, codomain_axes, domain_axes)
+    return projectto!(similar_map(raw, codomain_axes, domain_axes), raw)
+end
+
+"""
+    checked_project_map(raw, codomain_axes, domain_axes; kwargs...) -> dest
+
+Allocate a map-shaped array via [`similar_map`](@ref) and project `raw`
+into it with [`checked_projectto!`](@ref). Keyword arguments are forwarded
+to [`checked_projectto!`](@ref).
+"""
+function checked_project_map(raw, codomain_axes, domain_axes; kwargs...)
+    return checked_projectto!(
+        similar_map(raw, codomain_axes, domain_axes), raw; kwargs...
+    )
+end

src/similar_map.jl

@@ -0,0 +1,27 @@
+"""
+    similar_map(prototype, [T,] codomain_axes, domain_axes) -> M
+
+Allocate an array shaped as a linear map from `domain_axes` to
+`codomain_axes` with element type `T` (defaulting to `eltype(prototype)`),
+using `prototype` to determine the array backend. Defaults to
+`similar(prototype, T, (codomain_axes..., conj.(domain_axes)...))`.
+
+# Examples
+
+```jldoctest
+julia> using TensorAlgebra: similar_map
+
+julia> cod, dom = (Base.OneTo(2), Base.OneTo(3)), (Base.OneTo(4), Base.OneTo(5));
+
+julia> M = similar_map(randn(3), Float32, cod, dom);
+
+julia> eltype(M), size(M)
+(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