Replace contract_network with prod_tensors and an explicit ContractionTree

Project.toml

@@ -1,12 +1,13 @@
 name = "ITensorNetworksNext"
 uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c"
-version = "0.10.10"
+version = "0.11.0"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
 projects = ["benchmark", "dev", "docs", "examples", "test"]
 
 [deps]
+AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 AlgorithmsInterface = "d1e3940c-cd12-4505-8585-b0a4b322527d"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
@@ -30,13 +31,14 @@ OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 ITensorNetworksNextOMEinsumContractionOrdersExt = "OMEinsumContractionOrders"
 
 [compat]
+AbstractTrees = "0.4.5"
 Adapt = "4.3"
 AlgorithmsInterface = "0.1"
 Combinatorics = "1"
 DataGraphs = "0.6"
 Dictionaries = "0.4.5"
 Graphs = "1.13.1"
-ITensorBase = "0.15.5"
+ITensorBase = "0.16"
 LinearAlgebra = "1.10"
 MacroTools = "0.5.16"
 MatrixAlgebraKit = "0.6"

docs/Project.toml

@@ -10,5 +10,5 @@ path = ".."
 [compat]
 Documenter = "1"
 ITensorFormatter = "0.2.27"
-ITensorNetworksNext = "0.10"
+ITensorNetworksNext = "0.11"
 Literate = "2"

examples/Project.toml

@@ -5,4 +5,4 @@ ITensorNetworksNext = "302f2e75-49f0-4526-aef7-d8ba550cb06c"
 path = ".."
 
 [compat]
-ITensorNetworksNext = "0.10"
+ITensorNetworksNext = "0.11"

ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl

@@ -1,10 +1,37 @@
 module ITensorNetworksNextOMEinsumContractionOrdersExt
 
-using ITensorNetworksNext: ITensorNetworksNext
-using OMEinsumContractionOrders: CodeOptimizer
+using ITensorBase: inds, name
+using ITensorNetworksNext:
+    ITensorNetworksNext, ContractionTree, contraction_tree, prod_tensors
+using OMEinsumContractionOrders:
+    OMEinsumContractionOrders, CodeOptimizer, EinCode, NestedEinsum, optimize_code
 
-function ITensorNetworksNext.contraction_order(alg::CodeOptimizer, tn)
-    return ITensorNetworksNext._contraction_order(alg, tn)
+# Rebuild a `ContractionTree` from an optimized `NestedEinsum`, mapping each leaf's
+# `tensorindex` back to the key it came from.
+function nested_einsum_to_tree(ks, code::NestedEinsum)
+    # A leaf holds the 1-based index of its input tensor; internal nodes hold `-1`.
+    code.tensorindex != -1 && return ContractionTree(ks[code.tensorindex])
+    return reduce(ContractionTree, map(Base.Fix1(nested_einsum_to_tree, ks), code.args))
+end
+
+# Find a contraction order with any OMEinsumContractionOrders optimizer (`GreedyMethod`,
+# `TreeSA`, `KaHyParBipartite`, ...) by forwarding to `optimize_code`.
+function ITensorNetworksNext.contraction_tree(alg::CodeOptimizer, tensors)
+    ks = collect(keys(tensors))
+    ixs = [map(name, inds(tensors[k])) for k in ks]
+    all_inds = reduce(vcat, ixs)
+    labels = unique(all_inds)
+    size_dict = Dict(name(i) => length(i) for k in ks for i in inds(tensors[k]))
+    # Open indices (appearing on a single tensor) are the output of the network.
+    iy = filter(i -> count(==(i), all_inds) == 1, labels)
+    code = optimize_code(EinCode(ixs, iy), size_dict, alg)
+    return nested_einsum_to_tree(ks, code)
+end
+
+# `CodeOptimizer` cannot subtype `ContractionTreeAlgorithm`, so the order-taking form of
+# `prod_tensors` is extended for it here.
+function ITensorNetworksNext.prod_tensors(tensors, alg::CodeOptimizer)
+    return prod_tensors(tensors, contraction_tree(alg, tensors))
 end
 
 end

src/ITensorNetworksNext.jl

@@ -16,7 +16,8 @@ include("tensornetwork.jl")
 include("normnetwork.jl")
 include("normnetworkview.jl")
 include("ITensorNetworkGenerators/ITensorNetworkGenerators.jl")
-include("contract_network.jl")
+include("contraction_tree.jl")
+include("prod_tensors.jl")
 
 include("beliefpropagation/messagecache.jl")
 include("beliefpropagation/beliefpropagation.jl")

src/beliefpropagation/beliefpropagation.jl

@@ -232,16 +232,16 @@ function message_update!(cache, factors, edge; alg = nothing, kwargs...)
     )
 end
 
-@kwdef struct SimpleMessageUpdate{ContractionAlg} <: MessageUpdateAlgorithm
+@kwdef struct SimpleMessageUpdate{OrderAlg} <: MessageUpdateAlgorithm
     normalize::Bool = true
-    contraction_alg::ContractionAlg = Exact()
+    order_alg::OrderAlg = Greedy()
 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).
+# The tensors making up the factor at `vertex`, as separate operands to contract. A
+# `NormNetwork`'s factor is the product of its ket and bra layers. Handing that product over as one
+# operand would hide 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)]
@@ -251,11 +251,8 @@ end
 # `edge`.
 function updated_message(algorithm::SimpleMessageUpdate, cache, factors, edge)
     messages = collect(incoming_messages(cache, edge))
