Drop LFN/BFN identity_messages; route auto-init through QFN only

mtfishman · claude · mtfishman · commit 6d09a7186b36 · 2026-05-13T13:12:39.000-04:00
- Make identity_messages a QFN-only method. The QFN version pairs each
  ket Index with its bra counterpart explicitly via dual_index_map(fn),
  so the construction stays correct when multiple link indices share
  an edge (where commoninds-zip ordering across the two layers is not
  guaranteed).
- Drop the function-tag argument from initialize_cache. Auto-init is a
  property of (algorithm, network type) only, and QFN is the only form
  network where identity_messages is canonical, so a single
  initialize_cache(alg::bp, fn::QFN) specialization is enough.
- Change norm_sqr_network(psi) to return QuadraticFormNetwork(psi),
  matching its name. Route LinearAlgebra.norm_sqr and normalize through
  it so they pick up the QFN auto-init on loopy graphs.
- test_apply.jl is structurally LFN-based (apply's local-env expectations
  don't generalize to QFN), so it now builds its own LFN messages
  explicitly via a small _lfn_identity_messages helper.

Co-Authored-By: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/src/abstractitensornetwork.jl b/src/abstractitensornetwork.jl
@@ -517,7 +517,7 @@ function inner_network(
     return BilinearFormNetwork(A, x, y; kwargs...)
 end
 
-norm_sqr_network(ψ::AbstractITensorNetwork) = inner_network(ψ, ψ)
+norm_sqr_network(ψ::AbstractITensorNetwork) = QuadraticFormNetwork(ψ)
 
 #
 # Printing
diff --git a/src/contract.jl b/src/contract.jl
@@ -47,7 +47,7 @@ function logscalar(
         cache_update_kwargs = (;)
     )
     if isnothing(cache!)
-        cache! = Ref(initialize_cache(scalar, alg, tn; cache_construction_kwargs...))
+        cache! = Ref(initialize_cache(alg, tn; cache_construction_kwargs...))
     end
 
     if update_cache
diff --git a/src/environment.jl b/src/environment.jl
@@ -28,9 +28,7 @@ function environment(
         cache_update_kwargs = (;)
     )
     if isnothing(cache!)
-        cache! = Ref(
-            initialize_cache(environment, alg, ptn; cache_construction_kwargs...)
-        )
+        cache! = Ref(initialize_cache(alg, ptn; cache_construction_kwargs...))
     end
 
     if update_cache
diff --git a/src/expect.jl b/src/expect.jl
@@ -30,7 +30,7 @@ function expect(
     )
     ψIψ = QuadraticFormNetwork(ψ)
     if isnothing(cache!)
-        cache! = Ref(initialize_cache(expect, alg, ψIψ; cache_construction_kwargs...))
+        cache! = Ref(initialize_cache(alg, ψIψ; cache_construction_kwargs...))
     end
 
     if update_cache
diff --git a/src/formnetworks/bilinearformnetwork.jl b/src/formnetworks/bilinearformnetwork.jl
@@ -1,11 +1,8 @@
 using Adapt: adapt
 using DataGraphs: DataGraphs, set_vertex_data!
-using Dictionaries: Dictionary, set!
 using ITensors.NDTensors: datatype, denseblocks
-using ITensors: ITensor, Index, Op, commoninds, dag, delta, prime, sim
+using ITensors: ITensor, Index, Op, dag, delta, prime, sim
 using NamedGraphs.GraphsExtensions: disjoint_union
-using NamedGraphs.PartitionedGraphs:
-    PartitionedGraph, QuotientEdge, partitioned_vertices, quotientedges
 
 default_dual_site_index_map = prime
 default_dual_link_index_map = sim
@@ -108,50 +105,3 @@ function update(
     tensornetwork(blf)[ket_vertex(blf, original_ket_state_vertex)] = ket_state
     return blf
 end
-
-# Initial BP messages from bra↔ket pairings on each quotient edge.
-# Errors when the operator subnet has its own inter-vertex links (the
-# multi-site-operator case): those legs are operator-internal and have
-# no bra/ket pair, so no canonical identity initialization exists — the
-# caller must supply `messages` explicitly.
-function identity_messages(fn::BilinearFormNetwork, ptn::PartitionedGraph)
-    pairings = Dictionary{QuotientEdge, Pair{Vector{Index}, Vector{Index}}}()
-    tn = tensornetwork(fn)
-    pv = partitioned_vertices(ptn)
-    ket_s = ket_vertex_suffix(fn)
-    for pe in quotientedges(ptn)
-        src_orig = unique(first.(filter(v -> last(v) == ket_s, pv[parent(src(pe))])))
-        dst_orig = unique(first.(filter(v -> last(v) == ket_s, pv[parent(dst(pe))])))
-        for v_from in src_orig, v_to in dst_orig
-            op_inds = commoninds(
-                tn[operator_vertex(fn, v_from)], tn[operator_vertex(fn, v_to)]
-            )
-            if !isempty(op_inds)
-                error(
-                    "BilinearFormNetwork: operator-internal cross-Index between " *
-                        "$v_from and $v_to has no bra/ket pair; supply `messages` " *
-                        "explicitly to BP."
-                )
-            end
-        end
-        for (from_orig, to_orig, e) in (
-                (src_orig, dst_orig, pe),
-                (dst_orig, src_orig, reverse(pe)),
-            )
-            bras = Index[]
-            kets = Index[]
-            for v_from in from_orig, v_to in to_orig
-                append!(
-                    bras,
-                    commoninds(tn[bra_vertex(fn, v_from)], tn[bra_vertex(fn, v_to)])
-                )
-                append!(
-                    kets,
-                    commoninds(tn[ket_vertex(fn, v_from)], tn[ket_vertex(fn, v_to)])
-                )
-            end
-            set!(pairings, e, bras => kets)
-        end
-    end
-    return identity_messages(scalartype(tn), pairings)
-end
diff --git a/src/formnetworks/linearformnetwork.jl b/src/formnetworks/linearformnetwork.jl
@@ -1,9 +1,6 @@
 using DataGraphs: DataGraphs, set_vertex_data!
-using Dictionaries: Dictionary, set!
-using ITensors: ITensor, Index, commoninds, dag, prime
+using ITensors: ITensor, dag, prime
 using NamedGraphs.GraphsExtensions: disjoint_union
-using NamedGraphs.PartitionedGraphs:
-    PartitionedGraph, QuotientEdge, partitioned_vertices, quotientedges
 
 default_dual_link_index_map = prime
 
@@ -60,42 +57,3 @@ function update(lf::LinearFormNetwork, original_ket_state_vertex, ket_state::ITe
     tensornetwork(lf)[ket_vertex(blf, original_ket_state_vertex)] = ket_state
     return lf
 end
-
-# Initial BP messages on each quotient edge built from the bra↔ket leg
-# pairs crossing that edge: legs are taken from the `from`-partition side
-# of each layer (so the bra leg and the ket leg carry opposite directions
-# because the bra layer is `dag(dual_link_index_map(ket))`), and the
-# forward/reverse messages use opposite-end views so each one's open
-# legs face the correct receiving partition when read during BP updates.
-# Iteration is over Cartesian products of the original ket-graph vertices
-# in each partition, so this works for arbitrary partitionings (per-vertex
-# or coarser groupings such as whole columns).
-function identity_messages(fn::LinearFormNetwork, ptn::PartitionedGraph)
-    pairings = Dictionary{QuotientEdge, Pair{Vector{Index}, Vector{Index}}}()
-    tn = tensornetwork(fn)
-    pv = partitioned_vertices(ptn)
-    ket_s = ket_vertex_suffix(fn)
-    for pe in quotientedges(ptn)
-        src_orig = unique(first.(filter(v -> last(v) == ket_s, pv[parent(src(pe))])))
-        dst_orig = unique(first.(filter(v -> last(v) == ket_s, pv[parent(dst(pe))])))
-        for (from_orig, to_orig, e) in (
-                (src_orig, dst_orig, pe),
-                (dst_orig, src_orig, reverse(pe)),
-            )
-            bras = Index[]
-            kets = Index[]
-            for v_from in from_orig, v_to in to_orig
-                append!(
-                    bras,
-                    commoninds(tn[bra_vertex(fn, v_from)], tn[bra_vertex(fn, v_to)])
-                )
-                append!(
-                    kets,
-                    commoninds(tn[ket_vertex(fn, v_from)], tn[ket_vertex(fn, v_to)])
-                )
-            end
-            set!(pairings, e, bras => kets)
-        end
-    end
-    return identity_messages(scalartype(tn), pairings)
-end
diff --git a/src/formnetworks/quadraticformnetwork.jl b/src/formnetworks/quadraticformnetwork.jl
@@ -1,5 +1,8 @@
 using DataGraphs: DataGraphs, set_vertex_data!, underlying_graph, vertex_data
-using NamedGraphs.PartitionedGraphs: PartitionedGraph
+using Dictionaries: Dictionary, set!
+using ITensors: Index, commoninds, dag
+using NamedGraphs.PartitionedGraphs:
+    PartitionedGraph, QuotientEdge, partitioned_vertices, quotientedges
 
 default_index_map = prime
 default_inv_index_map = noprime
@@ -86,8 +89,37 @@ function QuadraticFormNetwork(
     return QuadraticFormNetwork(blf, dual_index_map, dual_inv_index_map)
 end
 
+# Build initial BP messages on each quotient edge as `delta(bra, ket)`
+# pairs, one per ket link Index crossing the cut. The bra-side counterpart
+# of each ket Index is computed explicitly via `dual_index_map(fn)`, so
+# the pairing is correct even when multiple link indices share an edge
+# (where `commoninds`-zip ordering between layers is not guaranteed).
 function identity_messages(fn::QuadraticFormNetwork, ptn::PartitionedGraph)
-    return identity_messages(bilinear_formnetwork(fn), ptn)
+    pairings = Dictionary{QuotientEdge, Pair{Vector{Index}, Vector{Index}}}()
+    tn = tensornetwork(fn)
+    map_idx = dual_index_map(fn)
+    pv = partitioned_vertices(ptn)
+    ket_s = ket_vertex_suffix(fn)
+    for pe in quotientedges(ptn)
+        src_orig = unique(first.(filter(v -> last(v) == ket_s, pv[parent(src(pe))])))
+        dst_orig = unique(first.(filter(v -> last(v) == ket_s, pv[parent(dst(pe))])))
+        for (from_orig, to_orig, e) in (
+                (src_orig, dst_orig, pe),
+                (dst_orig, src_orig, reverse(pe)),
+            )
+            kets = Index[]
+            bras = Index[]
+            for v_from in from_orig, v_to in to_orig
+                cur_kets = collect(
+                    commoninds(tn[ket_vertex(fn, v_from)], tn[ket_vertex(fn, v_to)])
+                )
+                append!(kets, cur_kets)
+                append!(bras, dag.(map_idx.(cur_kets)))
+            end
+            set!(pairings, e, bras => kets)
+        end
+    end
+    return identity_messages(scalartype(tn), pairings)
 end
 
 function update(qf::QuadraticFormNetwork, original_state_vertex, ket_state::ITensor)
diff --git a/src/initialize_cache.jl b/src/initialize_cache.jl
@@ -1,32 +1,24 @@
 using Dictionaries: Dictionary
 using Graphs: is_tree
 using ITensors.NDTensors: @Algorithm_str, Algorithm
-using ITensors: scalar
-using LinearAlgebra: normalize
 using NamedGraphs.PartitionedGraphs:
     AbstractPartitionedGraph, PartitionedGraph, quotient_graph
 
-# Build a cache appropriate for `f` on `tn` using algorithm `alg`. The
-# `f` tag carries the calling context (e.g. `scalar`, `normalize`,
-# `expect`, `rescale`, `environment`) so per-purpose methods can inject
-# context-specific initialization. The fallback constructs a plain
-# `BeliefPropagationCache` with no message defaults.
-function initialize_cache(f, alg::Algorithm"bp", tn::AbstractITensorNetwork; kwargs...)
+# Build a cache for algorithm `alg` on `tn`. The fallback constructs a
+# plain `BeliefPropagationCache` with no message defaults; the
+# `QuadraticFormNetwork` specialization injects `identity_messages` on
+# loopy quotient graphs (canonical for the structurally ψ-vs-ψ case).
+function initialize_cache(alg::Algorithm"bp", tn::AbstractITensorNetwork; kwargs...)
     return BeliefPropagationCache(tn; kwargs...)
 end
 
-function initialize_cache(
-        f, alg::Algorithm"bp", ptn::AbstractPartitionedGraph; kwargs...
-    )
+function initialize_cache(alg::Algorithm"bp", ptn::AbstractPartitionedGraph; kwargs...)
     return BeliefPropagationCache(ptn; kwargs...)
 end
 
-# Core helper: build a BPC on a form network with `identity_messages`
-# on loopy quotient graphs (empty messages on trees). Used by the
-# per-purpose specializations below where the form network is
-# structurally ψ-vs-ψ, so `identity_messages(fn, ptn)` is canonical.
-function _bp_cache_identity_messages(
-        fn::AbstractFormNetwork;
+function initialize_cache(
+        alg::Algorithm"bp",
+        fn::QuadraticFormNetwork;
         partitioned_vertices = default_partitioned_vertices(fn),
         messages = nothing
     )
@@ -36,27 +28,3 @@ function _bp_cache_identity_messages(
     end
     return BeliefPropagationCache(ptn; messages)
 end
-
-function initialize_cache(
-        ::typeof(scalar), alg::Algorithm"bp", fn::QuadraticFormNetwork; kwargs...
-    )
-    return _bp_cache_identity_messages(fn; kwargs...)
-end
-
-function initialize_cache(
-        ::typeof(normalize), alg::Algorithm"bp", fn::AbstractFormNetwork; kwargs...
-    )
-    return _bp_cache_identity_messages(fn; kwargs...)
-end
-
-function initialize_cache(
-        ::typeof(rescale), alg::Algorithm"bp", fn::AbstractFormNetwork; kwargs...
-    )
-    return _bp_cache_identity_messages(fn; kwargs...)
-end
-
-function initialize_cache(
-        ::typeof(expect), alg::Algorithm"bp", fn::QuadraticFormNetwork; kwargs...
-    )
-    return _bp_cache_identity_messages(fn; kwargs...)
-end
diff --git a/src/inner.jl b/src/inner.jl
@@ -173,7 +173,9 @@ end
 
 # TODO: rename `sqnorm` to match https://github.com/JuliaStats/Distances.jl,
 # or `norm_sqr` to match `LinearAlgebra.norm_sqr`
-LinearAlgebra.norm_sqr(ψ::AbstractITensorNetwork; kwargs...) = inner(ψ, ψ; kwargs...)
+function LinearAlgebra.norm_sqr(ψ::AbstractITensorNetwork; kwargs...)
+    return scalar(norm_sqr_network(ψ); kwargs...)
+end
 
 function LinearAlgebra.norm(ψ::AbstractITensorNetwork; kwargs...)
     return sqrt(abs(real(norm_sqr(ψ; kwargs...))))
diff --git a/src/normalize.jl b/src/normalize.jl
@@ -21,7 +21,7 @@ function rescale(
         kwargs...
     )
     if isnothing(cache!)
-        cache! = Ref(initialize_cache(rescale, alg, tn; cache_construction_kwargs...))
+        cache! = Ref(initialize_cache(alg, tn; cache_construction_kwargs...))
     end
 
     if update_cache
@@ -55,7 +55,7 @@ end
 function LinearAlgebra.normalize(
         alg::Algorithm"exact", tn::AbstractITensorNetwork; kwargs...
     )
-    logn = logscalar(alg, inner_network(tn, tn); kwargs...)
+    logn = logscalar(alg, norm_sqr_network(tn); kwargs...)
     c = inv(exp(logn / (2 * length(vertices(tn)))))
     return map(t -> c * t, tn)
 end
@@ -68,10 +68,9 @@ function LinearAlgebra.normalize(
         cache_update_kwargs = (;),
         cache_construction_kwargs = (;)
     )
-    norm_tn = inner_network(tn, tn)
+    norm_tn = norm_sqr_network(tn)
     if isnothing(cache!)
-        cache! =
-            Ref(initialize_cache(normalize, alg, norm_tn; cache_construction_kwargs...))
+        cache! = Ref(initialize_cache(alg, norm_tn; cache_construction_kwargs...))
     end
     vs = collect(vertices(tn))
     verts = vcat([ket_vertex(norm_tn, v) for v in vs], [bra_vertex(norm_tn, v) for v in vs])
diff --git a/test/test_apply.jl b/test/test_apply.jl
@@ -1,15 +1,43 @@
 using Compat: Compat
-using Graphs: vertices
+using Dictionaries: Dictionary, set!
+using Graphs: dst, src, vertices
 using ITensorNetworks: BeliefPropagationCache, apply, environment, identity_messages,
-    norm_sqr_network, siteinds, update
-using ITensors: ITensors, ITensor, inner, op
+    inner_network, siteinds, update
+using ITensors: ITensors, ITensor, Index, commoninds, dag, inner, op, prime
 using NamedGraphs.NamedGraphGenerators: named_grid
-using NamedGraphs.PartitionedGraphs: PartitionedGraph
+using NamedGraphs.PartitionedGraphs:
+    PartitionedGraph, QuotientEdge, partitioned_vertices, quotientedges
 using SplitApplyCombine: group
 using StableRNGs: StableRNG
 using TensorOperations: TensorOperations
 using Test: @test, @test_throws, @testset
 include("utils.jl")
+
+# Build identity-style initial messages on each quotient edge of an LFN
+# (where bra = `dag(prime(ket))`) by pairing each ket Index `k` with its
+# bra counterpart `dag(prime(k))` directly, so the construction is
+# robust to multiple link indices per edge.
+function _lfn_identity_messages(ψψ, ptn::PartitionedGraph)
+    pairings = Dictionary{QuotientEdge, Pair{Vector{Index}, Vector{Index}}}()
+    pv = partitioned_vertices(ptn)
+    for pe in quotientedges(ptn)
+        src_orig = unique(first.(filter(v -> last(v) == "ket", pv[parent(src(pe))])))
+        dst_orig = unique(first.(filter(v -> last(v) == "ket", pv[parent(dst(pe))])))
+        for (from_orig, to_orig, e) in (
+                (src_orig, dst_orig, pe),
+                (dst_orig, src_orig, reverse(pe)),
+            )
+            kets = Index[]
+            for v_from in from_orig, v_to in to_orig
+                append!(kets, commoninds(ψψ[(v_from, "ket")], ψψ[(v_to, "ket")]))
+            end
+            bras = dag.(prime.(kets))
+            set!(pairings, e, bras => kets)
+        end
+    end
+    return identity_messages(ITensors.scalartype(ψψ), pairings)
+end
+
 @testset "apply" begin
     g_dims = (2, 2)
     g = named_grid(g_dims)
@@ -18,18 +46,19 @@ include("utils.jl")
     rng = StableRNG(1234)
     ψ = random_tensornetwork(rng, s; link_space = χ)
     v1, v2 = (2, 2), (1, 2)
-    ψψ = norm_sqr_network(ψ)
+    ψψ = inner_network(ψ, ψ)
     # Simple Belief Propagation grouping (one bra+ket per partition) gives
     # a product environment around `[v1, v2]`, which is what `apply` requires.
     ptn_SBP = PartitionedGraph(ψψ, group(v -> v[1], vertices(ψψ)))
-    bp_cache = BeliefPropagationCache(ptn_SBP; messages = identity_messages(ψψ, ptn_SBP))
+    bp_cache =
+        BeliefPropagationCache(ptn_SBP; messages = _lfn_identity_messages(ψψ, ptn_SBP))
     bp_cache = update(bp_cache; maxiter = 20)
     envsSBP = environment(bp_cache, [(v1, "bra"), (v1, "ket"), (v2, "bra"), (v2, "ket")])
     # Column-grouping (one whole column per partition) gives a non-product
     # environment; `apply` should reject it.
     ptn_col = PartitionedGraph(ψψ, group(v -> v[1][1], vertices(ψψ)))
     bp_cache_col =
-        BeliefPropagationCache(ptn_col; messages = identity_messages(ψψ, ptn_col))
+        BeliefPropagationCache(ptn_col; messages = _lfn_identity_messages(ψψ, ptn_col))
     bp_cache_col = update(bp_cache_col; maxiter = 20)
     envsGBP = environment(
         bp_cache_col, [(v1, "bra"), (v1, "ket"), (v2, "bra"), (v2, "ket")]
diff --git a/test/test_belief_propagation.jl b/test/test_belief_propagation.jl

Original file line number	Diff line number	Diff line change
`@@ -517,7 +517,7 @@ function inner_network(`
`517`	`517`	`return BilinearFormNetwork(A, x, y; kwargs...)`
`518`	`518`	`end`
`519`	`519`
`520`		`-norm_sqr_network(ψ::AbstractITensorNetwork) = inner_network(ψ, ψ)`
	`520`	`+norm_sqr_network(ψ::AbstractITensorNetwork) = QuadraticFormNetwork(ψ)`
`521`	`521`
`522`	`522`	`#`
`523`	`523`	`# Printing`
Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ function logscalar(`
`47`	`47`	`cache_update_kwargs = (;)`
`48`	`48`	`)`
`49`	`49`	`if isnothing(cache!)`
`50`		`- cache! = Ref(initialize_cache(scalar, alg, tn; cache_construction_kwargs...))`
	`50`	`+ cache! = Ref(initialize_cache(alg, tn; cache_construction_kwargs...))`
`51`	`51`	`end`
`52`	`52`
`53`	`53`	`if update_cache`
Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,7 @@ function expect(`
`30`	`30`	`)`
`31`	`31`	`ψIψ = QuadraticFormNetwork(ψ)`
`32`	`32`	`if isnothing(cache!)`
`33`		`- cache! = Ref(initialize_cache(expect, alg, ψIψ; cache_construction_kwargs...))`
	`33`	`+ cache! = Ref(initialize_cache(alg, ψIψ; cache_construction_kwargs...))`
`34`	`34`	`end`
`35`	`35`
`36`	`36`	`if update_cache`