diff --git a/Project.toml b/Project.toml index 36c86768..5b863052 100644 --- a/Project.toml +++ b/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 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" diff --git a/docs/Project.toml b/docs/Project.toml index 6adb17f4..eee8bc4e 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -10,5 +10,5 @@ path = ".." [compat] Documenter = "1" ITensorFormatter = "0.2.27" -ITensorNetworksNext = "0.10" +ITensorNetworksNext = "0.11" Literate = "2" diff --git a/examples/Project.toml b/examples/Project.toml index 406118b8..0fa22c56 100644 --- a/examples/Project.toml +++ b/examples/Project.toml @@ -5,4 +5,4 @@ ITensorNetworksNext = "302f2e75-49f0-4526-aef7-d8ba550cb06c" path = ".." [compat] -ITensorNetworksNext = "0.10" +ITensorNetworksNext = "0.11" diff --git a/ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl b/ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl index abede5e5..ca80cf7c 100644 --- a/ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl +++ b/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 diff --git a/src/ITensorNetworksNext.jl b/src/ITensorNetworksNext.jl index fe1b2540..0298a85f 100644 --- a/src/ITensorNetworksNext.jl +++ b/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") diff --git a/src/beliefpropagation/beliefpropagation.jl b/src/beliefpropagation/beliefpropagation.jl index 3741fb3a..fa387253 100644 --- a/src/beliefpropagation/beliefpropagation.jl +++ b/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. diff --git a/src/beliefpropagation/messagecache.jl b/src/beliefpropagation/messagecache.jl index 650a5858..ec91cd58 100644 --- a/src/beliefpropagation/messagecache.jl +++ b/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)) diff --git a/src/contract_network.jl b/src/contract_network.jl deleted file mode 100644 index f13b36c8..00000000 --- a/src/contract_network.jl +++ /dev/null @@ -1,71 +0,0 @@ -using Base.Broadcast: materialize -using Base: @kwdef -using ITensorBase: EvaluationOrderAlgorithm, Greedy, Mul, lazy, optimize_evaluation_order, - substitute, symnamedtensor - -# `contract_network` -@kwdef struct Exact{Order, OrderAlg} - order::Order = nothing - order_alg::OrderAlg = Greedy() -end - -function contract_network(alg, tn) - return throw(ArgumentError("`contract_network` algorithm `$(alg)` not implemented.")) -end -function contract_network(tn; alg = Exact()) - return contract_network(alg, tn) -end - -# `contract_network(::Exact, ...)` -function get_order(alg::Exact, tn) - # Allow specifying either an explicit `order` or an `order_alg` to compute one. - order = if !isnothing(alg.order) - alg.order - else - contraction_order(tn; alg = alg.order_alg) - end - # Contraction order may or may not have indices attached, canonicalize the format - # by attaching indices. - subs = - Dict(symnamedtensor(i) => symnamedtensor(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) - T = mapreduce(typeof, promote_type, tn) - syms_to_ts = Dict( - symnamedtensor(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 - -# `contraction_order` -function contraction_order end -function contraction_order(tn; alg = Greedy()) - return contraction_order(alg, tn) -end -# Convert the tensor network to a flat symbolic multiplication expression. -struct Flat end -function contraction_order(alg::Flat, tn) - # Same as: `reduce((a, b) -> *(a, b; flatten = true), syms)`. - syms = vec([symnamedtensor(i, Tuple(axes(tn[i]))) for i in keys(tn)]) - return lazy(Mul(syms)) -end -struct LeftAssociative end -function contraction_order(alg::LeftAssociative, tn) - return prod(i -> symnamedtensor(i, Tuple(axes(tn[i]))), keys(tn)) -end -# Internal implementation shared with the OMEinsumContractionOrders extension. -function _contraction_order(alg, tn) - s = contraction_order(Flat(), tn) - return optimize_evaluation_order(s; alg) -end -function contraction_order(alg::EvaluationOrderAlgorithm, tn) - return _contraction_order(alg, tn) -end diff --git a/src/contraction_tree.jl b/src/contraction_tree.jl new file mode 100644 index 00000000..0fe6d9e1 --- /dev/null +++ b/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 diff --git a/src/normnetwork.jl b/src/normnetwork.jl index 46725507..d735b22f 100644 --- a/src/normnetwork.jl +++ b/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 diff --git a/src/prod_tensors.jl b/src/prod_tensors.jl new file mode 100644 index 00000000..a9d8a879 --- /dev/null +++ b/src/prod_tensors.jl @@ -0,0 +1,89 @@ +using AbstractTrees: nodevalue +using Combinatorics: combinations +using ITensorBase: inds + +""" + ContractionTreeAlgorithm + +Supertype of the strategies for finding a contraction order. See [`contraction_tree`](@ref). +""" +abstract type ContractionTreeAlgorithm <: AbstractAlgorithm end + +""" + prod_tensors(tensors, tree::ContractionTree) + prod_tensors(tensors, alg) + prod_tensors(tensors) + +Contract `tensors` in the pairwise order given by `tree`, whose leaf labels index into +`tensors`. Given an order algorithm instead, find a tree with it first. Given neither, contract +in key order, folding from the left. + +`tensors` is anything indexable by those labels, so a `Vector` pairs with a tree over positions +and a tensor network pairs with a tree over vertices. +""" +function prod_tensors end +prod_tensors(tensors) = prod_tensors(tensors, left_associative_tree(keys(tensors))) +function prod_tensors(tensors, tree::ContractionTree) + isleaf(tree) && return tensors[nodevalue(tree)] + return prod_tensors(tensors, tree[1]) * prod_tensors(tensors, tree[2]) +end +function prod_tensors(tensors, alg::ContractionTreeAlgorithm) + return prod_tensors(tensors, contraction_tree(alg, tensors)) +end + +# The tree that folds `labels` from the left, so `(a, b, c)` gives `((a, b), c)`. +function left_associative_tree(labels) + isempty(labels) && throw(ArgumentError("No tensors to contract.")) + return reduce(ContractionTree, map(ContractionTree, labels)) +end + +""" + contraction_tree(tensors; alg = Greedy()) + contraction_tree(alg, tensors) + +Find a pairwise contraction order for `tensors`, returned as a [`ContractionTree`](@ref) over +their keys. + +Only the index names and lengths are consulted, never the entries, so an order found once can +be replayed against any network of the same shape. +""" +function contraction_tree end +contraction_tree(tensors; alg = Greedy()) = contraction_tree(alg, tensors) +function contraction_tree(alg, tensors) + return throw(ArgumentError("Contraction order algorithm `$(alg)` not implemented.")) +end + +""" + Greedy + +Repeatedly contract the cheapest available pair, measured as the product of the lengths of all +indices involved. Outer products are taken only once nothing else is left. +""" +struct Greedy <: ContractionTreeAlgorithm end + +function contraction_tree(::Greedy, tensors) + ks = collect(keys(tensors)) + isempty(ks) && throw(ArgumentError("No tensors to contract.")) + trees = map(ContractionTree, ks) + # The inds of each pending operand. An intermediate never materializes, so its inds are + # tracked here rather than read off a tensor. + is = [inds(tensors[k]) for k in ks] + while length(trees) > 1 + i1, i2 = argmin(combinations(eachindex(trees), 2)) do (i, j) + # Defer outer products: with nothing contracted they cost the full product of both + # operands, and taking one early inflates every contraction that follows. + isdisjoint(is[i], is[j]) && return typemax(Int) + # Every index on either operand is iterated once, shared ones counted once. + return prod(length, union(is[i], is[j]); init = 1) + end + # Shared indices are summed over, so the result carries the symmetric difference. + contracted = symdiff(is[i1], is[i2]) + tree = ContractionTree(trees[i1], trees[i2]) + # Remove the pair by position. Removing it by value would also drop any other operand + # equal to it, silently losing a tensor from a network with repeated tensors. + keep = [i for i in eachindex(trees) if i ∉ (i1, i2)] + trees = [trees[keep]; [tree]] + is = [is[keep]; [contracted]] + end + return only(trees) +end diff --git a/test/Project.toml b/test/Project.toml index 1d3a475e..d761064c 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,5 @@ [deps] +AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c" AlgorithmsInterface = "d1e3940c-cd12-4505-8585-b0a4b322527d" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" DataGraphs = "b5a273c3-7e6c-41f6-98bd-8d7f1525a36a" @@ -23,14 +24,15 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" path = ".." [compat] +AbstractTrees = "0.4.5" AlgorithmsInterface = "0.1" Aqua = "0.8.14" DataGraphs = "0.6" Dictionaries = "0.4.5" GradedArrays = "0.17" Graphs = "1.13.1" -ITensorBase = "0.15.5" -ITensorNetworksNext = "0.10" +ITensorBase = "0.16" +ITensorNetworksNext = "0.11" ITensorPkgSkeleton = "0.3.42" MatrixAlgebraKit = "0.6" NamedGraphs = "0.14" diff --git a/test/test_beliefpropagation.jl b/test/test_beliefpropagation.jl index f838f5d3..1591e26f 100644 --- a/test/test_beliefpropagation.jl +++ b/test/test_beliefpropagation.jl @@ -1,18 +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, isdual using Graphs: AbstractGraph, add_vertex!, dst, edges, has_edge, has_vertex, ne, nv, rem_edge!, src, vertices -using ITensorBase: - Greedy, ITensor, Index, apply, inds, name, noprime, outputnames, prime, state -using ITensorNetworksNext: ITensorNetworksNext, Exact, ITensorNetwork, MessageCache, - NormNetwork, SimpleMessageUpdate, StopWhenConverged, beliefpropagation, - bethe_free_energy, bethe_free_entropy, bratensor, contract_network, contraction_order, +using ITensorBase: ITensor, Index, apply, inds, name, noprime, outputnames, prime, state +using ITensorNetworksNext: ITensorNetworksNext, ContractionTreeAlgorithm, Greedy, + ITensorNetwork, MessageCache, NormNetwork, SimpleMessageUpdate, StopWhenConverged, + beliefpropagation, bethe_free_energy, bethe_free_entropy, bratensor, contraction_tree, edge_scalar, edge_scalars, factor_tensors, incoming_messages, insertlink!, kettensor, - linkaxes, linkinds, message_environment, messagecache, region_scalar, subgraph, - tensornetwork, updated_message, vertex_scalar, vertex_scalars + linkaxes, linkinds, message_environment, messagecache, prod_tensors, region_scalar, + subgraph, tensornetwork, updated_message, vertex_scalar, vertex_scalars using LinearAlgebra: LinearAlgebra, norm, tr using NamedGraphs: NamedEdge, all_edges, incident_edges, named_comb_tree, named_grid, named_path_graph, vertextype @@ -43,13 +41,13 @@ 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 +# Records how many operands each contraction is given, then orders them greedily. +struct RecordOperands <: ContractionTreeAlgorithm counts::Vector{Int} end -function ITensorNetworksNext.contraction_order(alg::RecordOperands, tn) +function ITensorNetworksNext.contraction_tree(alg::RecordOperands, tn) push!(alg.counts, length(tn)) - return contraction_order(tn; alg = Greedy()) + return contraction_tree(tn; alg = Greedy()) end @testset "Belief propagation" begin @@ -345,7 +343,7 @@ end for (edge, rest) in ((1 => 2, 2:4), (4 => 3, 1:3)) layers = [[kettensor(nn, v) for v in rest]; [bratensor(nn, v) for v in rest]] - z_rest = contract_network([state(ones[edge]); layers])[] + z_rest = prod_tensors([state(ones[edge]); layers], Greedy())[] @test z_rest ≈ norm(prod([network[v] for v in rest]))^2 rtol = eps(real(T))^(1 / 3) end @@ -369,26 +367,24 @@ 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 doubled vertex splits into its two layers. + @test factor_tensors(nn, v) == [kettensor(nn, v), bratensor(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. + # The message update passes the layers 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)) - ) + algorithm = SimpleMessageUpdate(; 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]]]) + @test message ≈ + prod_tensors([messages; [kettensor(nn, v) * bratensor(nn, v)]], Greedy()) end end end diff --git a/test/test_contraction_tree.jl b/test/test_contraction_tree.jl new file mode 100644 index 00000000..54538553 --- /dev/null +++ b/test/test_contraction_tree.jl @@ -0,0 +1,132 @@ +using AbstractTrees: AbstractTrees, Leaves +using Graphs: edges, vertices +using ITensorBase: Index +using ITensorNetworksNext: Branch, ContractionTree, Greedy, contraction_tree, isleaf, + left_associative_tree, prod_tensors, tensornetwork +using NamedGraphs: incident_edges, named_grid +using OMEinsumContractionOrders: GreedyMethod, TreeSA +using Test: @test, @testset + +@testset "ContractionTree" begin + @testset "leaves and branches" begin + l = ContractionTree(:a) + @test isleaf(l) + @test AbstractTrees.nodevalue(l) == :a + @test AbstractTrees.children(l) == () + + # (a, (b, c)) + t = ContractionTree( + ContractionTree(:a), + ContractionTree(ContractionTree(:b), ContractionTree(:c)) + ) + @test !isleaf(t) + @test length(AbstractTrees.children(t)) == 2 + @test t[1] == ContractionTree(:a) + @test isleaf(t[1]) + @test !isleaf(t[2]) + @test AbstractTrees.nodevalue(t[2][1]) == :b + @test AbstractTrees.nodevalue(t[2][2]) == :c + end + + @testset "a branch comes from `Branch`, never from a bare tuple" begin + a, b = ContractionTree(1), ContractionTree(2) + @test ContractionTree(Branch(a, b)) == ContractionTree(a, b) + @test !isleaf(ContractionTree(Branch(a, b))) + + # A bare 2-tuple is a label, even when its entries are themselves trees. This is what + # keeps a tuple-valued vertex label from being read as children. + @test isleaf(ContractionTree((a, b))) + @test AbstractTrees.nodevalue(ContractionTree((a, b))) == (a, b) + end + + @testset "equality and hashing" begin + t1 = ContractionTree(ContractionTree(:a), ContractionTree(:b)) + t2 = ContractionTree(ContractionTree(:a), ContractionTree(:b)) + t3 = ContractionTree(ContractionTree(:b), ContractionTree(:a)) + @test t1 == t2 + @test hash(t1) == hash(t2) + @test t1 != t3 + # a leaf whose label collides structurally is not equal to a branch + @test ContractionTree(:a) != + ContractionTree(ContractionTree(:a), ContractionTree(:a)) + end + + @testset "arbitrary label types, including Vector-valued leaves" begin + # The ambiguity that nested `Vector` cannot represent: a leaf whose label is itself a + # vector. Leaf/branch is a type distinction, so this is unambiguous. + l = ContractionTree([1, 2]) + @test isleaf(l) + @test AbstractTrees.nodevalue(l) == [1, 2] + + t = ContractionTree(ContractionTree([1, 2]), ContractionTree([3, 4])) + @test !isleaf(t) + @test AbstractTrees.nodevalue(t[1]) == [1, 2] + @test AbstractTrees.nodevalue(t[2]) == [3, 4] + end + + @testset "tuple-valued vertex labels are leaves, not branches" begin + # A `Tuple`-valued label (e.g. grid coordinates) must not be mistaken for a branch. + # Leaves and branches are separate types, so a label can never be read as children. + l = ContractionTree((1, 2)) + @test isleaf(l) + @test AbstractTrees.nodevalue(l) == (1, 2) + + t = ContractionTree(ContractionTree((1, 2)), ContractionTree((3, 4))) + @test !isleaf(t) + @test isleaf(t[1]) + @test AbstractTrees.nodevalue(t[1]) == (1, 2) + @test AbstractTrees.nodevalue(t[2]) == (3, 4) + end + + @testset "show" begin + t = ContractionTree( + ContractionTree(1), + ContractionTree(ContractionTree(2), ContractionTree(3)) + ) + @test sprint(show, t) == "(1, (2, 3))" + @test sprint(show, ContractionTree(:a)) == ":a" + end + + @testset "`prod_tensors` follows the tree" begin + i, j, k = Index(2), Index(2), Index(5) + A = [1.0 1.0; 0.5 1.0][i, j] + B = [2.0, 1.0][i] + C = [5.0, 1.0][j] + D = [-2.0, 3.0, 4.0, 5.0, 1.0][k] + ts = [A, B, C, D] + + leaves = map(ContractionTree, 1:4) + left = left_associative_tree(1:4) + @test left == reduce(ContractionTree, leaves) + balanced = ContractionTree( + ContractionTree(leaves[1], leaves[2]), ContractionTree(leaves[3], leaves[4]) + ) + @test prod_tensors(ts, left) ≈ ((A * B) * C) * D + @test prod_tensors(ts, balanced) ≈ (A * B) * (C * D) + end + + @testset "`contraction_tree` over a grid" begin + g = named_grid((2, 3)) + l = Dict(e => Index(2) for e in edges(g)) + l = merge(l, Dict(reverse(e) => l[e] for e in edges(g))) + tn = tensornetwork(vertices(g)) do v + return randn(Tuple(map(e -> l[e], incident_edges(g, v)))) + end + + for alg in (Greedy(), GreedyMethod(), TreeSA()) + tree = contraction_tree(tn; alg) + # Every vertex is contracted exactly once, under its own key. + @test length(collect(Leaves(tree))) == length(vertices(g)) + @test issetequal(map(AbstractTrees.nodevalue, Leaves(tree)), vertices(g)) + @test prod_tensors(tn, tree)[] ≈ prod_tensors(tn, Greedy())[] + end + + # A tree depends only on the network's shape, so it replays against a second network + # built over the same graph and indices. + tree = contraction_tree(tn; alg = Greedy()) + tn2 = tensornetwork(vertices(g)) do v + return randn(Tuple(map(e -> l[e], incident_edges(g, v)))) + end + @test prod_tensors(tn2, tree)[] ≈ prod_tensors(tn2, Greedy())[] + end +end diff --git a/test/test_itensornetworkgenerators.jl b/test/test_itensornetworkgenerators.jl index be154eb1..0dc9f4b7 100644 --- a/test/test_itensornetworkgenerators.jl +++ b/test/test_itensornetworkgenerators.jl @@ -2,7 +2,7 @@ using DataGraphs: DataGraph using Graphs: edges, ne, nv, vertices using ITensorBase: Index, inds using ITensorNetworksNext.ITensorNetworkGenerators: delta, delta_network, ising_network -using ITensorNetworksNext: contract_network +using ITensorNetworksNext: Greedy, prod_tensors using NamedGraphs: arranged_edges, incident_edges, named_grid using Test: @test, @testset @@ -46,7 +46,7 @@ using Test: @test, @testset @test issetequal(is, inds(tn[v])) @test tn[v] ≠ delta(Float64, Tuple(is)) end - z = contract_network(tn)[] + z = prod_tensors(tn, Greedy())[] f = -log(z) / (β * nv(g)) f_analytic = TestUtils.f_1d_ising(β, 4; periodic) @test f ≈ f_analytic @@ -67,7 +67,7 @@ using Test: @test, @testset @test issetequal(is, inds(tn[v])) @test tn[v] ≠ delta(Float64, Tuple(is)) end - z = contract_network(tn)[] + z = prod_tensors(tn, Greedy())[] f = -log(z) / (β * nv(g)) f_inf = TestUtils.f_2d_ising(β) @test f ≈ f_inf rtol = 1.0e-1 diff --git a/test/test_normnetwork.jl b/test/test_normnetwork.jl index 5007be69..4528b220 100644 --- a/test/test_normnetwork.jl +++ b/test/test_normnetwork.jl @@ -1,11 +1,10 @@ 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, contract_network, indmap, kettensor, normnetwork, - tensornetwork +using ITensorBase: ITensor, Index, IndexName, conj, inds, name, setname, uniquename +using ITensorNetworksNext: BraView, Greedy, ITensorNetwork, KetView, NormNetwork, braname, + bratensor, conj_bratensor, contraction_tree, flatten_network, indmap, kettensor, + normnetwork, prod_tensors, tensornetwork using LinearAlgebra: norm using NamedGraphs: NamedEdge, incident_edges, named_grid, named_path_graph using Test: @test, @test_throws, @testset @@ -23,6 +22,10 @@ function random_state(::Type{T}, g; d = 2, χ = 2) where {T} return tn, l, s end +# A norm network carries no tensor at a vertex, so contracting it goes through the flattened +# network, whose two layers are separate operands. +contract_norm(nn) = prod_tensors(flatten_network(nn), Greedy())[] + @testset "`NormNetwork`" begin @testset "Basics" begin g = named_path_graph(3) @@ -37,8 +40,10 @@ end @test issetequal(vertices(nn), vertices(tn)) @test issetequal(edges(nn), edges(tn)) - # `eltype` is the type of the (lazy double-layer) vertex data. - @test eltype(nn) === typeof(nn[1]) + # The double-layer tensor at a vertex has no representation, so neither indexing nor + # `eltype` is defined. + @test_throws ErrorException nn[1] + @test_throws ErrorException eltype(nn) # Vertex data is assigned wherever the underlying network is. @test is_vertex_assigned(nn, 1) @@ -126,6 +131,29 @@ end @test !isinsertable(bv) end + @testset "`flatten_network`" begin + g = named_path_graph(3) + tn, l, s = random_state(Float64, g) + nn = NormNetwork(tn) + flat = flatten_network(nn) + + # One vertex per layer, so the vertex count doubles and each carries its layer's tensor. + @test length(vertices(flat)) == 2 * length(vertices(nn)) + # Each vertex's two layers stay adjacent, in ket-then-bra order. + @test collect(vertices(flat)) == + collect(Iterators.flatten(((v, :ket), (v, :bra)) for v in vertices(nn))) + @test flat[(2, :ket)] === kettensor(nn, 2) + @test flat[(2, :bra)] == bratensor(nn, 2) + + # A norm network carries no tensor at a vertex, so it has to be flattened first. + @test_throws ErrorException prod_tensors(nn, Greedy()) + + # A tree is keyed by the vertices of the network it was built over, so one built over + # the flattened network replays against it. + tree = contraction_tree(flat) + @test prod_tensors(flat, tree)[] ≈ prod_tensors(flat, Greedy())[] + end + @testset "contraction / physics" begin @testset "single normalized tensor contracts to 1" begin s = Index(3) @@ -135,7 +163,7 @@ end nn = NormNetwork(tn) # ⟨ψ|ψ⟩ for a single normalized site tensor is 1. - @test contract_network(nn)[] ≈ 1 + @test contract_norm(nn) ≈ 1 end @testset "$T" for T in (Float64, ComplexF64) @@ -143,19 +171,18 @@ end tn, l, s = random_state(T, g) # The norm network contracts to ⟨tn|tn⟩ = ‖prod(tn)‖², a real nonnegative number. - z = contract_network(NormNetwork(tn))[] + z = contract_norm(NormNetwork(tn)) @test z ≈ norm(prod(tn))^2 @test imag(z) ≈ 0 atol = 1.0e-12 * abs(z) @test real(z) > 0 # Rescaling a single tensor by 1/√z normalizes the state, so ⟨tn|tn⟩ = 1. tn[first(vertices(tn))] = tn[first(vertices(tn))] / sqrt(real(z)) - @test contract_network(NormNetwork(tn))[] ≈ 1 + @test contract_norm(NormNetwork(tn)) ≈ 1 # The contracted norm does not depend on the chosen bra-layer name map. custom = map(uniquename, keys(tn.dimname_vertices)) - @test contract_network(normnetwork(tn, custom))[] ≈ - contract_network(NormNetwork(tn))[] + @test contract_norm(normnetwork(tn, custom)) ≈ contract_norm(NormNetwork(tn)) end end end diff --git a/test/test_contract_network.jl b/test/test_prod_tensors.jl similarity index 67% rename from test/test_contract_network.jl rename to test/test_prod_tensors.jl index 667b7d54..b31d12bd 100644 --- a/test/test_contract_network.jl +++ b/test/test_prod_tensors.jl @@ -1,15 +1,12 @@ using Graphs: edges, vertices -using ITensorBase: - Greedy, Index, NamedTensorOperator, inputnames, operator, outputnames, state -using ITensorNetworksNext: Exact, ITensorNetwork, LeftAssociative, contract_network, - linkinds, siteinds, tensornetwork +using ITensorBase: Index, NamedTensorOperator, inputnames, operator, outputnames, state +using ITensorNetworksNext: + Greedy, ITensorNetwork, linkinds, prod_tensors, siteinds, tensornetwork using NamedGraphs: incident_edges, named_grid using OMEinsumContractionOrders: ExhaustiveSearch, GreedyMethod, TreeSA using Test: @test, @testset -@testset "contract_network" begin - orderalg = order_alg -> Exact(; order_alg) - +@testset "prod_tensors" begin @testset "Contract Vectors of ITensors" begin i, j, k = Index(2), Index(2), Index(5) A = [1.0 1.0; 0.5 1.0][i, j] @@ -17,11 +14,12 @@ using Test: @test, @testset C = [5.0, 1.0][j] D = [-2.0, 3.0, 4.0, 5.0, 1.0][k] - ABCD_1 = contract_network([A, B, C, D]; alg = orderalg(LeftAssociative())) - ABCD_2 = contract_network([A, B, C, D]; alg = orderalg(Greedy())) - ABCD_3 = contract_network([A, B, C, D]; alg = orderalg(ExhaustiveSearch())) - ABCD_4 = contract_network([A, B, C, D]; alg = orderalg(GreedyMethod())) - ABCD_5 = contract_network([A, B, C, D]; alg = orderalg(TreeSA())) + ts = [A, B, C, D] + ABCD_1 = prod_tensors(ts) + ABCD_2 = prod_tensors(ts, Greedy()) + ABCD_3 = prod_tensors(ts, ExhaustiveSearch()) + ABCD_4 = prod_tensors(ts, GreedyMethod()) + ABCD_5 = prod_tensors(ts, TreeSA()) @test ABCD_1 == ABCD_2 == ABCD_3 @test ABCD_1 ≈ ABCD_4 @test ABCD_1 ≈ ABCD_5 @@ -37,11 +35,11 @@ using Test: @test, @testset return randn(Tuple(is)) end - z1 = contract_network(tn; alg = orderalg(LeftAssociative()))[] - z2 = contract_network(tn; alg = orderalg(Greedy()))[] - z3 = contract_network(tn; alg = orderalg(ExhaustiveSearch()))[] - z4 = contract_network(tn; alg = orderalg(GreedyMethod()))[] - z5 = contract_network(tn; alg = orderalg(TreeSA()))[] + z1 = prod_tensors(tn)[] + z2 = prod_tensors(tn, Greedy())[] + z3 = prod_tensors(tn, ExhaustiveSearch())[] + z4 = prod_tensors(tn, GreedyMethod())[] + z5 = prod_tensors(tn, TreeSA())[] @test abs(z1 - z2) / abs(z1) <= 1.0e3 * eps(Float64) @test abs(z1 - z3) / abs(z1) <= 1.0e3 * eps(Float64) @@ -59,7 +57,7 @@ using Test: @test, @testset # 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]) + r = prod_tensors([o, t]) @test r isa NamedTensorOperator @test state(r) ≈ state(o * t) @test outputnames(r) == outputnames(o * t) @@ -67,18 +65,18 @@ using Test: @test, @testset # An all-operator network is likewise preserved. o2 = operator(randn(2, 2), (j,), (k,)) - r2 = contract_network([o, o2]) + r2 = prod_tensors([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)[] + @test prod_tensors([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) + @test !(prod_tensors([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). @@ -86,7 +84,7 @@ using Test: @test, @testset 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]) + rb = prod_tensors([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 diff --git a/test/test_tensornetwork.jl b/test/test_tensornetwork.jl index d9639475..0beaad01 100644 --- a/test/test_tensornetwork.jl +++ b/test/test_tensornetwork.jl @@ -2,7 +2,7 @@ using DataGraphs: DataGraph, assigned_edge_data, assigned_vertex_data, underlying_graph, vertex_data using Graphs: add_edge!, add_vertex!, dst, edges, edgetype, has_edge, has_vertex, is_directed, ne, nv, rem_edge!, rem_vertex!, src, vertices -using ITensorBase: Index, LazyITensor, inds, operator +using ITensorBase: Index, inds, operator using ITensorNetworksNext: ITensorNetwork, has_ind, linkaxes, linkinds, linknames, operator_support, siteaxes, siteinds, sitenames, tensornetwork using NamedGraphs: convert_vertextype, incident_edges, named_grid, named_path_graph,