-    # 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
-    )
+    tensors = [messages; factor_tensors(factors, src(edge))]
+    return prod_tensors(tensors, algorithm.order_alg)
 end
 
 # Single-layer network: the message is a plain bond vector, normalized by its entrywise sum.

src/beliefpropagation/messagecache.jl

@@ -132,10 +132,8 @@ end
 
 function vertex_scalar(factors, messages, vertex; kwargs...)
     in_messages = incoming_edge_data(messages, [vertex])
-    # 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...)[]
+    return prod_tensors(tensors, contraction_tree(tensors; kwargs...))[]
 end
 
 vertex_scalars(factors, messages) = vertex_scalars(factors, messages, keys(factors))

src/contraction_tree.jl

@@ -0,0 +1,59 @@
+using AbstractTrees: AbstractTrees
+
+"""
+    Branch(left, right)
+
+The children of a branch node, as held by a [`ContractionTree`](@ref).
+
+Giving the children a type of their own is what keeps a branch distinguishable from a leaf: a
+vertex label can be a 2-tuple, so a bare tuple of children would be ambiguous, but no label is
+a `Branch` unless it is deliberately made one.
+"""
+struct Branch{T}
+    children::NTuple{2, T}
+end
+Branch(left, right) = Branch((left, right))
+
+"""
+    ContractionTree{V}
+
+A binary tree representing a contraction order over leaf labels of type `V`. A node is either a
+leaf holding one label or a [`Branch`](@ref) holding its two child trees.
+"""
+struct ContractionTree{V}
+    node::Union{V, Branch{ContractionTree{V}}}
+end
+
+function ContractionTree(children::Branch{ContractionTree{V}}) where {V}
+    return ContractionTree{V}(children)
+end
+function ContractionTree(left::ContractionTree{V}, right::ContractionTree{V}) where {V}
+    return ContractionTree{V}(Branch(left, right))
+end
+
+isleaf(t::ContractionTree) = !(t.node isa Branch)
+
+AbstractTrees.nodevalue(t::ContractionTree) = t.node
+AbstractTrees.children(t::ContractionTree) = isleaf(t) ? () : t.node.children
+
+Base.getindex(t::ContractionTree, i::Int) = t.node.children[i]
+
+Base.:(==)(a::ContractionTree, b::ContractionTree) = a.node == b.node
+Base.hash(t::ContractionTree, h::UInt) = hash(t.node, hash(:ContractionTree, h))
+
+function Base.show(io::IO, t::ContractionTree)
+    if isleaf(t)
+        show(io, AbstractTrees.nodevalue(t))
+    else
+        print(io, "(")
+        show(io, t[1])
+        print(io, ", ")
+        show(io, t[2])
+        print(io, ")")
+    end
+    return nothing
+end
+function AbstractTrees.printnode(io::IO, t::ContractionTree)
+    isleaf(t) && show(io, AbstractTrees.nodevalue(t))
+    return nothing
+end

src/normnetwork.jl

@@ -1,5 +1,5 @@
-using Dictionaries: Dictionary
-using ITensorBase: LazyNamedTensor, lazy, rename, setname, similar_operator, uniquename
+using Dictionaries: Dictionary, dictionary
+using ITensorBase: rename, setname, similar_operator, uniquename
 using ITensorNetworksNext
 
 """
@@ -26,7 +26,11 @@ struct NormNetwork{T, V, I} <: AbstractITensorNetwork{T, V}
     end
 end
 
-Base.eltype(::Type{<:NormNetwork{T, V, I}}) where {T, V, I} = LazyNamedTensor{I, T}
+function Base.eltype(::Type{<:NormNetwork})
+    return error(
+        "`eltype` of a `NormNetwork` is not defined, since the double-layer tensor at a vertex has no representation of its own. Use `kettensor` and `bratensor` to reach the individual layers."
+    )
+end
 
 function NormNetwork(tn::ITensorNetwork)
     return NormNetwork(tn, map(uniquename, keys(tn.dimname_vertices)))
@@ -46,10 +50,9 @@ NamedGraphs.encoded_graph(nn::NormNetwork) = encoded_graph(nn.ket)
 # ==================================== DataGraphs.jl ===================================== #
 
 function DataGraphs.get_vertex_data(nn::NormNetwork, vertex)
-    A = kettensor(nn, vertex)
-    B = conj_bratensor(nn, vertex)
-    # TODO: implement and use a lazy `conj` via `LazyNamedDimsArrays` here?
-    return lazy(A) * lazy(conj(B))
+    return error(
+        "Indexing a `NormNetwork` is not defined, since the double-layer tensor at a vertex has no representation of its own. Use `kettensor` and `bratensor` to reach the individual layers."
+    )
 end
 
 function DataGraphs.is_vertex_assigned(nn::NormNetwork, vertex)
@@ -80,6 +83,24 @@ end
 
 bratensor(nn::NormNetwork, vertex) = conj(conj_bratensor(nn, vertex))
 
+"""
+    flatten_network(nn::NormNetwork) -> ITensorNetwork
+
+Expand a norm network into a plain tensor network carrying one vertex per layer, so that vertex
+`v` of `nn` becomes the two vertices `(v, :ket)` and `(v, :bra)` and `nv` doubles. A vertex's
+two layers stay adjacent in the vertex order, ket first.
+"""
+function flatten_network(nn::NormNetwork)
+    return ITensorNetwork(
+        dictionary(
+            Iterators.flatten(
+                ((v, :ket) => kettensor(nn, v), (v, :bra) => bratensor(nn, v))
+                    for v in vertices(nn)
+            )
+        )
+    )
+end
+
 """
     normnetwork(tn::ITensorNetwork, [braname]) -> NormNetwork