Add gram_eigh_full named overloads + fix Bijection ordering

Project.toml

@@ -1,6 +1,6 @@
 name = "NamedDimsArrays"
 uuid = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
-version = "0.15.4"
+version = "0.15.5"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -12,6 +12,7 @@ Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 FunctionImplementations = "7c7cc465-9c6a-495f-bdd1-f42428e86d0c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
@@ -41,6 +42,7 @@ FillArrays = "1.13"
 FunctionImplementations = "0.4"
 LinearAlgebra = "1.10"
 Mooncake = "0.4.202, 0.5"
+OrderedCollections = "1.6"
 Random = "1.10"
 SimpleTraits = "0.9.4"
 TensorAlgebra = "0.8, 0.9"

src/nameddimsoperator.jl

@@ -1,3 +1,5 @@
+using OrderedCollections: OrderedDict
+
 # Named dimension operator minimal interface.
 
 # Choi state representation of the named operator.
@@ -64,13 +66,13 @@ function product(x::AbstractNamedDimsArray, y::AbstractNamedDimsArray)
 end
 
 struct Bijection{Codomain, Domain} <: AbstractDict{Domain, Codomain}
-    domain_to_codomain::Dict{Domain, Codomain}
-    codomain_to_domain::Dict{Codomain, Domain}
+    domain_to_codomain::OrderedDict{Domain, Codomain}
+    codomain_to_domain::OrderedDict{Codomain, Domain}
 end
 function Bijection(domain, codomain)
     pairs = domain .=> codomain
-    domain_to_codomain = Dict(pairs)
-    codomain_to_domain = Dict(reverse(kv) for kv in pairs)
+    domain_to_codomain = OrderedDict(pairs)
+    codomain_to_domain = OrderedDict(reverse(kv) for kv in pairs)
     return Bijection(domain_to_codomain, codomain_to_domain)
 end
 function Base.get(b::Bijection, k, default)
@@ -79,8 +81,10 @@ end
 function inverse(b::Bijection)
     return Bijection(b.codomain_to_domain, b.domain_to_codomain)
 end
+# Both accessors iterate `codomain_to_domain` so that successive calls return
+# values in lock-step positional order (codomain[i] paired with domain[i]).
 function codomain(b::Bijection)
-    return values(b.domain_to_codomain)
+    return keys(b.codomain_to_domain)
 end
 function domain(b::Bijection)
     return values(b.codomain_to_domain)
@@ -118,6 +122,90 @@ for f in MATRIX_FUNCTIONS
         end
     end
 end
