Contract networks mixing operators and plain tensors

Project.toml

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

src/beliefpropagation/beliefpropagation.jl

@@ -244,10 +244,7 @@ end
 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)
+    return contract_network([messages; [factor]]; alg = algorithm.contraction_alg)
 end
 
 # Single-layer network: the message is a plain bond vector, normalized by its entrywise sum.
@@ -261,11 +258,11 @@ function message_update!(algorithm::SimpleMessageUpdate, cache, factors, edge)
     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.
+# `NormNetwork`: the message is a doubled (ket/bra) bond operator. Contracting a plain vertex factor
+# with the incoming messages leaves the surviving bond legs dangling, so assign the ket/bra pairing
+# the norm network gives this edge (the same convention as `similar_message_environment`). 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(

src/beliefpropagation/messagecache.jl

@@ -138,10 +138,7 @@ end
 
 function vertex_scalar(factors, messages, vertex; kwargs...)
     in_messages = incoming_edge_data(messages, [vertex])
-    # 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))]
+    tensors = [[factors[vertex]]; collect(in_messages)]
     return contract_network(tensors; kwargs...)[]
 end
 

src/contract_network.jl

@@ -26,12 +26,20 @@ function get_order(alg::Exact, tn)
     end
     # Contraction order may or may not have indices attached, canonicalize the format
     # by attaching indices.
-    subs = Dict(symnameddims(i) => symnameddims(i, Tuple(axes(tn[i]))) for i in keys(tn))
+    subs = Dict(symnameddims(i) => symnameddims(i, Tuple(axes(t))) for (i, t) in pairs(tn))
     return substitute(order, subs)
 end
+# Promote the operands to their common type before lowering to the lazy expression, so every lazy
+# operand shares one concrete type. Otherwise a network of mixed types (a plain tensor is a trivial
+# operator, so mixing operators and plain tensors is the common case) widens the symbolic `Mul`
+# container to a `UnionAll` it cannot construct. `promote_type`/`convert` keep an all-plain network
+# at the plain type (the promotion is a no-op), so its fast path is unchanged.
 function contract_network(alg::Exact, tn)
     order = get_order(alg, tn)
-    syms_to_ts = Dict(symnameddims(i, Tuple(axes(tn[i]))) => lazy(tn[i]) for i in keys(tn))
+    T = mapreduce(typeof, promote_type, tn)
+    syms_to_ts = Dict(
+        symnameddims(i, Tuple(axes(t))) => lazy(convert(T, t)) for (i, t) in pairs(tn)
+    )
     tn_expression = substitute(order, syms_to_ts)
     return materialize(tn_expression)
 end

test/test_contract_network.jl

@@ -1,5 +1,6 @@
 using Graphs: edges, vertices
-using ITensorBase: Greedy, Index
+using ITensorBase:
+    Greedy, Index, NamedTensorOperator, inputnames, operator, outputnames, state
 using ITensorNetworksNext: Exact, ITensorNetwork, LeftAssociative, contract_network,
     linkinds, siteinds, tensornetwork
 using NamedGraphs.GraphsExtensions: arranged_edges, incident_edges
@@ -51,4 +52,43 @@ using Test: @test, @testset
         @test z1 ≈ z4
         @test z1 ≈ z5
     end
+
+    @testset "Contract network with operators" begin
+        i, j, k = Index(2), Index(2), Index(2)
+        o = operator(randn(2, 2), (i,), (j,))     # output i, input j
+
+        # A network mixing an operator with plain tensors previously threw a `convert`
+        # `MethodError`; it now contracts, stays an operator, and matches the binary product.
+        t = randn(2, 2)[j, k]
+        r = contract_network([o, t])
+        @test r isa NamedTensorOperator
+        @test state(r) ≈ state(o * t)
+        @test outputnames(r) == outputnames(o * t)
+        @test inputnames(r) == inputnames(o * t)
+
+        # An all-operator network is likewise preserved.
+        o2 = operator(randn(2, 2), (j,), (k,))
+        r2 = contract_network([o, o2])
+        @test r2 isa NamedTensorOperator
+        @test state(r2) ≈ state(o * o2)
+
+        # A fully-contracted operator network reads out as a scalar via `[]`.
+        f = randn(2, 2)[i, j]
+        @test contract_network([o, f])[] ≈ (o * f)[]
+
+        # An all-plain network is unaffected: it is not promoted to an operator.
+        a = randn(2, 2)[i, j]
+        b = randn(2, 2)[j, k]
+        @test !(contract_network([a, b]) isa NamedTensorOperator)
+
+        # Pairing is order-independent: a branching network with a surviving output/input pair
+        # matches the binary product under any fold order (greedy contraction included).
+        ip, mp, m, x = Index(2), Index(2), Index(2), Index(2)
+        op = operator(randn(2, 2, 2, 2), (ip, mp), (i, m))
+        u = randn(2, 2)[mp, x]
+        w = randn(2, 2)[m, x]
+        rb = contract_network([op, u, w])
+        @test outputnames(rb) == outputnames((op * u) * w) == outputnames(op * (u * w))
+        @test inputnames(rb) == inputnames((op * u) * w) == inputnames(op * (u * w))
+    end
 end

test/test_normnetwork.jl

@@ -1,22 +1,17 @@
-using Base.Broadcast: materialize
 using DataGraphs: is_vertex_assigned
 using Dictionaries: isinsertable, issettable
 using Graphs: edges, vertices
 using ITensorBase:
     ITensor, Index, IndexName, LazyITensor, conj, inds, name, setname, uniquename
 using ITensorNetworksNext: BraView, ITensorNetwork, KetView, NormNetwork, braname,
-    bratensor, conj_bratensor, indmap, kettensor, normnetwork, tensornetwork
+    bratensor, conj_bratensor, contract_network, indmap, kettensor, normnetwork,
+    tensornetwork
 using LinearAlgebra: norm
 using NamedGraphs.GraphsExtensions: incident_edges
 using NamedGraphs.NamedGraphGenerators: named_grid, named_path_graph
 using NamedGraphs: NamedEdge
 using Test: @test, @test_throws, @testset
 
-# Contract a (possibly double-layer) network into a single tensor by multiplying all
-# of its vertex tensors together. For a `NormNetwork` the vertex data are lazy products
-# `ket * conj(bra)`, so the result is a lazy expression that we materialize.
-contract_network(tn) = materialize(prod(tn))
-
 # Build a random `ITensorNetwork` state on the graph `g` with site dimension `d` and
 # bond dimension `χ`.
 function random_state(::Type{T}, g; d = 2, χ = 2) where {T}