Backport the norm-network belief-propagation contraction fix to v0.10

Project.toml

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

src/beliefpropagation/beliefpropagation.jl

@@ -237,12 +237,25 @@ end
     contraction_alg::ContractionAlg = Exact()
 end
 
+# The tensors making up the factor at `vertex`, as separate operands for `contract_network`. A
+# `NormNetwork`'s factor is a lazy `ket * conj(bra)` product, and the contraction order sees each
+# operand as one node carrying only its outer axes — which hides the physical index the two layers
+# share, forcing the doubled vertex to be formed before any message is absorbed (χ^(2 * degree)
+# rather than the χ^(degree + 1) an interleaved order reaches).
+factor_tensors(factors, vertex) = [factors[vertex]]
+function factor_tensors(factors::NormNetwork, vertex)
+    return [kettensor(factors, vertex), bratensor(factors, vertex)]
+end
+
 # 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)]
-    return contract_network([messages; [factor]]; alg = algorithm.contraction_alg)
+    # TODO: Remove `factor_tensors` once `contract_network` handles lazy tensors in
+    # contraction sequences properly.
+    return contract_network(
+        [messages; factor_tensors(factors, src(edge))]; alg = algorithm.contraction_alg
+    )
 end
 
 # Single-layer network: the message is a plain bond vector, normalized by its entrywise sum.

src/beliefpropagation/messagecache.jl

@@ -132,7 +132,9 @@ end
 
 function vertex_scalar(factors, messages, vertex; kwargs...)
     in_messages = incoming_edge_data(messages, [vertex])
-    tensors = [[factors[vertex]]; collect(in_messages)]
+    # TODO: Remove `factor_tensors` once `contract_network` handles lazy tensors in
+    # contraction sequences properly.
+    tensors = [factor_tensors(factors, vertex); collect(in_messages)]
     return contract_network(tensors; kwargs...)[]
 end
 

test/test_beliefpropagation.jl

@@ -1,14 +1,16 @@
 import AlgorithmsInterface as AI
+using Base.Broadcast: materialize
 using DataGraphs: DataGraphs, DataGraph, edge_data, edge_data_type
 using Dictionaries: Dictionary, dictionary, set!
 using GradedArrays: U1, gradedrange
 using Graphs: AbstractGraph, add_vertex!, dst, edges, has_edge, has_vertex, nv, rem_edge!,
     src, vertices
-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 ITensorBase: Greedy, ITensor, Index, inds, name, noprime, outputnames, prime
+using ITensorNetworksNext: ITensorNetworksNext, Exact, ITensorNetwork, MessageCache,
+    NormNetwork, SimpleMessageUpdate, StopWhenConverged, beliefpropagation,
+    bethe_free_energy, contract_network, contraction_order, edge_scalar, factor_tensors,
+    incoming_messages, insertlink!, linkinds, message_environment, messagecache,
+    region_scalar, subgraph, tensornetwork, updated_message, vertex_scalar, vertex_scalars
 using LinearAlgebra: LinearAlgebra
 using NamedGraphs: NamedEdge, all_edges, incident_edges, named_comb_tree, named_grid,
     named_path_graph, vertextype
@@ -38,6 +40,15 @@ function spin_ice_tensornetwork(g)
     return ITensorNetwork(ts)
 end
 
+# Records how many operands each `contract_network` call is given, then orders them greedily.
+struct RecordOperands
+    counts::Vector{Int}
+end
+function ITensorNetworksNext.contraction_order(alg::RecordOperands, tn)
+    push!(alg.counts, length(tn))
+    return contraction_order(tn; alg = Greedy())
+end
+
 @testset "Belief propagation" begin
     @testset "`MessageCache`" begin
         @testset "Basics" begin
@@ -277,4 +288,44 @@ end
             @test z_bp ≈ z_exact rtol = eps(real(T))^(1 / 3)
         end
     end
+
+    @testset "Doubled-vertex contraction operands" begin
+        site_ranges = (
+            "plain" => 2,
+            "U1" => gradedrange([U1(0) => 1, U1(1) => 1]),
+        )
+        @testset "$label" for (label, site_range) in site_ranges
+            rng = StableRNG(1234)
+            g = named_grid((3, 3))
+            network = tensornetwork(vertices(g)) do v
+                return randn(rng, (Index(site_range),))
+            end
+            for edge in edges(g)
+                insertlink!(network, edge)
+            end
+            nn = NormNetwork(network)
+            v = (2, 2)
+
+            # A doubled vertex splits into its two layers, and the split is faithful.
+            @test length(factor_tensors(nn, v)) == 2
+            @test prod(factor_tensors(nn, v)) ≈ materialize(nn[v])
+            # A single-layer network's factor is a single operand.
+            @test factor_tensors(network, v) == [network[v]]
+
+            # The message update passes the layers to `contract_network` as separate operands, so
+            # the contraction order can interleave the incoming messages between them, and the
+            # result matches contracting the doubled vertex as one operand.
+            counts = Int[]
+            algorithm = SimpleMessageUpdate(;
+                contraction_alg = Exact(; order_alg = RecordOperands(counts))
+            )
+            cache = message_environment(one, nn)
+            edge = NamedEdge(v => (2, 3))
+            messages = collect(incoming_messages(cache, edge))
+            message = updated_message(algorithm, cache, nn, edge)
+            # `v` has degree 4, so 3 incoming messages plus the ket and bra layers.
+            @test only(counts) == 5
+            @test message ≈ contract_network([messages; [nn[v]]])
+        end
+    end
 end