Refactor initialize_cache to one(similar_operator(...)) form

Replace the dense `identity_map` helper with two composable primitives:

  * `similar_operator(prototype, codomain_axes)` — undef
    `NamedDimsOperator` with codomain = input axes, domain = same axes
    fresh-renamed. Backend / eltype propagates from `prototype` via
    `Base.similar`.
  * `Base.one(::AbstractNamedDimsOperator)` — identity operator via
    matricize → fill with `I` → unmatricize → rewrap.

`initialize_cache` reduces to `state(one(similar_operator(factor,
linkaxes(iterate, edge))))` per edge.

Whitelist `Base.one` in `test_aqua.jl` as a stand-in extension that
will move upstream into NDA's `MATRIX_FUNCTIONS` operator-extensions
loop.

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

Author: mtfishman

src/apply/apply_operators.jl

@@ -5,8 +5,8 @@ import TensorAlgebra as TA
 using Base: @kwdef
 using Graphs: dst, src, vertices
 using LinearAlgebra: norm
-using NamedDimsArrays:
-    AbstractNamedDimsArray, dimnames, domainnames, nameddims, randname, replacedimnames
+using NamedDimsArrays: AbstractNamedDimsArray, dimnames, domainnames, nameddims, randname,
+    replacedimnames, state
 using NamedGraphs.GraphsExtensions: all_edges, boundary_edges
 
 # === NestedAlgorithm framework ===
@@ -159,20 +159,13 @@ function AI.initialize_state!(
     return state
 end
 
-# Identity-message cache: trivial Vidal-gauge initialization where each bond
-# carries the identity 2-leg map (= √I = I, in sqrt-message form). Stored
-# in a `SqrtMessageCache` so the BP simple update knows to use the messages
-# as gauge-in factors directly and skip the √ step.
+# Initialize the BP message cache to identity square-root messages.
 function initialize_cache(
-        problem::ApplyOperatorProblem, ::BPApplyGate, iterate::AbstractTensorNetwork
+        ::ApplyOperatorProblem, ::BPApplyGate, iterate::AbstractTensorNetwork
     )
-    T = eltype(iterate[first(vertices(iterate))])
     return sqrtmessagecache(all_edges(iterate)) do edge
-        bond_name = only(linknames(iterate, edge))
-        bond_axis = only(linkaxes(iterate, edge))
-        fresh_name = randname(bond_name)
-        A = identity_map(T, (bond_axis,), (bond_axis,))
-        return nameddims(A, (fresh_name, bond_name))
+        factor = iterate[dst(edge)]
+        return state(one(similar_operator(factor, linkaxes(iterate, edge))))
     end
 end
 

src/apply/tensoralgebra.jl

@@ -22,7 +22,8 @@
 import MatrixAlgebraKit as MAK
 import TensorAlgebra as TA
 using LinearAlgebra: I
-using NamedDimsArrays: AbstractNamedDimsArray, denamed, dimnames, name, nameddims, randname
+using NamedDimsArrays: AbstractNamedDimsArray, AbstractNamedDimsOperator, codomainnames,
+    denamed, dimnames, domainnames, name, nameddims, operator, randname, setname, state
 
 # === N-d / TensorAlgebra layer ===
 
@@ -69,37 +70,26 @@ function MAK.inv_regularized(
     return MAK.inv_regularized(a, codomain_names, domain_names; kwargs...)
 end
 
-# === identity_map ===
-#
-# 2k-leg identity *map* (pairwise δ per (co_i, dom_i)):
-# `I_{co_1, dom_1} ⊗ … ⊗ I_{co_k, dom_k}` reshaped to a 2k-leg tensor.
-#
-# Local stand-in: dense-only. Eventual home is `TensorAlgebra.jl` with
-# an `AbstractNamedDimsArray` overload and axis-type dispatch for the
-# graded / FusionTensor specializations (see
-# `gate_application/Overview.md` in `ITensorDevelopmentPlans`).
-
-function identity_map(::Type{T}, codomain_axes, domain_axes) where {T}
+function similar_operator(prototype::AbstractNamedDimsArray, codomain_axes)
     co_axes = Tuple(codomain_axes)
-    dom_axes = Tuple(domain_axes)
-    co_lens = length.(co_axes)
-    dom_lens = length.(dom_axes)
-    n_co = prod(co_lens; init = 1)
-    n_dom = prod(dom_lens; init = 1)
-    return reshape(Matrix{T}(I, n_co, n_dom), (co_lens..., dom_lens...))
+    dom_axes = setname.(co_axes, randname.(name.(co_axes)))
+    A = similar(denamed(prototype), (co_axes..., dom_axes...))
+    return operator(A, collect(name.(co_axes)), collect(name.(dom_axes)))
 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`.
+function Base.one(a::AbstractNamedDimsOperator)
+    co = codomainnames(a)
+    dom = domainnames(a)
+    A = state(a)
+    A_denamed = denamed(A)
+    style = TA.FusionStyle(A_denamed)
+    ndims_co = Val(length(co))
+    A_mat = TA.matricize(style, A_denamed, ndims_co)
+    id_mat = similar(A_mat)
+    copyto!(id_mat, I)
+    biperm = TA.trivialbiperm(ndims_co, Val(ndims(A_denamed)))
+    co_axes, dom_axes = TA.blocks(axes(A_denamed)[biperm])
+    id_denamed = TA.unmatricize(style, id_mat, co_axes, dom_axes)
+    id_nda = nameddims(id_denamed, dimnames(A))
+    return operator(id_nda, co, dom)
+end

test/test_aqua.jl

@@ -4,12 +4,17 @@ using MatrixAlgebraKit: MatrixAlgebraKit
 using Test: @testset
 
 @testset "Code quality (Aqua.jl)" begin
-    # `MatrixAlgebraKit.inv_regularized` is locally extended for
-    # `AbstractNamedDimsArray` as a stand-in until the corresponding method
-    # moves into `NamedDimsArrays.jl`. Whitelist it for the piracy check.
+    # Stand-in Base / MAK extensions on `AbstractNamedDimsArray` /
+    # `AbstractNamedDimsOperator` that will move upstream into
+    # `NamedDimsArrays.jl` (or its operator extensions). Whitelist them
+    # for the piracy check until the upstream PRs land:
+    # * `MAK.inv_regularized` — N-d pseudo-inverse for named arrays.
+    # * `Base.one` on `AbstractNamedDimsOperator` — identity operator,
+    #   analog of the existing `Base.sqrt` / `Base.exp` / … extensions
+    #   already defined in NDA's `MATRIX_FUNCTIONS` loop.
     Aqua.test_all(
         ITensorNetworksNext;
         persistent_tasks = false,
-        piracies = (; treat_as_own = [MatrixAlgebraKit.inv_regularized])
+        piracies = (; treat_as_own = [MatrixAlgebraKit.inv_regularized, Base.one])
     )
 end