Dispatch belief-propagation message updates on NormNetwork

Project.toml

@@ -1,6 +1,6 @@
 name = "ITensorNetworksNext"
 uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c"
-version = "0.9.10"
+version = "0.9.11"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

src/beliefpropagation/beliefpropagation.jl

@@ -3,8 +3,7 @@ using .AlgorithmsInterfaceExtensions:
 using AlgorithmsInterface: AlgorithmsInterface as AI
 using DataGraphs: edge_data
 using Graphs: AbstractEdge, edges, edgetype, has_edge, vertices
-using ITensorBase:
-    AbstractITensor, NamedTensorOperator, inputnames, operator, outputnames, state
+using ITensorBase: AbstractITensor, operator, state
 using LinearAlgebra: norm, normalize, tr
 using NamedGraphs.GraphsExtensions:
     add_edges!, boundary_edges, forest_cover_edge_sequence, subgraph
@@ -240,37 +239,42 @@ end
     contraction_alg::ContractionAlg = Exact()
 end
 
-function message_update!(algorithm::SimpleMessageUpdate, cache, factors, edge)
+# Contract the incoming messages into the source factor to form the (unnormalized) new message on
+# `edge`.
+function updated_message(algorithm::SimpleMessageUpdate, cache, factors, edge)
     messages = collect(incoming_messages(cache, edge))
     factor = factors[src(edge)]
+    # TODO: `contract_network` can't currently contract a mix of operator and plain operands, so
+    # unwrap operator-valued messages with `state` first. Remove the `state.` once `contract_network`
+    # handles operator operands.
+    return contract_network([state.(messages); [factor]]; alg = algorithm.contraction_alg)
+end
 
-    # `contract_network` works on plain named arrays, so unwrap any operator messages to
-    # their underlying tensors before contracting (fermionic signs ride on the graded
-    # arrays, so nothing is lost).
-    message_tensors = state.(messages)
-    new_message = contract_network(
-        [message_tensors; [factor]]; alg = algorithm.contraction_alg
-    )
-
-    # `contract_network` drops the bra/ket operator structure, so restore it from the
-    # existing message. A doubled (ket/bra) message is then a bond operator: normalize by
-    # its trace, which is sign-correct for fermionic bonds (the entrywise `sum` can flip
-    # the odd-parity block's sign). A single-layer message stays a vector with no bra/ket
-    # pairing, so fall back to the entrywise sum there.
-    old_message = cache[edge]
-    if old_message isa NamedTensorOperator
-        new_message =
-            operator(new_message, outputnames(old_message), inputnames(old_message))
+# Single-layer network: the message is a plain bond vector, normalized by its entrywise sum.
+function message_update!(algorithm::SimpleMessageUpdate, cache, factors, edge)
+    new_message = updated_message(algorithm, cache, factors, edge)
+    if algorithm.normalize
+        message_norm = sum(new_message)
+        iszero(message_norm) || (new_message /= message_norm)
     end
+    cache[edge] = new_message
+    return cache
+end
 
+# `NormNetwork`: the message is a doubled (ket/bra) bond operator. `contract_network` drops the
+# operator structure, so re-wrap the result with the ket/bra names the norm network assigns to this
+# edge (the same convention as `similar_message_environment`) rather than reconstructing them from
+# the old message. Normalize by the trace, which is sign-correct on fermionic bonds where the
+# entrywise `sum` can flip the odd-parity block's sign.
+function message_update!(algorithm::SimpleMessageUpdate, cache, factors::NormNetwork, edge)
+    new_tensor = updated_message(algorithm, cache, factors, edge)
+    new_message = operator(
+        new_tensor, linknames(KetView(factors), edge), linknames(BraView(factors), edge)
+    )
     if algorithm.normalize
-        message_norm =
-            new_message isa NamedTensorOperator ? tr(new_message) : sum(new_message)
-        if !iszero(message_norm)
-            new_message /= message_norm
-        end
+        message_norm = tr(new_message)
+        iszero(message_norm) || (new_message /= message_norm)
     end
-
     cache[edge] = new_message
     return cache
 end

src/beliefpropagation/messagecache.jl

@@ -3,7 +3,7 @@ using DataGraphs: DataGraphs, AbstractDataGraph, AbstractEdgeDataGraph, edge_dat
     vertex_data_type
 using Dictionaries: Dictionary, delete!, getindices, set!
 using Graphs: AbstractGraph, connected_components, is_directed, is_tree
-using ITensorBase: unnamed
+using ITensorBase: state, unnamed
 using NamedGraphs.GraphsExtensions: IsDirected, boundary_edges, default_root_vertex,
     directed_graph, forest_cover, in_incident_edges, post_order_dfs_edges, undirected_graph,
     vertextype
@@ -138,7 +138,10 @@ end
 
 function vertex_scalar(factors, messages, vertex; kwargs...)
     in_messages = incoming_edge_data(messages, [vertex])
