From 16b77faf1b729938e5eaa137ee956216af10690b Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 13 May 2026 10:10:33 -0400 Subject: [PATCH 1/7] Drop default_message; explicit identity_messages discipline Loopy BP on QN-graded networks was silently NaN-ing because the single-leg delta(i) initial messages collapsed half the QN sectors to empty blocks. Replace the implicit default_message family with an explicit two-leg delta(b, k) identity_messages that pairs bra and ket link inds, keeping the QN sectors aligned through the contractions. BeliefPropagationCache(ptn) now starts with an empty messages dict. Form-network identity_messages(fn, ptn) builds the bra/ket pairings from cross-partition vertex pairs and works for both per-vertex and coarser partitionings. The only auto-init is on QuadraticFormNetwork (structurally psi-vs-psi, so identity_messages is canonical); asymmetric LFN/BFN entries fall through to the generic path and the caller must supply messages on loopy graphs. Adds a 4-cycle QN-conserving regression test that pins the old failure mode by going through scalar(QuadraticFormNetwork(psi); alg = "bp"). Co-Authored-By: Claude Opus 4.7 (1M context) --- src/caches/abstractbeliefpropagationcache.jl | 49 +++++++++-------- src/caches/beliefpropagationcache.jl | 7 +-- src/expect.jl | 16 +++++- src/formnetworks/bilinearformnetwork.jl | 52 +++++++++++++++++- src/formnetworks/linearformnetwork.jl | 44 ++++++++++++++- src/formnetworks/quadraticformnetwork.jl | 32 +++++++++++ src/inner.jl | 2 +- src/normalize.jl | 56 ++++++++++++++++++-- test/test_apply.jl | 12 +++-- test/test_belief_propagation.jl | 52 ++++++++++++++---- test/test_forms.jl | 11 ++-- test/test_normalize.jl | 12 ++++- test/utils.jl | 39 +++----------- 13 files changed, 296 insertions(+), 88 deletions(-) diff --git a/src/caches/abstractbeliefpropagationcache.jl b/src/caches/abstractbeliefpropagationcache.jl index d682b2a7..e153cfef 100644 --- a/src/caches/abstractbeliefpropagationcache.jl +++ b/src/caches/abstractbeliefpropagationcache.jl @@ -1,5 +1,6 @@ using Adapt: Adapt, adapt, adapt_structure using DataGraphs: DataGraphs, underlying_graph, vertex_data +using Dictionaries: Dictionary, dictionary using Graphs: Graphs, IsDirected, dst, src using ITensors: commoninds, delta, dir using LinearAlgebra: diag, dot @@ -34,14 +35,31 @@ function message_diff(message_a::Vector{ITensor}, message_b::Vector{ITensor}) return 1 - f end -function default_message(datatype::Type{<:AbstractArray}, inds_e) - return [adapt(datatype, denseblocks(delta(i))) for i in inds_e] +""" + identity_messages([elt::Type,] pairings) -> Dictionary{QuotientEdge, Vector{ITensor}} + +Build a message dictionary from a `pairings` mapping that lists, for each +quotient edge `e`, the `(bra_inds, ket_inds)` pair that defines the +identity-from-bra-to-ket message on that edge. Each entry is constructed +as `[delta(elt, bra_ind_i, ket_ind_i) for i in ...]`; the bra-side and +ket-side Indices are expected to already carry opposite directions so the +two-leg delta has well-defined QN flow (the bra side of a form network is +constructed via `dag` of the dual-mapped ket, which gives this naturally). + +`pairings` is anything supporting `pairs(...)` iteration yielding +`e => (bra_inds => ket_inds)` entries (e.g. a `Dictionary` or `Dict`). +`elt` defaults to `Float64`; form-network wrappers fill it in from +`scalartype(tensornetwork(fn))`. +""" +function identity_messages(elt::Type, pairings) + return dictionary( + e => ITensor[delta(elt, b, k) for (b, k) in zip(bras, kets)] + for (e, (bras, kets)) in pairs(pairings) + ) end -function default_message(elt::Type{<:Number}, inds_e) - return default_message(Vector{elt}, inds_e) -end -default_messages(ptn::PartitionedGraph) = Dictionary() +identity_messages(pairings) = identity_messages(Float64, pairings) + @traitfn default_bp_maxiter(g::::(!IsDirected)) = is_tree(g) ? 1 : nothing @traitfn function default_bp_maxiter(g::::IsDirected) return default_bp_maxiter(undirected_graph(underlying_graph(g))) @@ -50,11 +68,6 @@ default_partitioned_vertices(ψ::AbstractITensorNetwork) = group(v -> v, vertice partitioned_tensornetwork(bpc::AbstractBeliefPropagationCache) = not_implemented() messages(bpc::AbstractBeliefPropagationCache) = not_implemented() -function default_message( - bpc::AbstractBeliefPropagationCache, edge::QuotientEdge; kwargs... - ) - return not_implemented() -end default_update_alg(bpc::AbstractBeliefPropagationCache) = not_implemented() default_message_update_alg(bpc::AbstractBeliefPropagationCache) = not_implemented() Base.copy(bpc::AbstractBeliefPropagationCache) = not_implemented() @@ -162,11 +175,6 @@ function PartitionedGraphs.quotientedge( return PartitionedGraphs.quotientedge(partitioned_tensornetwork(bpc), edge) end -function linkinds(bpc::AbstractBeliefPropagationCache, pe::QuotientEdge) - pitn = partitioned_tensornetwork(bpc) - return commoninds(subgraph(pitn, src(pe)), subgraph(pitn, dst(pe))) -end - NDTensors.scalartype(bpc::AbstractBeliefPropagationCache) = scalartype(tensornetwork(bpc)) """ @@ -187,12 +195,11 @@ function update_factor(bpc, vertex, factor) return bpc end -function message(bpc::AbstractBeliefPropagationCache, edge::QuotientEdge; kwargs...) - mts = messages(bpc) - return get(() -> default_message(bpc, edge; kwargs...), mts, edge) +function message(bpc::AbstractBeliefPropagationCache, edge::QuotientEdge) + return messages(bpc)[edge] end -function messages(bpc::AbstractBeliefPropagationCache, edges; kwargs...) - return map(edge -> message(bpc, edge; kwargs...), edges) +function messages(bpc::AbstractBeliefPropagationCache, edges) + return map(edge -> message(bpc, edge), edges) end function set_messages!(bpc::AbstractBeliefPropagationCache, quotientedges_messages) ms = messages(bpc) diff --git a/src/caches/beliefpropagationcache.jl b/src/caches/beliefpropagationcache.jl index 0bb9ff93..e1a7b143 100644 --- a/src/caches/beliefpropagationcache.jl +++ b/src/caches/beliefpropagationcache.jl @@ -1,4 +1,5 @@ using DataGraphs: DataGraphs, set_vertex_data! +using Dictionaries: Dictionary using Graphs: IsDirected using ITensors: dir using LinearAlgebra: diag, dot @@ -24,7 +25,7 @@ struct BeliefPropagationCache{V, PV, PTN <: AbstractPartitionedGraph{V, PV}, MTS end #Constructors... -function BeliefPropagationCache(ptn::PartitionedGraph; messages = default_messages(ptn)) +function BeliefPropagationCache(ptn::PartitionedGraph; messages = Dictionary()) return BeliefPropagationCache(ptn, messages) end @@ -51,10 +52,6 @@ end messages(bp_cache::BeliefPropagationCache) = bp_cache.messages -function default_message(bp_cache::BeliefPropagationCache, edge::QuotientEdge) - return default_message(datatype(bp_cache), linkinds(bp_cache, edge)) -end - function Base.copy(bp_cache::BeliefPropagationCache) return BeliefPropagationCache( copy(partitioned_tensornetwork(bp_cache)), copy(messages(bp_cache)) diff --git a/src/expect.jl b/src/expect.jl index 0a831b9e..30e0f100 100644 --- a/src/expect.jl +++ b/src/expect.jl @@ -1,5 +1,7 @@ using Dictionaries: Dictionary, set! +using Graphs: is_tree using ITensors: Op, contract, op, which_op +using NamedGraphs.PartitionedGraphs: PartitionedGraph, quotient_graph default_expect_alg() = "bp" @@ -20,7 +22,7 @@ function expect(ψIψ::AbstractFormNetwork, op::Op; kwargs...) end function expect( - alg::Algorithm, + alg::Algorithm"bp", ψ::AbstractITensorNetwork, ops; (cache!) = nothing, @@ -31,7 +33,17 @@ function expect( ) ψIψ = QuadraticFormNetwork(ψ) if isnothing(cache!) - cache! = Ref(cache(alg, ψIψ; cache_construction_kwargs...)) + pv = get( + cache_construction_kwargs, :partitioned_vertices, + default_partitioned_vertices(ψIψ) + ) + ptn = PartitionedGraph(ψIψ, pv) + messages = get(cache_construction_kwargs, :messages, nothing) + if isnothing(messages) + messages = + is_tree(quotient_graph(ptn)) ? Dictionary() : identity_messages(ψIψ, ptn) + end + cache! = Ref(BeliefPropagationCache(ptn; messages)) end if update_cache diff --git a/src/formnetworks/bilinearformnetwork.jl b/src/formnetworks/bilinearformnetwork.jl index 513a3d97..f24ecc2f 100644 --- a/src/formnetworks/bilinearformnetwork.jl +++ b/src/formnetworks/bilinearformnetwork.jl @@ -1,8 +1,11 @@ using Adapt: adapt using DataGraphs: DataGraphs, set_vertex_data! +using Dictionaries: Dictionary, set! using ITensors.NDTensors: datatype, denseblocks -using ITensors: ITensor, Op, delta, prime, sim +using ITensors: ITensor, Index, Op, commoninds, 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 @@ -105,3 +108,50 @@ 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 index 654d858f..7e0a6b50 100644 --- a/src/formnetworks/linearformnetwork.jl +++ b/src/formnetworks/linearformnetwork.jl @@ -1,6 +1,9 @@ using DataGraphs: DataGraphs, set_vertex_data! -using ITensors: ITensor, prime +using Dictionaries: Dictionary, set! +using ITensors: ITensor, Index, commoninds, dag, prime using NamedGraphs.GraphsExtensions: disjoint_union +using NamedGraphs.PartitionedGraphs: + PartitionedGraph, QuotientEdge, partitioned_vertices, quotientedges default_dual_link_index_map = prime @@ -57,3 +60,42 @@ 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 index eb4b2674..b7eb6b7c 100644 --- a/src/formnetworks/quadraticformnetwork.jl +++ b/src/formnetworks/quadraticformnetwork.jl @@ -1,4 +1,8 @@ using DataGraphs: DataGraphs, set_vertex_data!, underlying_graph, vertex_data +using Dictionaries: Dictionary +using Graphs: is_tree +using ITensors.NDTensors: @Algorithm_str, Algorithm +using NamedGraphs.PartitionedGraphs: PartitionedGraph, quotient_graph default_index_map = prime default_inv_index_map = noprime @@ -85,6 +89,34 @@ function QuadraticFormNetwork( return QuadraticFormNetwork(blf, dual_index_map, dual_inv_index_map) end +function identity_messages(fn::QuadraticFormNetwork, ptn::PartitionedGraph) + return identity_messages(bilinear_formnetwork(fn), ptn) +end + +# QFN is structurally ψ-vs-ψ, so identity messages are the canonical +# initialization on a loopy quotient graph. Asymmetric form networks +# (general LFN/BFN built from ϕ ≠ ψ) don't have this guarantee, so this +# specialization is QFN-only — `inner(ϕ, ψ; alg = "bp")` still falls +# back to the generic path and requires the caller to supply messages. +function logscalar( + alg::Algorithm"bp", + fn::QuadraticFormNetwork; + (cache!) = nothing, + update_cache = isnothing(cache!), + cache_update_kwargs = (;) + ) + if isnothing(cache!) + pv = default_partitioned_vertices(fn) + ptn = PartitionedGraph(fn, pv) + messages = is_tree(quotient_graph(ptn)) ? Dictionary() : identity_messages(fn, ptn) + cache! = Ref(BeliefPropagationCache(ptn; messages)) + end + if update_cache + cache![] = update(cache![]; cache_update_kwargs...) + end + return logscalar(cache![]) +end + function update(qf::QuadraticFormNetwork, original_state_vertex, ket_state::ITensor) state_inds = inds(ket_state) bra_state = replaceinds(dag(ket_state), state_inds, dual_index_map(qf).(state_inds)) diff --git a/src/inner.jl b/src/inner.jl index 3733de48..e6d7c2cf 100644 --- a/src/inner.jl +++ b/src/inner.jl @@ -1,4 +1,4 @@ -using ITensors: inner, scalar +using ITensors: inner, scalar, sim using LinearAlgebra: norm, norm_sqr default_contract_alg(tns::Tuple) = "bp" diff --git a/src/normalize.jl b/src/normalize.jl index 83d77cec..6f41dcf0 100644 --- a/src/normalize.jl +++ b/src/normalize.jl @@ -1,4 +1,7 @@ +using Dictionaries: Dictionary +using Graphs: is_tree using LinearAlgebra: normalize +using NamedGraphs.PartitionedGraphs: PartitionedGraph, quotient_graph function rescale(tn::AbstractITensorNetwork; alg = "exact", kwargs...) return rescale(Algorithm(alg), tn; kwargs...) @@ -10,6 +13,39 @@ function rescale(alg::Algorithm"exact", tn::AbstractITensorNetwork; kwargs...) return map(t -> c * t, tn) end +function rescale( + alg::Algorithm"bp", + tn::AbstractFormNetwork, + args...; + (cache!) = nothing, + cache_construction_kwargs = (;), + update_cache = isnothing(cache!), + cache_update_kwargs = (;), + kwargs... + ) + if isnothing(cache!) + pv = get( + cache_construction_kwargs, :partitioned_vertices, + default_partitioned_vertices(tn) + ) + ptn = PartitionedGraph(tn, pv) + messages = get(cache_construction_kwargs, :messages, nothing) + if isnothing(messages) + messages = + is_tree(quotient_graph(ptn)) ? Dictionary() : identity_messages(tn, ptn) + end + cache! = Ref(BeliefPropagationCache(ptn; messages)) + end + + if update_cache + cache![] = update(cache![]; cache_update_kwargs...) + end + + cache![] = rescale(cache![], args...; kwargs...) + + return tensornetwork(cache![]) +end + function rescale( alg::Algorithm, tn::AbstractITensorNetwork, @@ -61,20 +97,30 @@ function LinearAlgebra.normalize( end function LinearAlgebra.normalize( - alg::Algorithm, + alg::Algorithm"bp", tn::AbstractITensorNetwork; (cache!) = nothing, - cache_construction_function = tn -> - cache(alg, tn; default_cache_construction_kwargs(alg, tn)...), update_cache = isnothing(cache!), cache_update_kwargs = (;), cache_construction_kwargs = (;) ) norm_tn = inner_network(tn, tn) if isnothing(cache!) - cache! = Ref(cache(alg, norm_tn; cache_construction_kwargs...)) + pv = get( + cache_construction_kwargs, :partitioned_vertices, + default_partitioned_vertices(norm_tn) + ) + ptn = PartitionedGraph(norm_tn, pv) + messages = get(cache_construction_kwargs, :messages, nothing) + if isnothing(messages) + messages = if is_tree(quotient_graph(ptn)) + Dictionary() + else + identity_messages(norm_tn, ptn) + end + end + cache! = Ref(BeliefPropagationCache(ptn; messages)) 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]) norm_tn = rescale(alg, norm_tn; verts, cache!, update_cache, cache_update_kwargs) diff --git a/test/test_apply.jl b/test/test_apply.jl index 2c928b03..555b746b 100644 --- a/test/test_apply.jl +++ b/test/test_apply.jl @@ -1,9 +1,10 @@ using Compat: Compat using Graphs: vertices -using ITensorNetworks: - BeliefPropagationCache, apply, environment, norm_sqr_network, siteinds, update +using ITensorNetworks: BeliefPropagationCache, apply, environment, identity_messages, + norm_sqr_network, siteinds, update using ITensors: ITensors, ITensor, inner, op using NamedGraphs.NamedGraphGenerators: named_grid +using NamedGraphs.PartitionedGraphs: PartitionedGraph using SplitApplyCombine: group using StableRNGs: StableRNG using TensorOperations: TensorOperations @@ -20,12 +21,15 @@ include("utils.jl") ψψ = norm_sqr_network(ψ) # Simple Belief Propagation grouping (one bra+ket per partition) gives # a product environment around `[v1, v2]`, which is what `apply` requires. - bp_cache = BeliefPropagationCache(ψψ, group(v -> v[1], vertices(ψψ))) + ptn_SBP = PartitionedGraph(ψψ, group(v -> v[1], vertices(ψψ))) + bp_cache = BeliefPropagationCache(ptn_SBP; messages = 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. - bp_cache_col = BeliefPropagationCache(ψψ, group(v -> v[1][1], vertices(ψψ))) + ptn_col = PartitionedGraph(ψψ, group(v -> v[1][1], vertices(ψψ))) + bp_cache_col = + BeliefPropagationCache(ptn_col; messages = 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 index 29aafc68..19f327a7 100644 --- a/test/test_belief_propagation.jl +++ b/test/test_belief_propagation.jl @@ -1,16 +1,17 @@ using Compat: Compat -using Graphs: vertices -using ITensorNetworks: ITensorNetworks, BeliefPropagationCache, contract, - contraction_sequence, environment, inner_network, message, message_diff, - partitioned_tensornetwork, scalar, siteinds, tensornetwork, update, update_factor, - updated_message, ⊗ +using Dictionaries: Dictionary, set! +using Graphs: dst, src, vertices +using ITensorNetworks: ITensorNetworks, BeliefPropagationCache, QuadraticFormNetwork, + contract, contraction_sequence, environment, identity_messages, inner_network, message, + message_diff, partitioned_tensornetwork, scalar, siteinds, tensornetwork, update, + update_factor, updated_message, ⊗ include("utils.jl") using ITensors.NDTensors: array -using ITensors: ITensors, Algorithm, ITensor, combiner, commoninds, dag, inds, inner, op, - prime, random_itensor +using ITensors: ITensors, Algorithm, ITensor, Index, combiner, commoninds, dag, inds, inner, + op, prime, random_itensor using LinearAlgebra: eigvals, tr using NamedGraphs.NamedGraphGenerators: named_comb_tree, named_grid -using NamedGraphs.PartitionedGraphs: quotientedges +using NamedGraphs.PartitionedGraphs: PartitionedGraph, QuotientEdge, quotientedges using NamedGraphs: NamedEdge, NamedGraph using SplitApplyCombine: group using StableRNGs: StableRNG @@ -28,7 +29,25 @@ using Test: @test, @testset rng = StableRNG(1234) ψ = random_tensornetwork(rng, elt, s; link_space = χ) ψψ = ψ ⊗ prime(dag(ψ); sites = []) - bpc = BeliefPropagationCache(ψψ, group(v -> first(v), vertices(ψψ))) + pv = group(v -> first(v), vertices(ψψ)) + ptn = PartitionedGraph(ψψ, pv) + # Build identity initial messages by hand: layer 1 is the ket + # (original ψ) and layer 2 is the bra (`dag(prime(ψ; sites = []))`). + # Pair the bra/ket link inds ordinally on each quotient edge so the + # resulting `delta(b, dag(k))` has opposite QN flow. + pairings = Dictionary{QuotientEdge, Pair{Vector{Index}, Vector{Index}}}() + for pe in quotientedges(ptn) + v_src, v_dst = parent(src(pe)), parent(dst(pe)) + for (v_from, v_to, e) in ( + (v_src, v_dst, pe), + (v_dst, v_src, reverse(pe)), + ) + bras = collect(commoninds(ψψ[(v_from, 2)], ψψ[(v_to, 2)])) + kets = collect(commoninds(ψψ[(v_from, 1)], ψψ[(v_to, 1)])) + set!(pairings, e, bras => kets) + end + end + bpc = BeliefPropagationCache(ptn; messages = identity_messages(elt, pairings)) #Test updating the tensors in the cache vket, vbra = ((1, 1), 1), ((1, 1), 2) @@ -92,3 +111,18 @@ using Test: @test, @testset @test iszero(scalar(ψ; alg = "bp")) end end + +# Regression guard for `identity_messages`: a tiny QN-graded loopy +# network (4-cycle, S=1/2 sites) where the old single-leg `delta(i)` +# initialization collapsed messages to empty blocks and BP produced +# NaN. Two-leg `delta(b, k)` keeps the QN sectors aligned, so the +# QFN-routed `scalar` (which auto-uses `identity_messages` because QFN +# is structurally ψ-vs-ψ) must come back finite and nonzero. +@testset "QN-graded loopy BP — identity_messages regression" begin + g = named_grid((2, 2); periodic = true) + s = siteinds("S=1/2", g; conserve_qns = true) + ψ = productstate(v -> isodd(sum(v)) ? "↑" : "↓", s) + n2 = scalar(QuadraticFormNetwork(ψ); alg = "bp", cache_update_kwargs = (; maxiter = 25)) + @test isfinite(n2) + @test !iszero(n2) +end diff --git a/test/test_forms.jl b/test/test_forms.jl index 3d71fdae..40923d45 100644 --- a/test/test_forms.jl +++ b/test/test_forms.jl @@ -59,15 +59,14 @@ include("utils.jl") ∂qf_∂v = only(environment(qf, state_vertices(qf, [v]); alg = "exact")) @test (∂qf_∂v) * (qf[ket_vertex(qf, v)] * qf[bra_vertex(qf, v)]) ≈ contract(qf) - ∂qf_∂v_bp = environment(qf, state_vertices(qf, [v]); alg = "bp", update_cache = false) - ∂qf_∂v_bp = contract(∂qf_∂v_bp) - ∂qf_∂v_bp /= norm(∂qf_∂v_bp) - ∂qf_∂v /= norm(∂qf_∂v) - @test ∂qf_∂v_bp != ∂qf_∂v - + # `environment(::AbstractFormNetwork, …; alg = "bp")` builds messages + # via the form-network's `identity_messages` path and updates the + # cache; on a tree, BP is exact so the BP env matches the exact env + # after one sweep. ∂qf_∂v_bp = environment(qf, state_vertices(qf, [v]); alg = "bp", update_cache = true) ∂qf_∂v_bp = contract(∂qf_∂v_bp) ∂qf_∂v_bp /= norm(∂qf_∂v_bp) + ∂qf_∂v /= norm(∂qf_∂v) @test ∂qf_∂v_bp ≈ ∂qf_∂v #Test having non-uniform number of site indices per vertex diff --git a/test/test_normalize.jl b/test/test_normalize.jl index 9fa7e06b..aa4014a4 100644 --- a/test/test_normalize.jl +++ b/test/test_normalize.jl @@ -1,9 +1,11 @@ using Graphs: SimpleGraph, uniform_tree -using ITensorNetworks: BeliefPropagationCache, QuadraticFormNetwork, edge_scalars, messages, +using ITensorNetworks: BeliefPropagationCache, QuadraticFormNetwork, + default_partitioned_vertices, edge_scalars, identity_messages, messages, norm_sqr_network, rescale, scalartype, siteinds, vertex_scalars using ITensors: dag, inner, scalar using LinearAlgebra: normalize using NamedGraphs.NamedGraphGenerators: named_comb_tree, named_grid +using NamedGraphs.PartitionedGraphs: PartitionedGraph using NamedGraphs: NamedGraph using StableRNGs: StableRNG using TensorOperations: TensorOperations @@ -37,7 +39,13 @@ include("utils.jl") ψ = normalize(x; alg = "exact") @test scalar(norm_sqr_network(ψ); alg = "exact") ≈ 1.0 - ψIψ_bpc = Ref(BeliefPropagationCache(QuadraticFormNetwork(x))) + # Loopy BP needs explicit initial messages; the form-network's + # `identity_messages` gives a properly-flowing pair-delta starting + # point even for QN-graded indices. + qfn = QuadraticFormNetwork(x) + pv = default_partitioned_vertices(qfn) + ptn = PartitionedGraph(qfn, pv) + ψIψ_bpc = Ref(BeliefPropagationCache(ptn; messages = identity_messages(qfn, ptn))) ψ = normalize( x; alg = "bp", (cache!) = ψIψ_bpc, update_cache = true, cache_update_kwargs = (; maxiter = 20) diff --git a/test/utils.jl b/test/utils.jl index e6492676..ddb54e3f 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -5,24 +5,14 @@ using DataGraphs: underlying_graph, vertex_data using Dictionaries: Indices -using Graphs: AbstractGraph, dst, edges, src, vertices +using Graphs: AbstractGraph, add_edge!, dst, edges, src, vertices using ITensorNetworks: ITensorNetwork, IndsNetwork using ITensors.NDTensors: dim -using ITensors: ITensors, ITensor, Index, QN, dag, hasqns, inds, itensor, onehot +using ITensors: ITensors, ITensor, Index, QN, dag, hasqns, inds, itensor using NamedGraphs.GraphsExtensions: incident_edges using NamedGraphs: NamedGraph using Random: Random, AbstractRNG -# Pick a length-1 space matching `tn`'s Index style: a dense `1` if no -# vertex tensor has QN-graded indices, otherwise a single-block `[QN() => 1]`. -# Used by `_add_edge!` below to thread a trivial Link Index between two -# tensors without disturbing the QN structure (real `Index` constructors -# need a space, not just a dimension). -function _trivial_link_space(tn) - any(hasqns ∘ inds, (tn[v] for v in vertices(tn))) || return 1 - return [QN() => 1] -end - # --- random_tensornetwork ---------------------------------------------------- # Core: at each vertex of `graph`, place an `itensor(randn(rng, eltype, ...), inds_v)` @@ -83,27 +73,14 @@ end # --- productstate ------------------------------------------------------------- -# Thread a fresh trivial-QN dim-1 link Index between `src(edge)` and `dst(edge)` -# by outer-producting `onehot` basis vectors into each endpoint tensor. This is -# the productstate-specific analog of `ITensorNetworks.add_edge!` — the latter -# uses QR, which would push site-state QN flux into the link and leave BP's -# `default_message` (single-index `delta`) with no compatible block. `onehot` -# defaults to `Float64`, so we typecast it to `elt` to keep `productstate(elt, -# ...)` element-type-preserving. Reverse-map reconciliation on the second -# `setindex!` brings the graph edge along for free. -function _add_edge!(elt::Type, tn, edge) - iₑ = Index(_trivial_link_space(tn), "Link") - X = ITensors.convert_eltype(elt, onehot(iₑ => 1)) - tn[src(edge)] = tn[src(edge)] * X - tn[dst(edge)] = tn[dst(edge)] * dag(X) - return tn -end - # Build a product-state ITensorNetwork: start from a site-only TN (one tensor # per vertex with just the on-site state vector, looked up via # `ITensors.state(name, site_index)`), then add each edge from the original -# IndsNetwork via `_add_edge!`. `state` is anything indexable by vertex (dict, -# dictionary, array, ...); the `Function` method just materializes a Dict first. +# IndsNetwork via `Graphs.add_edge!`. The latter threads a fresh link `Index` +# via QR; QN flux on the link is fine here because BP messages now pair the +# bra and ket legs explicitly via `identity_messages`. `state` is anything +# indexable by vertex (dict, dictionary, array, ...); the `Function` method +# just materializes a Dict first. function productstate(elt::Type, state::Function, s::IndsNetwork) return productstate(elt, Dict(v => state(v) for v in vertices(s)), s) end @@ -115,7 +92,7 @@ function productstate(elt::Type, state, s::IndsNetwork) end tn = ITensorNetwork(ts) for e in edges(s) - _add_edge!(elt, tn, e) + add_edge!(tn, e) end return tn end From 2fc91cd3d55806d2cfc61bcc2af7885d9ec35038 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 13 May 2026 11:03:12 -0400 Subject: [PATCH 2/7] Replace cache(alg, tn) with initialize_cache(f, alg, tn) Move the per-purpose initialization logic out of the algorithm wrappers and into a single `initialize_cache(f, alg, tn; kwargs...)` dispatch. The `f` tag carries the calling context (`scalar`, `normalize`, `expect`, `rescale`, `environment`) so specializations can encode context-specific initialization without the wrappers having to know about form networks. A single private helper `_bp_cache_identity_messages(fn)` does the "identity messages on loopy quotient graph" construction; per-purpose methods for `scalar`/`normalize`/`expect`/`rescale` delegate to it, dispatched on `QuadraticFormNetwork` or `AbstractFormNetwork` as appropriate. Drops `cache(alg, tn)`, `default_cache_construction_kwargs`, and the type-driven `logscalar(::bp, ::QFN)` specialization in favor of this single dispatch point. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/ITensorNetworks.jl | 1 + src/caches/beliefpropagationcache.jl | 12 ----- src/contract.jl | 4 +- src/environment.jl | 6 ++- src/expect.jl | 17 +------ src/formnetworks/quadraticformnetwork.jl | 29 +---------- src/initialize_cache.jl | 62 ++++++++++++++++++++++++ src/normalize.jl | 58 ++-------------------- 8 files changed, 77 insertions(+), 112 deletions(-) create mode 100644 src/initialize_cache.jl diff --git a/src/ITensorNetworks.jl b/src/ITensorNetworks.jl index b048a941..64f8c07d 100644 --- a/src/ITensorNetworks.jl +++ b/src/ITensorNetworks.jl @@ -55,6 +55,7 @@ include("inner.jl") include("normalize.jl") include("expect.jl") include("environment.jl") +include("initialize_cache.jl") include("exports.jl") end diff --git a/src/caches/beliefpropagationcache.jl b/src/caches/beliefpropagationcache.jl index e1a7b143..6ee4eb8c 100644 --- a/src/caches/beliefpropagationcache.jl +++ b/src/caches/beliefpropagationcache.jl @@ -10,14 +10,6 @@ using NamedGraphs.PartitionedGraphs: AbstractPartitionedGraph, PartitionedGraph, using SimpleTraits: SimpleTraits, @traitfn, Not using SplitApplyCombine: group -function default_cache_construction_kwargs(alg::Algorithm"bp", ψ::AbstractITensorNetwork) - return (; partitioned_vertices = default_partitioned_vertices(ψ)) -end - -function default_cache_construction_kwargs(alg::Algorithm"bp", pg::PartitionedGraph) - return (;) -end - struct BeliefPropagationCache{V, PV, PTN <: AbstractPartitionedGraph{V, PV}, MTS} <: AbstractBeliefPropagationCache{V, PV} partitioned_tensornetwork::PTN @@ -42,10 +34,6 @@ function BeliefPropagationCache( return BeliefPropagationCache(tn, partitioned_vertices; kwargs...) end -function cache(alg::Algorithm"bp", tn; kwargs...) - return BeliefPropagationCache(tn; kwargs...) -end - function partitioned_tensornetwork(bp_cache::BeliefPropagationCache) return bp_cache.partitioned_tensornetwork end diff --git a/src/contract.jl b/src/contract.jl index 8c6f4f57..d59b9372 100644 --- a/src/contract.jl +++ b/src/contract.jl @@ -42,12 +42,12 @@ function logscalar( alg::Algorithm, tn::AbstractITensorNetwork; (cache!) = nothing, - cache_construction_kwargs = default_cache_construction_kwargs(alg, tn), + cache_construction_kwargs = (;), update_cache = isnothing(cache!), cache_update_kwargs = (;) ) if isnothing(cache!) - cache! = Ref(cache(alg, tn; cache_construction_kwargs...)) + cache! = Ref(initialize_cache(scalar, alg, tn; cache_construction_kwargs...)) end if update_cache diff --git a/src/environment.jl b/src/environment.jl index b7a0f6d4..b5a58d91 100644 --- a/src/environment.jl +++ b/src/environment.jl @@ -24,11 +24,13 @@ function environment( vertices::Vector; (cache!) = nothing, update_cache = isnothing(cache!), - cache_construction_kwargs = default_cache_construction_kwargs(alg, ptn), + cache_construction_kwargs = (;), cache_update_kwargs = (;) ) if isnothing(cache!) - cache! = Ref(cache(alg, ptn; cache_construction_kwargs...)) + cache! = Ref( + initialize_cache(environment, alg, ptn; cache_construction_kwargs...) + ) end if update_cache diff --git a/src/expect.jl b/src/expect.jl index 30e0f100..922e1964 100644 --- a/src/expect.jl +++ b/src/expect.jl @@ -1,7 +1,4 @@ -using Dictionaries: Dictionary, set! -using Graphs: is_tree using ITensors: Op, contract, op, which_op -using NamedGraphs.PartitionedGraphs: PartitionedGraph, quotient_graph default_expect_alg() = "bp" @@ -22,7 +19,7 @@ function expect(ψIψ::AbstractFormNetwork, op::Op; kwargs...) end function expect( - alg::Algorithm"bp", + alg::Algorithm, ψ::AbstractITensorNetwork, ops; (cache!) = nothing, @@ -33,17 +30,7 @@ function expect( ) ψIψ = QuadraticFormNetwork(ψ) if isnothing(cache!) - pv = get( - cache_construction_kwargs, :partitioned_vertices, - default_partitioned_vertices(ψIψ) - ) - ptn = PartitionedGraph(ψIψ, pv) - messages = get(cache_construction_kwargs, :messages, nothing) - if isnothing(messages) - messages = - is_tree(quotient_graph(ptn)) ? Dictionary() : identity_messages(ψIψ, ptn) - end - cache! = Ref(BeliefPropagationCache(ptn; messages)) + cache! = Ref(initialize_cache(expect, alg, ψIψ; cache_construction_kwargs...)) end if update_cache diff --git a/src/formnetworks/quadraticformnetwork.jl b/src/formnetworks/quadraticformnetwork.jl index b7eb6b7c..38097bda 100644 --- a/src/formnetworks/quadraticformnetwork.jl +++ b/src/formnetworks/quadraticformnetwork.jl @@ -1,8 +1,5 @@ using DataGraphs: DataGraphs, set_vertex_data!, underlying_graph, vertex_data -using Dictionaries: Dictionary -using Graphs: is_tree -using ITensors.NDTensors: @Algorithm_str, Algorithm -using NamedGraphs.PartitionedGraphs: PartitionedGraph, quotient_graph +using NamedGraphs.PartitionedGraphs: PartitionedGraph default_index_map = prime default_inv_index_map = noprime @@ -93,30 +90,6 @@ function identity_messages(fn::QuadraticFormNetwork, ptn::PartitionedGraph) return identity_messages(bilinear_formnetwork(fn), ptn) end -# QFN is structurally ψ-vs-ψ, so identity messages are the canonical -# initialization on a loopy quotient graph. Asymmetric form networks -# (general LFN/BFN built from ϕ ≠ ψ) don't have this guarantee, so this -# specialization is QFN-only — `inner(ϕ, ψ; alg = "bp")` still falls -# back to the generic path and requires the caller to supply messages. -function logscalar( - alg::Algorithm"bp", - fn::QuadraticFormNetwork; - (cache!) = nothing, - update_cache = isnothing(cache!), - cache_update_kwargs = (;) - ) - if isnothing(cache!) - pv = default_partitioned_vertices(fn) - ptn = PartitionedGraph(fn, pv) - messages = is_tree(quotient_graph(ptn)) ? Dictionary() : identity_messages(fn, ptn) - cache! = Ref(BeliefPropagationCache(ptn; messages)) - end - if update_cache - cache![] = update(cache![]; cache_update_kwargs...) - end - return logscalar(cache![]) -end - function update(qf::QuadraticFormNetwork, original_state_vertex, ket_state::ITensor) state_inds = inds(ket_state) bra_state = replaceinds(dag(ket_state), state_inds, dual_index_map(qf).(state_inds)) diff --git a/src/initialize_cache.jl b/src/initialize_cache.jl new file mode 100644 index 00000000..4fb7bda1 --- /dev/null +++ b/src/initialize_cache.jl @@ -0,0 +1,62 @@ +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...) + return BeliefPropagationCache(tn; kwargs...) +end + +function initialize_cache( + f, 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; + partitioned_vertices = default_partitioned_vertices(fn), + messages = nothing + ) + ptn = PartitionedGraph(fn, partitioned_vertices) + if isnothing(messages) + messages = is_tree(quotient_graph(ptn)) ? Dictionary() : identity_messages(fn, ptn) + 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/normalize.jl b/src/normalize.jl index 6f41dcf0..4fa9cd65 100644 --- a/src/normalize.jl +++ b/src/normalize.jl @@ -1,7 +1,4 @@ -using Dictionaries: Dictionary -using Graphs: is_tree using LinearAlgebra: normalize -using NamedGraphs.PartitionedGraphs: PartitionedGraph, quotient_graph function rescale(tn::AbstractITensorNetwork; alg = "exact", kwargs...) return rescale(Algorithm(alg), tn; kwargs...) @@ -13,51 +10,18 @@ function rescale(alg::Algorithm"exact", tn::AbstractITensorNetwork; kwargs...) return map(t -> c * t, tn) end -function rescale( - alg::Algorithm"bp", - tn::AbstractFormNetwork, - args...; - (cache!) = nothing, - cache_construction_kwargs = (;), - update_cache = isnothing(cache!), - cache_update_kwargs = (;), - kwargs... - ) - if isnothing(cache!) - pv = get( - cache_construction_kwargs, :partitioned_vertices, - default_partitioned_vertices(tn) - ) - ptn = PartitionedGraph(tn, pv) - messages = get(cache_construction_kwargs, :messages, nothing) - if isnothing(messages) - messages = - is_tree(quotient_graph(ptn)) ? Dictionary() : identity_messages(tn, ptn) - end - cache! = Ref(BeliefPropagationCache(ptn; messages)) - end - - if update_cache - cache![] = update(cache![]; cache_update_kwargs...) - end - - cache![] = rescale(cache![], args...; kwargs...) - - return tensornetwork(cache![]) -end - function rescale( alg::Algorithm, tn::AbstractITensorNetwork, args...; (cache!) = nothing, - cache_construction_kwargs = default_cache_construction_kwargs(alg, tn), + cache_construction_kwargs = (;), update_cache = isnothing(cache!), cache_update_kwargs = (;), kwargs... ) if isnothing(cache!) - cache! = Ref(cache(alg, tn; cache_construction_kwargs...)) + cache! = Ref(initialize_cache(rescale, alg, tn; cache_construction_kwargs...)) end if update_cache @@ -97,7 +61,7 @@ function LinearAlgebra.normalize( end function LinearAlgebra.normalize( - alg::Algorithm"bp", + alg::Algorithm, tn::AbstractITensorNetwork; (cache!) = nothing, update_cache = isnothing(cache!), @@ -106,20 +70,8 @@ function LinearAlgebra.normalize( ) norm_tn = inner_network(tn, tn) if isnothing(cache!) - pv = get( - cache_construction_kwargs, :partitioned_vertices, - default_partitioned_vertices(norm_tn) - ) - ptn = PartitionedGraph(norm_tn, pv) - messages = get(cache_construction_kwargs, :messages, nothing) - if isnothing(messages) - messages = if is_tree(quotient_graph(ptn)) - Dictionary() - else - identity_messages(norm_tn, ptn) - end - end - cache! = Ref(BeliefPropagationCache(ptn; messages)) + cache! = + Ref(initialize_cache(normalize, 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]) From 6d09a7186b36a32d823b8b7fb59686d609ac8897 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 13 May 2026 13:12:39 -0400 Subject: [PATCH 3/7] Drop LFN/BFN identity_messages; route auto-init through QFN only - 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) --- src/abstractitensornetwork.jl | 2 +- src/contract.jl | 2 +- src/environment.jl | 4 +- src/expect.jl | 2 +- src/formnetworks/bilinearformnetwork.jl | 52 +----------------------- src/formnetworks/linearformnetwork.jl | 44 +------------------- src/formnetworks/quadraticformnetwork.jl | 36 +++++++++++++++- src/initialize_cache.jl | 50 ++++------------------- src/inner.jl | 4 +- src/normalize.jl | 9 ++-- test/test_apply.jl | 43 ++++++++++++++++---- test/test_belief_propagation.jl | 6 ++- 12 files changed, 97 insertions(+), 157 deletions(-) diff --git a/src/abstractitensornetwork.jl b/src/abstractitensornetwork.jl index faddc4c9..4cee74b1 100644 --- 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 index d59b9372..3592700e 100644 --- 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 index b5a58d91..e9a92bd8 100644 --- 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 index 922e1964..2ad9e774 100644 --- 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 index f24ecc2f..0404acf2 100644 --- 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 index 7e0a6b50..9e9fb4c4 100644 --- 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 index 38097bda..e79f0aff 100644 --- 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 index 4fb7bda1..037ee496 100644 --- 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 index e6d7c2cf..02760a7b 100644 --- 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 index 4fa9cd65..d15b08cb 100644 --- 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 index 555b746b..b90e7e21 100644 --- 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 index 19f327a7..abf00efa 100644 --- a/test/test_belief_propagation.jl +++ b/test/test_belief_propagation.jl @@ -42,8 +42,12 @@ using Test: @test, @testset (v_src, v_dst, pe), (v_dst, v_src, reverse(pe)), ) - bras = collect(commoninds(ψψ[(v_from, 2)], ψψ[(v_to, 2)])) + # ψψ layer 2 is `dag(prime(ψ; sites = []))`, so each bra + # link Index is `dag(prime(k))` for the matching ket k. + # Build the pair explicitly to avoid relying on + # `commoninds` ordering across the two layers. kets = collect(commoninds(ψψ[(v_from, 1)], ψψ[(v_to, 1)])) + bras = dag.(prime.(kets)) set!(pairings, e, bras => kets) end end From c8d03b9633c93b33b6529f260e42fb5a332ca6a2 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 13 May 2026 13:33:10 -0400 Subject: [PATCH 4/7] Test simplification: use norm_sqr_network for BP and apply tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both `test/test_belief_propagation.jl` and `test/test_apply.jl` now build their `ψψ` through `norm_sqr_network(ψ)` (a QuadraticFormNetwork), so the auto-init in `initialize_cache(::bp, ::QFN)` handles the identity-messages construction directly — no hand-rolled pairings or LFN-specific helpers in the tests. The 2-site RDM construction in `test_belief_propagation.jl` is also rewritten for the QFN layout: ask `environment` for all three layers (bra/ket/operator) at the selected vertices, then contract with only the bra and ket tensors at those vertices. Dropping the per-site identity operator leaves the bra (primed) and ket (unprimed) site inds open, which is exactly the reduced density matrix. Co-Authored-By: Claude Opus 4.7 (1M context) --- test/test_apply.jl | 77 +++++++++++--------------- test/test_belief_propagation.jl | 97 +++++++++++++-------------------- 2 files changed, 71 insertions(+), 103 deletions(-) diff --git a/test/test_apply.jl b/test/test_apply.jl index b90e7e21..eb651ae7 100644 --- a/test/test_apply.jl +++ b/test/test_apply.jl @@ -1,43 +1,15 @@ using Compat: Compat -using Dictionaries: Dictionary, set! -using Graphs: dst, src, vertices -using ITensorNetworks: BeliefPropagationCache, apply, environment, identity_messages, - inner_network, siteinds, update -using ITensors: ITensors, ITensor, Index, commoninds, dag, inner, op, prime +using Graphs: vertices +using ITensorNetworks: BeliefPropagationCache, apply, environment, initialize_cache, + norm_sqr_network, siteinds, update +using ITensors: ITensors, Algorithm, ITensor, inner, op using NamedGraphs.NamedGraphGenerators: named_grid -using NamedGraphs.PartitionedGraphs: - PartitionedGraph, QuotientEdge, partitioned_vertices, quotientedges +using NamedGraphs.PartitionedGraphs: PartitionedGraph 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) @@ -46,23 +18,38 @@ end rng = StableRNG(1234) ψ = random_tensornetwork(rng, s; link_space = χ) v1, v2 = (2, 2), (1, 2) - ψψ = inner_network(ψ, ψ) - # Simple Belief Propagation grouping (one bra+ket per partition) gives - # a product environment around `[v1, v2]`, which is what `apply` requires. + ψψ = norm_sqr_network(ψ) + # Vertices of `[v1, v2]` across all layers of `ψψ` (bra/ket/operator), + # so the environment around them is just the incoming BP messages — + # the per-site operator tensors aren't pulled in as central tensors. + env_verts(vs) = [ + (v, suffix) for v in vs for suffix in ("bra", "ket", "operator") + ] + # Simple Belief Propagation grouping (one bra/ket/operator triple 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 = _lfn_identity_messages(ψψ, ptn_SBP)) - bp_cache = update(bp_cache; maxiter = 20) - envsSBP = environment(bp_cache, [(v1, "bra"), (v1, "ket"), (v2, "bra"), (v2, "ket")]) + bp_cache = update( + initialize_cache( + Algorithm("bp"), + ψψ; + partitioned_vertices = ptn_SBP.partitioned_vertices + ); + maxiter = 20 + ) + envsSBP = environment(bp_cache, env_verts((v1, v2))) # 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 = _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")] + bp_cache_col = update( + initialize_cache( + Algorithm("bp"), + ψψ; + partitioned_vertices = ptn_col.partitioned_vertices + ); + maxiter = 20 ) + envsGBP = environment(bp_cache_col, env_verts((v1, v2))) inner_alg = "exact" ngates = 5 truncerr = 0.0 diff --git a/test/test_belief_propagation.jl b/test/test_belief_propagation.jl index abf00efa..8182b722 100644 --- a/test/test_belief_propagation.jl +++ b/test/test_belief_propagation.jl @@ -1,19 +1,16 @@ using Compat: Compat -using Dictionaries: Dictionary, set! -using Graphs: dst, src, vertices +using Graphs: vertices using ITensorNetworks: ITensorNetworks, BeliefPropagationCache, QuadraticFormNetwork, - contract, contraction_sequence, environment, identity_messages, inner_network, message, - message_diff, partitioned_tensornetwork, scalar, siteinds, tensornetwork, update, - update_factor, updated_message, ⊗ + bra_vertex, contract, contraction_sequence, default_partitioned_vertices, environment, + identity_messages, ket_vertex, message, message_diff, norm_sqr_network, operator_vertex, + partitioned_tensornetwork, scalar, siteinds, tensornetwork, update, update_factor, + updated_message include("utils.jl") using ITensors.NDTensors: array -using ITensors: ITensors, Algorithm, ITensor, Index, combiner, commoninds, dag, inds, inner, - op, prime, random_itensor +using ITensors: ITensors, ITensor, combiner, dag, inds, inner, prime, random_itensor using LinearAlgebra: eigvals, tr -using NamedGraphs.NamedGraphGenerators: named_comb_tree, named_grid -using NamedGraphs.PartitionedGraphs: PartitionedGraph, QuotientEdge, quotientedges -using NamedGraphs: NamedEdge, NamedGraph -using SplitApplyCombine: group +using NamedGraphs.NamedGraphGenerators: named_grid +using NamedGraphs.PartitionedGraphs: PartitionedGraph, quotientedges using StableRNGs: StableRNG using TensorOperations: TensorOperations using Test: @test, @testset @@ -28,38 +25,19 @@ using Test: @test, @testset χ = 2 rng = StableRNG(1234) ψ = random_tensornetwork(rng, elt, s; link_space = χ) - ψψ = ψ ⊗ prime(dag(ψ); sites = []) - pv = group(v -> first(v), vertices(ψψ)) - ptn = PartitionedGraph(ψψ, pv) - # Build identity initial messages by hand: layer 1 is the ket - # (original ψ) and layer 2 is the bra (`dag(prime(ψ; sites = []))`). - # Pair the bra/ket link inds ordinally on each quotient edge so the - # resulting `delta(b, dag(k))` has opposite QN flow. - pairings = Dictionary{QuotientEdge, Pair{Vector{Index}, Vector{Index}}}() - for pe in quotientedges(ptn) - v_src, v_dst = parent(src(pe)), parent(dst(pe)) - for (v_from, v_to, e) in ( - (v_src, v_dst, pe), - (v_dst, v_src, reverse(pe)), - ) - # ψψ layer 2 is `dag(prime(ψ; sites = []))`, so each bra - # link Index is `dag(prime(k))` for the matching ket k. - # Build the pair explicitly to avoid relying on - # `commoninds` ordering across the two layers. - kets = collect(commoninds(ψψ[(v_from, 1)], ψψ[(v_to, 1)])) - bras = dag.(prime.(kets)) - set!(pairings, e, bras => kets) - end - end - bpc = BeliefPropagationCache(ptn; messages = identity_messages(elt, pairings)) + ψψ = norm_sqr_network(ψ) + ptn = PartitionedGraph(ψψ, default_partitioned_vertices(ψψ)) + bpc = BeliefPropagationCache(ptn; messages = identity_messages(ψψ, ptn)) - #Test updating the tensors in the cache - vket, vbra = ((1, 1), 1), ((1, 1), 2) + # Test updating the tensors in the cache. QFN bra has both site + # and link inds primed (relative to the ket), so the bra-side + # tensor is constructed as `dag(prime(new_A))` directly. + v = (1, 1) + vket = ket_vertex(ψψ, v) + vbra = bra_vertex(ψψ, v) A = bpc[vket] new_A = random_itensor(elt, inds(A)) - new_A_dag = ITensors.replaceind( - dag(prime(new_A)), only(s[first(vket)])', only(s[first(vket)]) - ) + new_A_dag = dag(prime(new_A)) bpc[vket] = new_A bpc[vbra] = new_A_dag @test bpc[vket] == new_A @@ -73,28 +51,31 @@ using Test: @test, @testset @test eltype(only(message(bpc, pe))) == elt end #Test updating the underlying tensornetwork in the cache - v = first(vertices(ψψ)) + v_any = first(vertices(ψψ)) rng = StableRNG(1234) - new_tensor = random_itensor(rng, inds(ψψ[v])) - bpc_updated = update_factor(bpc, v, new_tensor) + new_tensor = random_itensor(rng, inds(ψψ[v_any])) + bpc_updated = update_factor(bpc, v_any, new_tensor) ψψ_updated = tensornetwork(bpc_updated) - @test ψψ_updated[v] == new_tensor + @test ψψ_updated[v_any] == new_tensor - #Test forming a two-site RDM. Check it has the correct size, trace 1 and is PSD + # Two-site RDM at `vs`: ask for the environment around all three + # layers (bra/ket/operator) at `vs` so the BP env is just incoming + # messages, then contract with only the bra and ket tensors at + # `vs` — dropping the per-site identity operator at those + # vertices leaves the bra (primed) and ket (unprimed) site inds + # open, which is exactly the RDM. vs = [(2, 2), (2, 3)] - - # Prime the bra-ket shared site indices on the ket side at the - # selected vertices, so the contracted RDM has open primed/unprimed - # legs there. Mutates a copy of `ψψ` in place; no graph edits. - ψψsplit = copy(ψψ) - for v in vs - common = commoninds(ψψsplit[(v, 1)], ψψsplit[(v, 2)]) - ψψsplit[(v, 2)] = prime(ψψsplit[(v, 2)], common) - end - env_tensors = environment(bpc, [(v, 2) for v in vs]) - rdm = - contract(vcat(env_tensors, ITensor[ψψsplit[vp] for vp in [(v, 2) for v in vs]])) - + env_vs = vcat( + [bra_vertex(ψψ, v) for v in vs], + [ket_vertex(ψψ, v) for v in vs], + [operator_vertex(ψψ, v) for v in vs] + ) + env_tensors = environment(bpc, env_vs) + local_tensors = vcat( + ITensor[ψψ[bra_vertex(ψψ, v)] for v in vs], + ITensor[ψψ[ket_vertex(ψψ, v)] for v in vs] + ) + rdm = contract(vcat(env_tensors, local_tensors)) rdm = array( (rdm * combiner(inds(rdm; plev = 0)...)) * combiner(inds(rdm; plev = 1)...) ) From 4564fbac3c8d767c2f3cbadd007b95c1c69f49ee Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 13 May 2026 13:45:25 -0400 Subject: [PATCH 5/7] Simplify initialize_cache fallback and drop unused PartitionedGraph in test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Collapse the `AbstractITensorNetwork` and `AbstractPartitionedGraph` fallbacks of `initialize_cache(::Algorithm"bp", ...)` into a single untyped method — both just forward to `BeliefPropagationCache`. In `test_apply.jl`, pass the partition dict from `group` directly to `initialize_cache` via the `partitioned_vertices` kwarg instead of constructing a throwaway `PartitionedGraph` only to read its `partitioned_vertices` field back out. --- src/initialize_cache.jl | 9 ++------- test/test_apply.jl | 17 ++++------------- 2 files changed, 6 insertions(+), 20 deletions(-) diff --git a/src/initialize_cache.jl b/src/initialize_cache.jl index 037ee496..27ed989c 100644 --- a/src/initialize_cache.jl +++ b/src/initialize_cache.jl @@ -1,21 +1,16 @@ using Dictionaries: Dictionary using Graphs: is_tree using ITensors.NDTensors: @Algorithm_str, Algorithm -using NamedGraphs.PartitionedGraphs: - AbstractPartitionedGraph, PartitionedGraph, quotient_graph +using NamedGraphs.PartitionedGraphs: PartitionedGraph, quotient_graph # 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...) +function initialize_cache(alg::Algorithm"bp", tn; kwargs...) return BeliefPropagationCache(tn; kwargs...) end -function initialize_cache(alg::Algorithm"bp", ptn::AbstractPartitionedGraph; kwargs...) - return BeliefPropagationCache(ptn; kwargs...) -end - function initialize_cache( alg::Algorithm"bp", fn::QuadraticFormNetwork; diff --git a/test/test_apply.jl b/test/test_apply.jl index eb651ae7..21a97086 100644 --- a/test/test_apply.jl +++ b/test/test_apply.jl @@ -4,7 +4,6 @@ using ITensorNetworks: BeliefPropagationCache, apply, environment, initialize_ca norm_sqr_network, siteinds, update using ITensors: ITensors, Algorithm, ITensor, inner, op using NamedGraphs.NamedGraphGenerators: named_grid -using NamedGraphs.PartitionedGraphs: PartitionedGraph using SplitApplyCombine: group using StableRNGs: StableRNG using TensorOperations: TensorOperations @@ -28,25 +27,17 @@ include("utils.jl") # Simple Belief Propagation grouping (one bra/ket/operator triple per # partition) gives a product environment around `[v1, v2]`, which is # what `apply` requires. - ptn_SBP = PartitionedGraph(ψψ, group(v -> v[1], vertices(ψψ))) + pv_SBP = group(v -> v[1], vertices(ψψ)) bp_cache = update( - initialize_cache( - Algorithm("bp"), - ψψ; - partitioned_vertices = ptn_SBP.partitioned_vertices - ); + initialize_cache(Algorithm("bp"), ψψ; partitioned_vertices = pv_SBP); maxiter = 20 ) envsSBP = environment(bp_cache, env_verts((v1, v2))) # 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(ψψ))) + pv_col = group(v -> v[1][1], vertices(ψψ)) bp_cache_col = update( - initialize_cache( - Algorithm("bp"), - ψψ; - partitioned_vertices = ptn_col.partitioned_vertices - ); + initialize_cache(Algorithm("bp"), ψψ; partitioned_vertices = pv_col); maxiter = 20 ) envsGBP = environment(bp_cache_col, env_verts((v1, v2))) From cf61f87ee809a792abf7d30e5f94e5798af667c5 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 13 May 2026 13:58:52 -0400 Subject: [PATCH 6/7] Inline identity_messages(elt, pairings) into the QFN method The pairings-based `identity_messages(elt, pairings)` was a one-caller helper that just zipped the bra/ket index lists into `delta(elt, b, k)` tensors per quotient edge. With the QFN method now the only caller, collapse it inline: build each message Vector{ITensor} directly while walking the ket commoninds, and drop the standalone helper plus its docstring/Float64 fallback. Also prunes the now-unused `dictionary`, `commoninds`, `delta` imports from `abstractbeliefpropagationcache.jl`. --- src/caches/abstractbeliefpropagationcache.jl | 29 ++------------------ src/formnetworks/quadraticformnetwork.jl | 20 ++++++-------- 2 files changed, 11 insertions(+), 38 deletions(-) diff --git a/src/caches/abstractbeliefpropagationcache.jl b/src/caches/abstractbeliefpropagationcache.jl index e153cfef..29eb6542 100644 --- a/src/caches/abstractbeliefpropagationcache.jl +++ b/src/caches/abstractbeliefpropagationcache.jl @@ -1,8 +1,8 @@ using Adapt: Adapt, adapt, adapt_structure using DataGraphs: DataGraphs, underlying_graph, vertex_data -using Dictionaries: Dictionary, dictionary +using Dictionaries: Dictionary using Graphs: Graphs, IsDirected, dst, src -using ITensors: commoninds, delta, dir +using ITensors: dir using LinearAlgebra: diag, dot using NDTensors: NDTensors using NamedGraphs.GraphsExtensions: subgraph @@ -35,31 +35,6 @@ function message_diff(message_a::Vector{ITensor}, message_b::Vector{ITensor}) return 1 - f end -""" - identity_messages([elt::Type,] pairings) -> Dictionary{QuotientEdge, Vector{ITensor}} - -Build a message dictionary from a `pairings` mapping that lists, for each -quotient edge `e`, the `(bra_inds, ket_inds)` pair that defines the -identity-from-bra-to-ket message on that edge. Each entry is constructed -as `[delta(elt, bra_ind_i, ket_ind_i) for i in ...]`; the bra-side and -ket-side Indices are expected to already carry opposite directions so the -two-leg delta has well-defined QN flow (the bra side of a form network is -constructed via `dag` of the dual-mapped ket, which gives this naturally). - -`pairings` is anything supporting `pairs(...)` iteration yielding -`e => (bra_inds => ket_inds)` entries (e.g. a `Dictionary` or `Dict`). -`elt` defaults to `Float64`; form-network wrappers fill it in from -`scalartype(tensornetwork(fn))`. -""" -function identity_messages(elt::Type, pairings) - return dictionary( - e => ITensor[delta(elt, b, k) for (b, k) in zip(bras, kets)] - for (e, (bras, kets)) in pairs(pairings) - ) -end - -identity_messages(pairings) = identity_messages(Float64, pairings) - @traitfn default_bp_maxiter(g::::(!IsDirected)) = is_tree(g) ? 1 : nothing @traitfn function default_bp_maxiter(g::::IsDirected) return default_bp_maxiter(undirected_graph(underlying_graph(g))) diff --git a/src/formnetworks/quadraticformnetwork.jl b/src/formnetworks/quadraticformnetwork.jl index e79f0aff..db473f9a 100644 --- a/src/formnetworks/quadraticformnetwork.jl +++ b/src/formnetworks/quadraticformnetwork.jl @@ -1,6 +1,6 @@ using DataGraphs: DataGraphs, set_vertex_data!, underlying_graph, vertex_data using Dictionaries: Dictionary, set! -using ITensors: Index, commoninds, dag +using ITensors: ITensor, commoninds, dag, delta using NamedGraphs.PartitionedGraphs: PartitionedGraph, QuotientEdge, partitioned_vertices, quotientedges @@ -95,8 +95,9 @@ end # 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) - pairings = Dictionary{QuotientEdge, Pair{Vector{Index}, Vector{Index}}}() + messages = Dictionary{QuotientEdge, Vector{ITensor}}() tn = tensornetwork(fn) + elt = scalartype(tn) map_idx = dual_index_map(fn) pv = partitioned_vertices(ptn) ket_s = ket_vertex_suffix(fn) @@ -107,19 +108,16 @@ function identity_messages(fn::QuadraticFormNetwork, ptn::PartitionedGraph) (src_orig, dst_orig, pe), (dst_orig, src_orig, reverse(pe)), ) - kets = Index[] - bras = Index[] + ms = ITensor[] 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))) + for k in commoninds(tn[ket_vertex(fn, v_from)], tn[ket_vertex(fn, v_to)]) + push!(ms, delta(elt, dag(map_idx(k)), k)) + end end - set!(pairings, e, bras => kets) + set!(messages, e, ms) end end - return identity_messages(scalartype(tn), pairings) + return messages end function update(qf::QuadraticFormNetwork, original_state_vertex, ket_state::ITensor) From 77d46f95571be32a821a8c915291abc474e4ef6f Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Thu, 14 May 2026 10:07:54 -0400 Subject: [PATCH 7/7] Drop PartitionedGraph from identity_messages signature Make `identity_messages(fn::QuadraticFormNetwork)` take a `partitioned_vertices` kwarg (defaulting to `default_partitioned_vertices(fn)`) instead of a constructed `PartitionedGraph`. The function builds the PartitionedGraph internally; callers no longer have to pre-construct one just to hand it in. `initialize_cache` and the test sites pass their existing partition dict directly. --- src/formnetworks/quadraticformnetwork.jl | 11 +++++++---- src/initialize_cache.jl | 6 +++++- test/test_belief_propagation.jl | 7 +++++-- test/test_normalize.jl | 7 ++++++- 4 files changed, 23 insertions(+), 8 deletions(-) diff --git a/src/formnetworks/quadraticformnetwork.jl b/src/formnetworks/quadraticformnetwork.jl index db473f9a..adba9a3b 100644 --- a/src/formnetworks/quadraticformnetwork.jl +++ b/src/formnetworks/quadraticformnetwork.jl @@ -1,8 +1,7 @@ using DataGraphs: DataGraphs, set_vertex_data!, underlying_graph, vertex_data using Dictionaries: Dictionary, set! using ITensors: ITensor, commoninds, dag, delta -using NamedGraphs.PartitionedGraphs: - PartitionedGraph, QuotientEdge, partitioned_vertices, quotientedges +using NamedGraphs.PartitionedGraphs: PartitionedGraph, QuotientEdge, quotientedges default_index_map = prime default_inv_index_map = noprime @@ -94,12 +93,16 @@ end # 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) +function identity_messages( + fn::QuadraticFormNetwork; + partitioned_vertices = default_partitioned_vertices(fn) + ) + ptn = PartitionedGraph(fn, partitioned_vertices) messages = Dictionary{QuotientEdge, Vector{ITensor}}() tn = tensornetwork(fn) elt = scalartype(tn) map_idx = dual_index_map(fn) - pv = partitioned_vertices(ptn) + pv = partitioned_vertices 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))]))) diff --git a/src/initialize_cache.jl b/src/initialize_cache.jl index 27ed989c..ca4bf71e 100644 --- a/src/initialize_cache.jl +++ b/src/initialize_cache.jl @@ -19,7 +19,11 @@ function initialize_cache( ) ptn = PartitionedGraph(fn, partitioned_vertices) if isnothing(messages) - messages = is_tree(quotient_graph(ptn)) ? Dictionary() : identity_messages(fn, ptn) + messages = if is_tree(quotient_graph(ptn)) + Dictionary() + else + identity_messages(fn; partitioned_vertices) + end end return BeliefPropagationCache(ptn; messages) end diff --git a/test/test_belief_propagation.jl b/test/test_belief_propagation.jl index 8182b722..67d6fd1e 100644 --- a/test/test_belief_propagation.jl +++ b/test/test_belief_propagation.jl @@ -26,8 +26,11 @@ using Test: @test, @testset rng = StableRNG(1234) ψ = random_tensornetwork(rng, elt, s; link_space = χ) ψψ = norm_sqr_network(ψ) - ptn = PartitionedGraph(ψψ, default_partitioned_vertices(ψψ)) - bpc = BeliefPropagationCache(ptn; messages = identity_messages(ψψ, ptn)) + pv = default_partitioned_vertices(ψψ) + ptn = PartitionedGraph(ψψ, pv) + bpc = BeliefPropagationCache( + ptn; messages = identity_messages(ψψ; partitioned_vertices = pv) + ) # Test updating the tensors in the cache. QFN bra has both site # and link inds primed (relative to the ket), so the bra-side diff --git a/test/test_normalize.jl b/test/test_normalize.jl index aa4014a4..47a478ac 100644 --- a/test/test_normalize.jl +++ b/test/test_normalize.jl @@ -45,7 +45,12 @@ include("utils.jl") qfn = QuadraticFormNetwork(x) pv = default_partitioned_vertices(qfn) ptn = PartitionedGraph(qfn, pv) - ψIψ_bpc = Ref(BeliefPropagationCache(ptn; messages = identity_messages(qfn, ptn))) + ψIψ_bpc = Ref( + BeliefPropagationCache( + ptn; + messages = identity_messages(qfn; partitioned_vertices = pv) + ) + ) ψ = normalize( x; alg = "bp", (cache!) = ψIψ_bpc, update_cache = true, cache_update_kwargs = (; maxiter = 20)