+
+# Operator entries for the gram factorizations defined in `tensoralgebra.jl`.
+# Placed here because `AbstractNamedDimsOperator` is defined below
+# `tensoralgebra.jl` in the include order.
+#
+# Per-method docstrings are factored out into `const` strings and attached
+# inside the `@eval` loop via `@doc`. This keeps the loop body uniform when
+# methods need distinct user-facing docs (including jldoctest examples) that
+# don't share enough structure to warrant `$($f)`-interpolation.
+
+const _gram_eigh_full_operator_docstring = """
+    TensorAlgebra.gram_eigh_full(a::AbstractNamedDimsOperator; kwargs...) -> x
+
+Gram factorization of a Hermitian positive semi-definite named operator
+`a`, returning `x` such that `conj(x) * x_dom ≈ state(a)`, where `x_dom`
+is `x` with its codomain dimension names replaced by the corresponding
+domain names of `a`. The codomain and domain partition is taken from
+`codomainnames(a)` and `domainnames(a)`.
+
+`kwargs` are forwarded to `TensorAlgebra.gram_eigh_full` on the
+underlying named array (e.g. `atol`, `rtol`).
+
+# Examples
+
+```jldoctest
+julia> using NamedDimsArrays: namedoneto, operator, replacedimnames, state
+
+julia> using TensorAlgebra: gram_eigh_full
+
+julia> i, j, k, l, aux = namedoneto.((2, 2, 2, 2, 8), ("i", "j", "k", "l", "aux"));
+
+julia> b = randn(aux, i, k);
+
+julia> a = operator(conj(b) * replacedimnames(b, "i" => "j", "k" => "l"), ("i", "k"), ("j", "l"));
+
+julia> x = gram_eigh_full(a);
+
+julia> conj(x) * replacedimnames(x, "i" => "j", "k" => "l") ≈ state(a)
+true
+```
+"""
+
+const _gram_eigh_full_with_pinv_operator_docstring = """
+    TensorAlgebra.gram_eigh_full_with_pinv(a::AbstractNamedDimsOperator; kwargs...) -> x, y
+
+Like `TensorAlgebra.gram_eigh_full`, but additionally returns a
+named array `y` such that `x * y` projects onto the rank subspace
+(equal to the identity when `a` is full rank). The codomain and domain
+partition is taken from `codomainnames(a)` and `domainnames(a)`.
+
+# Examples
+
+```jldoctest
+julia> using LinearAlgebra: I
+
+julia> using NamedDimsArrays: dename, dimnames, namedoneto, operator, replacedimnames
+
+julia> using TensorAlgebra: gram_eigh_full_with_pinv
+
+julia> i, j, k, l, aux = namedoneto.((2, 2, 2, 2, 8), ("i", "j", "k", "l", "aux"));
+
+julia> b = randn(aux, i, k);
+
+julia> a = operator(conj(b) * replacedimnames(b, "i" => "j", "k" => "l"), ("i", "k"), ("j", "l"));
+
+julia> x, y = gram_eigh_full_with_pinv(a);
+
+julia> rname = only(setdiff(dimnames(x), ("i", "k")));
+
+julia> reshape(dename(x, (rname, "i", "k")), :, 4) *
+       reshape(dename(y, ("i", "k", rname)), 4, :) ≈ I
+true
+```
+"""
+
+for f in (:gram_eigh_full, :gram_eigh_full_with_pinv)
+    doc_sym = Symbol("_", f, "_operator_docstring")
+    @eval begin
+        @doc $doc_sym function TA.$f(a::AbstractNamedDimsOperator; kwargs...)
+            return TA.$f(state(a), codomainnames(a), domainnames(a); kwargs...)
+        end
+    end
+end
+
 struct NamedDimsOperator{T, N, P <: AbstractNamedDimsArray{T, N}, D, C} <:
     AbstractNamedDimsOperator{T, N}
     parent::P

src/tensoralgebra.jl

@@ -350,6 +350,109 @@ function right_null_nameddims(a::AbstractArray, dimnames_codomain; kwargs...)
     return TA.right_null(a, codomain, domain; kwargs...)
 end
 