-    tensors = vcat([factors[vertex]], collect(in_messages))
+    # TODO: `contract_network` can't currently contract a mix of operator and plain operands, so
+    # unwrap operator-valued messages with `state` first. Remove the `state.` once `contract_network`
+    # handles operator operands.
+    tensors = [[factors[vertex]]; state.(collect(in_messages))]
     return contract_network(tensors; kwargs...)[]
 end
 
@@ -150,10 +153,8 @@ function vertex_scalars(factors, messages, vertices)
     return map(v -> vertex_scalar(factors, messages, v), vertices)
 end
 
-function edge_scalar(cache, edge; kwargs...)
-    m1 = cache[edge]
-    m2 = cache[reverse(edge)]
-    return contract_network([m1, m2]; kwargs...)[]
+function edge_scalar(cache, edge)
+    return (cache[edge] * cache[reverse(edge)])[]
 end
 
 edge_scalars(cache) = edge_scalars(cache, keys(cache))

test/test_beliefpropagation.jl

@@ -1,16 +1,19 @@
 import AlgorithmsInterface as AI
 using DataGraphs: DataGraphs, DataGraph, edge_data, edge_data_type
 using Dictionaries: Dictionary, dictionary, set!
+using GradedArrays: U1, gradedrange
 using Graphs: AbstractGraph, dst, edges, has_edge, src, vertices
-using ITensorBase: ITensor, Index, inds, name, noprime, prime
-using ITensorNetworksNext: ITensorNetworksNext, ITensorNetwork, MessageCache,
-    StopWhenConverged, bethe_free_energy, edge_scalar, incoming_messages, linkinds,
-    messagecache, region_scalar, subgraph, tensornetwork, vertex_scalar, vertex_scalars
+using ITensorBase: ITensor, Index, inds, name, noprime, outputnames, prime
+using ITensorNetworksNext: ITensorNetworksNext, ITensorNetwork, MessageCache, NormNetwork,
+    StopWhenConverged, beliefpropagation, bethe_free_energy, edge_scalar, incoming_messages,
+    insertlink!, linkinds, message_environment, messagecache, region_scalar, subgraph,
+    tensornetwork, vertex_scalar, vertex_scalars
 using LinearAlgebra: LinearAlgebra
 using NamedGraphs.GraphsExtensions: all_edges, arranged_edges, incident_edges, vertextype
 using NamedGraphs.NamedGraphGenerators: named_comb_tree, named_grid, named_path_graph
 using NamedGraphs: NamedEdge
 using StableRNGs: StableRNG
+using TensorKitSectors: FermionParity
 using Test: @test, @testset
 
 function spin_ice_tensornetwork(g)
@@ -171,7 +174,7 @@ end
                 edge => ones(T, Tuple(linkinds(tn, edge))) for edge in all_edges(g)
             )
 
-            cache = ITensorNetworksNext.beliefpropagation(
+            cache = beliefpropagation(
                 tn, messages; stopping_criterion = (; maxiter = 1)
             )
             z_bp = exp(bethe_free_energy(tn, cache))
@@ -194,7 +197,7 @@ end
                 edge => ones(T, Tuple(linkinds(tn, edge))) for edge in all_edges(g)
             )
 
-            cache = ITensorNetworksNext.beliefpropagation(
+            cache = beliefpropagation(
                 tn, messages; stopping_criterion = (; maxiter = 1)
             )
             z_bp = exp(bethe_free_energy(tn, cache))
@@ -213,7 +216,7 @@ end
                             for edge in all_edges(g)
                     )
 
-                    cache = ITensorNetworksNext.beliefpropagation(
+                    cache = beliefpropagation(
                         tn, messages;
                         stopping_criterion = (; maxiter = 10, tol = 1.0e-10)
                     )
@@ -225,4 +228,40 @@ end
             end
         end
     end
+
+    @testset "NormNetwork (operator-valued messages)" begin
+        site_ranges = (
+            "U1" => gradedrange([U1(0) => 1, U1(1) => 1]),
+            "FermionParity" =>
+                gradedrange([FermionParity(0) => 1, FermionParity(1) => 1]),
+        )
+        @testset "$label, T=$T" for (label, site_range) in site_ranges,
+                T in (Float64, ComplexF64)
+
+            rng = StableRNG(123)
+            g = named_path_graph(4)
+            site_axes = Dict(v => Index(site_range) for v in vertices(g))
+            network = tensornetwork(vertices(g)) do v
+                return randn(rng, T, (site_axes[v],))
+            end
+            for edge in edges(g)
+                insertlink!(network, edge)
+            end
+            nn = NormNetwork(network)
+
+            cache = beliefpropagation(
+                nn, message_environment(one, nn);
+                stopping_criterion = (; maxiter = 20, tol = 1.0e-10)
+            )
+
+            # Messages stay operator-valued end to end (a plain message has no output names).
+            @test all(msg -> !isempty(outputnames(msg)), edge_data(cache))
+
+            # Belief propagation is exact on a tree, including on the fermionic norm network.
+            ket = prod(network)
+            z_exact = (ket * conj(ket))[]
+            z_bp = exp(bethe_free_energy(nn, cache))
+            @test z_bp ≈ z_exact rtol = eps(real(T))^(1 / 3)
+        end
+    end
 end