Refactor Val{2} √S split via sqrt(S, co, dom) + replacedimnames

Drop the local `balanced_eigh_factorization` stand-in in favor of using
NamedDimsArrays' existing `Base.sqrt(::NDA, codomain, domain)` (single
matrix-sqrt named array) directly, splitting the result into two factors
at the call site via `replacedimnames`. The "transposition-via-relabel"
on `cache![v1 => v2]` (swap the codomain/domain name slots, then fresh)
ensures each directed sqrt-message has the correct arrow direction on
its matching leg; for dense backings sqrt_S equals its transpose so the
swap is numerically a no-op, but the distinction matters for graded /
fermionic axes.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: mtfishman

src/apply/apply_operators.jl

@@ -268,22 +268,42 @@ function apply_gate_bp_nsite!(
         S = S / norm(S)
     end
     name_v1, name_v2 = dimnames(S)
-    sqrt_S_v1, sqrt_S_v2 = balanced_eigh_factorization(S, (name_v1,), (name_v2,))
-    R_v1 = U_v1 * sqrt_S_v1
-    R_v2 = sqrt_S_v2 * U_v2
+    # `sqrt(S, (name_v1,), (name_v2,))` is NDA's matrix sqrt of `S` —
+    # a single 2-leg named array with dimnames `(name_v1, name_v2)`
+    # satisfying `sqrt_S * sqrt_S ≈ S` in the matrix algebra (each
+    # `sqrt_S` factor contracts on one of `S`'s legs). Eventual endpoint:
+    # 1-arg `sqrt(S)` once `TA.svd` returns `S` as a `NamedDimsOperator`.
+    sqrt_S = sqrt(S, (name_v1,), (name_v2,))
+    # Build R factors by absorbing `sqrt_S` on each side; the rebind on
+    # the v1 side picks `name_v1` as the new shared bond between
+    # `dest[v1]` and `dest[v2]`. With a `NamedDimsOperator` wrapper, the
+    # rebind becomes `apply(sqrt_S, U_v1)`.
+    R_v1 = replacedimnames(U_v1 * sqrt_S, name_v2 => name_v1)
+    R_v2 = sqrt_S * U_v2
 
     dest[v1] = prod([[Q_v1 * R_v1]; inv_sqrt_envs_v1])
     dest[v2] = prod([[Q_v2 * R_v2]; inv_sqrt_envs_v2])
 
-    # Reuse the two `sqrt_S` factors as new sqrt-messages, rebinding the
-    # outer (SVD-codomain / SVD-domain) leg to a fresh name per directed
-    # edge so the two messages don't share a leg name. Each direction
-    # picks the factor whose shared-bond arrow contracts with the
-    # receiving tensor: `sqrt_S_v1`'s bond arrow contracts with `dest[v2]`
-    # (v1 => v2), `sqrt_S_v2`'s with `dest[v1]` (v2 => v1). For dense
-    # backings the two factors carry the same data and the choice is
-    # invisible; the distinction matters for graded / fermionic axes.
-    cache![v1 => v2] = replacedimnames(sqrt_S_v1, name_v1 => randname(name_v1))
-    cache![v2 => v1] = replacedimnames(sqrt_S_v2, name_v2 => randname(name_v2))
+    # Both directed sqrt-messages derive from the same `sqrt_S`, but
+    # with different name-slot choices so each message's "matching" leg
+    # (name_v1, contracting with the receiving tensor) carries the
+    # correct arrow direction.
+    #
+    # `dest[v1]`'s name_v1 bond inherits the domain-side arrow of `S`
+    # (from the `name_v2 => name_v1` rebind in `R_v1`), and `dest[v2]`'s
+    # name_v1 bond inherits the codomain-side arrow (from `sqrt_S * U_v2`).
+    # So:
+    #   * `cache![v2 => v1]`'s matching leg needs the codomain-side arrow
+    #     → use sqrt_S's name_v1 leg directly; relabel name_v2 to fresh.
+    #   * `cache![v1 => v2]`'s matching leg needs the domain-side arrow
+    #     → swap roles: rename sqrt_S's name_v2 to name_v1, and the
+    #     original name_v1 (now the internal-rank slot) to a fresh name.
+    # For dense backings sqrt_S equals its transpose, so the two choices
+    # coincide numerically; the distinction matters for graded /
+    # fermionic axes.
+    cache![v1 => v2] = replacedimnames(
+        sqrt_S, name_v1 => randname(name_v1), name_v2 => name_v1
+    )
+    cache![v2 => v1] = replacedimnames(sqrt_S, name_v2 => randname(name_v2))
     return dest
 end

src/apply/tensoralgebra.jl

@@ -89,57 +89,17 @@ function identity_map(::Type{T}, codomain_axes, domain_axes) where {T}
     return reshape(Matrix{T}(I, n_co, n_dom), (co_lens..., dom_lens...))
 end
 
-# === balanced_eigh_factorization ===
-#
-# Balanced eigh-based factorization of a Hermitian PSD named array `a`:
-# returns `(X, Y)` with `X * Y ≈ a` via named contraction, sharing a
-# fresh-named bond. For k-codomain input, `X` has names
-# `(codomain..., new_bond)` and `Y` has names `(new_bond, domain...)`.
-#
-# Conceptually: `a = U Λ U†` via eigh, then split Λ = √Λ · √Λ symmetrically
-# between the two halves so `X = U √Λ` and `Y = √Λ U†`. For
-# diagonal-Hermitian-PSD input (the BP simple-update SVD-`S` case),
-# eigh is trivial and this reduces to the per-element √ split.
-#
-# Layered through `TA.matricize` → matrix `sqrt` → `TA.unmatricize`,
-# matching the shape of `inv_regularized` above. The N-d / TA layer
-# is namespaced locally (intended `TA.balanced_eigh_factorization`),
-# the named layer extends here. See `gate_application/Overview.md` in
-# `ITensorDevelopmentPlans` for the operator-design synthesis this
-# slots into (`balanced_eigh_factor` single-factor companion,
-# `cholesky_factor`, `positive_factor` umbrella).
-
-function balanced_eigh_factorization(
-        style::TA.FusionStyle, A::AbstractArray, ndims_codomain::Val
-    )
-    M = TA.matricize(style, A, ndims_codomain)
-    sqrtM = sqrt(M)
-    biperm = TA.trivialbiperm(ndims_codomain, Val(ndims(A)))
-    axes_codomain, axes_domain = TA.blocks(axes(A)[biperm])
-    bond_axis = axes(sqrtM, 2)
-    return (
-        TA.unmatricize(style, sqrtM, axes_codomain, (bond_axis,)),
-        TA.unmatricize(style, sqrtM, (bond_axis,), axes_domain),
-    )
-end
-function balanced_eigh_factorization(A::AbstractArray, ndims_codomain::Val)
-    return balanced_eigh_factorization(TA.FusionStyle(A), A, ndims_codomain)
-end
-
-function balanced_eigh_factorization(
-        a::AbstractNamedDimsArray, codomain_dimnames, domain_dimnames
-    )
-    codomain_names = collect(name.(codomain_dimnames))
-    domain_names = collect(name.(domain_dimnames))
-    biperm = TA.blockedperm_indexin(
-        Tuple.((dimnames(a), codomain_names, domain_names))...
-    )
-    perm_codomain, perm_domain = TA.blocks(biperm)
-    A_perm = TA.bipermutedims(denamed(a), perm_codomain, perm_domain)
-    X_denamed, Y_denamed = balanced_eigh_factorization(A_perm, Val(length(perm_codomain)))
-    new_bond = randname(first(codomain_names))
-    return (
-        nameddims(X_denamed, [codomain_names; [new_bond]]),
-        nameddims(Y_denamed, [[new_bond]; domain_names]),
-    )
-end
+# Note: the BP simple-update `√S` split uses NDA's existing
+# `Base.sqrt(::AbstractNamedDimsArray, codomain_dimnames,
+# domain_dimnames)` (matrix sqrt as a single named array) directly,
+# combined with explicit `replacedimnames` at the call site to split
+# the result into two factors sharing a fresh bond. See the comment in
+# `apply_gate_bp_nsite!` (Val{2} method) for the call-site
+# choreography. A tuple-returning `factorize_sqrt` primitive — splitting
+# a Hermitian PSD `M` into `(X, Y)` with a fresh shared bond — was
+# previously staged here as a local stand-in but isn't needed for the
+# current `√S` use case (K=1 codomain). It can be reintroduced when a
+# multi-codomain (K>1) factorization use case lands, alongside the
+# rest of the `factorize_<backend>` family
+# (`factorize_balanced_eigh`, `factorize_cholesky`) discussed in
+# `gate_application/Overview.md` in `ITensorDevelopmentPlans`.