+"""
+    TensorAlgebra.gram_eigh_full(a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs...) -> x
+
+Gram factorization of a Hermitian positive semi-definite named array `a`,
+returning `x` such that `a ≈ conj(x) * x_dom`, where `x_dom` is `x` with
+its codomain dimension names replaced by the corresponding domain names.
+The new rank dimension is given a fresh name.
+
+`kwargs` are forwarded to `TensorAlgebra.gram_eigh_full` on the underlying
+unnamed array (e.g. `atol`, `rtol`).
+
+# Examples
+
+```jldoctest
+julia> using NamedDimsArrays: dimnames, namedoneto, replacedimnames
+
+julia> using TensorAlgebra: gram_eigh_full
+
+julia> i, j, k, l, aux = namedoneto.((2, 2, 2, 2, 8), ("i", "j", "k", "l", "aux"));
+
+julia> b = randn(aux, i, k);
+
+julia> a = conj(b) * replacedimnames(b, "i" => "j", "k" => "l");
+
+julia> x = gram_eigh_full(a, (i, k), (j, l));
+
+julia> conj(x) * replacedimnames(x, "i" => "j", "k" => "l") ≈ a
+true
+```
+"""
+function TA.gram_eigh_full(
+        a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs...
+    )
+    return gram_eigh_full_nameddims(a, dimnames_codomain, dimnames_domain; kwargs...)
+end
+function gram_eigh_full_nameddims(
+        a::AbstractArray, dimnames_codomain, dimnames_domain; kwargs...
+    )
+    codomain = name.(dimnames_codomain)
+    domain = name.(dimnames_domain)
+    x_denamed = TA.gram_eigh_full(denamed(a), dimnames(a), codomain, domain; kwargs...)
+    name_x = randname(dimnames(a, 1))
+    dimnames_x = (name_x, codomain...)
+    return nameddims(x_denamed, dimnames_x)
+end
+
+"""
+    TensorAlgebra.gram_eigh_full_with_pinv(a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs...) -> x, y
+
+Like `TensorAlgebra.gram_eigh_full`, but additionally returns a
+named array `y` such that `x * y` projects onto the rank subspace
+(equal to the identity when `a` is full rank). `x` has the rank-name
+first, `y` has it last, both sharing the codomain dimension names of
+`a`.
+
+# Examples
+
+```jldoctest
+julia> using LinearAlgebra: I
+
+julia> using NamedDimsArrays: dename, dimnames, namedoneto, replacedimnames
+
+julia> using TensorAlgebra: gram_eigh_full_with_pinv
+
+julia> i, j, k, l, aux = namedoneto.((2, 2, 2, 2, 8), ("i", "j", "k", "l", "aux"));
+
+julia> b = randn(aux, i, k);
+
+julia> a = conj(b) * replacedimnames(b, "i" => "j", "k" => "l");
+
+julia> x, y = gram_eigh_full_with_pinv(a, (i, k), (j, l));
+
+julia> conj(x) * replacedimnames(x, "i" => "j", "k" => "l") ≈ a
+true
+
+julia> rname = only(setdiff(dimnames(x), ("i", "k")));
+
+julia> reshape(dename(x, (rname, "i", "k")), :, 4) *
+       reshape(dename(y, ("i", "k", rname)), 4, :) ≈ I
+true
+```
+"""
+function TA.gram_eigh_full_with_pinv(
+        a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs...
+    )
+    return gram_eigh_full_with_pinv_nameddims(
+        a, dimnames_codomain, dimnames_domain; kwargs...
+    )
+end
+function gram_eigh_full_with_pinv_nameddims(
+        a::AbstractArray, dimnames_codomain, dimnames_domain; kwargs...
+    )
+    codomain = name.(dimnames_codomain)
+    domain = name.(dimnames_domain)
+    x_denamed, y_denamed = TA.gram_eigh_full_with_pinv(
+        denamed(a), dimnames(a), codomain, domain; kwargs...
+    )
+    name_xy = randname(dimnames(a, 1))
+    dimnames_x = (name_xy, codomain...)
+    dimnames_y = (codomain..., name_xy)
+    return nameddims(x_denamed, dimnames_x), nameddims(y_denamed, dimnames_y)
+end
+
 const MATRIX_FUNCTIONS = [
     :exp, :cis, :log, :sqrt, :cbrt, :cos, :sin, :tan, :csc, :sec, :cot, :cosh, :sinh,
     :tanh,

test/test_operator.jl

@@ -1,5 +1,7 @@
+using LinearAlgebra: I, norm
 using NamedDimsArrays: NamedDimsArrays as NDA, NamedDimsArray, NamedDimsOperator, apply,
-    dimnames, namedoneto, operator, product, replacedimnames, state
+    denamed, dimnames, nameddims, namedoneto, operator, product, replacedimnames, state
+using TensorAlgebra: gram_eigh_full, gram_eigh_full_with_pinv
 using Test: @test, @testset
 
 @testset "operator" begin
@@ -39,3 +41,25 @@ using Test: @test, @testset
     @test issetequal(dimnames(ov), ("i", "j"))
     @test ov ≈ replacedimnames(o * v, "i'" => "i", "j'" => "j")
 end
+
+@testset "gram_eigh_full on AbstractNamedDimsOperator" begin
+    n = 5
+    B = randn(n, n)
+    A = B * B'  # Hermitian PSD
+    M_nda = nameddims(A, ("ket", "bra"))
+    M_op = operator(M_nda, ["ket"], ["bra"])
+
+    X_op = gram_eigh_full(M_op)
+    X_arr = gram_eigh_full(M_nda, ("ket",), ("bra",))
+    # Operator entry forwards to the named-array entry: same data, same shape.
+    @test size(parent(X_op)) == size(parent(X_arr))
+
+    Xp = parent(X_op)
+    @test Xp' * Xp ≈ A
+
+    X2, Y2 = gram_eigh_full_with_pinv(M_op)
+    Xp2 = parent(X2)
+    Yp2 = parent(Y2)
+    @test Xp2' * Xp2 ≈ A
+    @test Xp2 * Yp2 ≈ I(n)
+end

test/test_tensoralgebra.jl

@@ -1,8 +1,10 @@
-using LinearAlgebra: factorize, lq, norm, qr, svd
-using NamedDimsArrays: NamedDimsArrays, dename, denamed, inds, namedoneto
+using LinearAlgebra: LinearAlgebra, factorize, lq, norm, qr, svd
+using NamedDimsArrays:
+    NamedDimsArrays, dename, denamed, dimnames, inds, namedoneto, randname, replacedimnames
 using StableRNGs: StableRNG
-using TensorAlgebra: TensorAlgebra, contract, left_null, left_orth, left_polar, matricize,
-    orth, polar, right_null, right_orth, right_polar, unmatricize
+using TensorAlgebra: TensorAlgebra, contract, gram_eigh_full, gram_eigh_full_with_pinv,
+    left_null, left_orth, left_polar, matricize, orth, polar, right_null, right_orth,
+    right_polar, unmatricize
 using Test: @test, @test_broken, @testset
 
 @testset "TensorAlgebra (eltype=$(elt))" for elt in
@@ -119,4 +121,34 @@ using Test: @test, @test_broken, @testset
             @test norm(n * a) ≈ 0
         end
     end
+    @testset "gram_eigh_full" begin
+        # Build a Hermitian PSD a ≈ conj(b) * b over an aux dim, with codomain
+        # (i, k) and domain (j, l) sharing the same axis lengths.
+        i, j, k, l, aux = namedoneto.((2, 2, 2, 2, 5), ("i", "j", "k", "l", "aux"))
+        b = randn(elt, aux, i, k)
+        # conj(b) * b with the non-conjugated copy's (i, k) relabeled to
+        # (j, l) to form the operator-shaped Hermitian a ≈ X' * X.
+        b_dom = replacedimnames(b, "i" => "j", "k" => "l")
+        a = conj(b) * b_dom
+
+        let X = gram_eigh_full(a, (i, k), (j, l))
+            X_dom = replacedimnames(X, "i" => "j", "k" => "l")
+            @test (i, k) ⊆ inds(X)
+            @test conj(X) * X_dom ≈ a
+        end
+
+        let (X, Y) = gram_eigh_full_with_pinv(a, (i, k), (j, l))
+            rank_name = only(setdiff(dimnames(X), ("i", "k")))
+            @test rank_name == only(setdiff(dimnames(Y), ("i", "k")))
+            X_dom = replacedimnames(X, "i" => "j", "k" => "l")
+            @test conj(X) * X_dom ≈ a
+            # Rename one rank dimension so `X * Y` contracts only on
+            # the shared codomain names `(i, k)` and leaves a
+            # (rank × rank) named identity.
+            fresh_rank = randname(rank_name)
+            Y_fresh = replacedimnames(Y, rank_name => fresh_rank)
+            XYmat = dename(X * Y_fresh, (rank_name, fresh_rank))
+            @test XYmat ≈ LinearAlgebra.I(size(XYmat, 1))
+        end
+    end
 end