diff --git a/Project.toml b/Project.toml index dce5a90..6680757 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ITensorNetworksNext" uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c" -version = "0.9.12" +version = "0.9.13" authors = ["ITensor developers and contributors"] [workspace] diff --git a/src/contract_network.jl b/src/contract_network.jl index a48bdda..1e8f528 100644 --- a/src/contract_network.jl +++ b/src/contract_network.jl @@ -1,7 +1,7 @@ using Base.Broadcast: materialize using Base: @kwdef -using ITensorBase: EvaluationOrderAlgorithm, Greedy, Mul, lazy, optimize_evaluation_order, - substitute, symnameddims +using ITensorBase: EvaluationOrderAlgorithm, Greedy, Mul, ismul, lazy, + optimize_evaluation_order, substitute, symnameddims, to_mul_arguments # `contract_network` @kwdef struct Exact{Order, OrderAlg} @@ -29,19 +29,40 @@ function get_order(alg::Exact, tn) subs = Dict(symnameddims(i) => symnameddims(i, Tuple(axes(t))) for (i, t) in pairs(tn)) return substitute(order, subs) end +# The contraction leaves of an operand: a lazy product (the `NormNetwork` doubled vertex +# `lazy(ket) * lazy(conj(bra))`) contributes each factor, recursively; anything else is one leaf. +leaf_tensors(t) = ismul(t) ? mapreduce(leaf_tensors, vcat, to_mul_arguments(t)) : [t] + # 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. +# +# For a computed order, expand each promoted operand into its contraction leaves so the order +# optimizer sees each factor of a lazy product — and the physical index shared between the ket and +# bra of a doubled vertex — instead of one opaque node with only the outer bond legs. Without this +# the optimizer cannot interleave the other operands between the two layers and is forced to form +# the doubled `ket * conj(bra)` tensor first (χ^(2·degree)). Flattening the *promoted* operand keeps +# the operator/state semantics the promotion just established. An explicit operand-level `order` is +# honored over the operands as given, so it is not flattened. function contract_network(alg::Exact, tn) - order = get_order(alg, tn) + if !isnothing(alg.order) + # Explicit order: honor it over the operands as given, so it is not flattened. + order = get_order(alg, 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) + ) + return materialize(substitute(order, syms_to_ts)) + end + # Computed order: expand each promoted operand into its contraction leaves. 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) + leaves = collect(Iterators.flatten(leaf_tensors(lazy(convert(T, t))) for t in tn)) + order = get_order(alg, leaves) + syms_to_ts = + Dict(symnameddims(i, Tuple(axes(t))) => lazy(t) for (i, t) in pairs(leaves)) + return materialize(substitute(order, syms_to_ts)) end # `contraction_order` diff --git a/test/Project.toml b/test/Project.toml index a49dd6d..b83ab06 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -20,6 +20,7 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb" TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" TensorKitSectors = "13a9c161-d5da-41f0-bcbd-e1a08ae0647f" +TermInterface = "8ea1fca8-c5ef-4a55-8b96-4e9afe9c9a3c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [sources.ITensorNetworksNext] @@ -46,4 +47,5 @@ StableRNGs = "1" Suppressor = "0.2.8" TensorAlgebra = "0.16, 0.17" TensorKitSectors = "0.3" +TermInterface = "2" Test = "1.10" diff --git a/test/test_contract_network.jl b/test/test_contract_network.jl index e567a62..1804df4 100644 --- a/test/test_contract_network.jl +++ b/test/test_contract_network.jl @@ -1,11 +1,12 @@ using Graphs: edges, vertices -using ITensorBase: - Greedy, Index, NamedTensorOperator, inputnames, operator, outputnames, state +using ITensorBase: Greedy, ITensor, Index, NamedTensorOperator, conj, dimnames, inputnames, + lazy, operator, outputnames, state using ITensorNetworksNext: Exact, ITensorNetwork, LeftAssociative, contract_network, - linkinds, siteinds, tensornetwork + get_order, leaf_tensors, linkinds, siteinds, tensornetwork using NamedGraphs.GraphsExtensions: arranged_edges, incident_edges using NamedGraphs.NamedGraphGenerators: named_grid using OMEinsumContractionOrders: ExhaustiveSearch, GreedyMethod, TreeSA +using TermInterface: arguments, iscall using Test: @test, @testset @testset "contract_network" begin @@ -91,4 +92,52 @@ using Test: @test, @testset @test outputnames(rb) == outputnames((op * u) * w) == outputnames(op * (u * w)) @test inputnames(rb) == inputnames((op * u) * w) == inputnames(op * (u * w)) end + + @testset "Flatten lazy product operands into contraction leaves" begin + # A NormNetwork's doubled vertex is a lazy product `lazy(ket) * lazy(conj(bra))`. The + # contraction order must see each factor — and the physical index they share — as its own + # leaf, so the optimizer can interleave the other operands between the two layers instead of + # forming the doubled `ket * conj(bra)` tensor (χ^(2·degree)). + d, χ = 2, 8 + p = Index(d) + b1, b2, b3 = Index(χ), Index(χ), Index(χ) + c1, c2, c3 = Index(χ), Index(χ), Index(χ) + A = ITensor(randn(d, χ, χ, χ), (p, b1, b2, b3)) + B = ITensor(randn(d, χ, χ, χ), (p, c1, c2, c3)) + factor = lazy(A) * lazy(conj(B)) + + # `leaf_tensors` splits a lazy product into its factors, recursively; anything else is one leaf. + @test length(leaf_tensors(factor)) == 2 + @test length(leaf_tensors(lazy(A))) == 1 + @test length(leaf_tensors(lazy(factor) * lazy(A))) == 3 + + # Degree-3 doubled vertex with two incoming bond messages, contracted to the outgoing message. + msg1 = ITensor(randn(χ, χ), (b1, c1)) + msg2 = ITensor(randn(χ, χ), (b2, c2)) + net = [msg1, msg2, factor] + + # The network flattens to degree + 1 = 4 leaves: the ket and bra are separate contractible nodes. + T = mapreduce(typeof, promote_type, net) + leaves = collect(Iterators.flatten(leaf_tensors(lazy(convert(T, t))) for t in net)) + @test length(leaves) == 4 + + # Walk the computed order: a contraction node's open dims are the symmetric difference of its + # children's, its width their size product. The peak stays at the χ^(degree+1) path — no + # intermediate spans all 2·degree bond legs (the χ^(2·degree) doubled tensor). + order = get_order(Exact(), leaves) + idxsize = + Dict(n => s for t in (A, B, msg1, msg2) for (n, s) in zip(dimnames(t), size(t))) + width(dims) = isempty(dims) ? 1 : prod(idxsize[n] for n in dims) + function peakwidth(node) + iscall(node) || return (Set(dimnames(node)), width(dimnames(node))) + (da, pa), (db, pb) = peakwidth.(arguments(node)) + open = symdiff(da, db) + return open, max(pa, pb, width(open)) + end + @test last(peakwidth(order)) ≤ χ^(3 + 1) + @test last(peakwidth(order)) < χ^(2 * 3) + + # The flattened contraction matches forming the doubled vertex explicitly. + @test contract_network(net) ≈ contract_network([msg1, msg2, A, conj(B)]) + end end