Merge pull request #156 from numericalEFT/computgraph_pchou

Update FrontEnds API and Bugfix

src/computational_graph/feynmangraph.jl

@@ -489,7 +489,8 @@ julia> g.subgraphs
 ```
 """
 function feynman_diagram(subgraphs::Vector{FeynmanGraph{F,W}}, topology::Vector{Vector{Int}}, perm_noleg::Union{Vector{Int},Nothing}=nothing;
-    factor=one(_dtype.factor), weight=zero(_dtype.weight), name="", diagtype::DiagramType=GenericDiag(), is_signed::Bool=false) where {F,W}
+    contraction_orders::Union{Nothing,Vector{Vector{Int}}}=nothing, factor=one(F), weight=zero(W),
+    name="", diagtype::DiagramType=GenericDiag(), is_signed::Bool=false) where {F,W}
 
     # external_ops = OperatorProduct(operators[external]) # the external operators for the building diagram after contractions
     contraction = collect(Iterators.flatten(topology))
@@ -498,7 +499,11 @@ function feynman_diagram(subgraphs::Vector{FeynmanGraph{F,W}}, topology::Vector{
     vertices, all_external_legs = OperatorProduct[], Bool[]
     external_leg, external_noleg = Int[], Int[] # index all leg/nonleg external operators
     ind = 0
+
+    orders_length = length(orders(subgraphs[1]))
+    diag_orders = zeros(Int, orders_length)
     for g in subgraphs
+        diag_orders += orders(g)
         diagram_type(g) == Propagator && continue  # exclude propagator subgraph to avoid double counting.
         push!(vertices, external_operators(g))
         append!(all_external_legs, external_legs(g))
@@ -534,13 +539,22 @@ function feynman_diagram(subgraphs::Vector{FeynmanGraph{F,W}}, topology::Vector{
         sign = 1
     end
 
-    for connection in topology
-        push!(subgraphs, propagator(operators[connection]))
+    if isnothing(contraction_orders)
+        for (i, connection) in enumerate(topology)
+            push!(subgraphs, propagator(operators[connection]; orders=zeros(Int, orders_length)))
+        end
+    else
+        for (i, connection) in enumerate(topology)
+            propagator_orders = zeros(Int, orders_length)
+            propagator_orders[eachindex(contraction_orders[i])] = contraction_orders[i]
+            push!(subgraphs, propagator(operators[connection]; orders=propagator_orders))
+            diag_orders += propagator_orders
+        end
     end
     _external_indices = union(external_leg, external_noleg)
     _external_legs = append!([true for i in eachindex(external_leg)], [false for i in eachindex(external_noleg)])
     return FeynmanGraph(subgraphs; topology=topology, external_indices=_external_indices, external_legs=_external_legs, vertices=vertices,
-        name=name, diagtype=diagtype, operator=Prod(), factor=factor * sign, weight=weight)
+        orders=diag_orders, name=name, diagtype=diagtype, operator=Prod(), factor=factor * sign, weight=weight)
 end
 
 # do nothing when already a OperatorProduct; 
@@ -554,13 +568,18 @@ _extract_vertex(::Type{<:FeynmanGraph}, g) = OperatorProduct(external_operators(
 
     Create a Propagator-type FeynmanGraph from given OperatorProduct or Vector{QuantumOperator} `ops`, including two quantum operators.
 """
-function propagator(ops::Union{OperatorProduct,Vector{QuantumOperator}};
+function propagator(ops::Union{OperatorProduct,Vector{QuantumOperator}}; orders::Union{Nothing,Vector{Int}}=nothing,
     name="", factor=one(_dtype.factor), weight=zero(_dtype.weight), operator=Sum())
     @assert length(ops) == 2
     @assert adjoint(ops[1].operator) == ops[2].operator
     sign, perm = correlator_order(OperatorProduct(ops))
-    return FeynmanGraph(FeynmanGraph[]; topology=[[1, 2]], external_indices=perm, external_legs=[true, true], vertices=OperatorProduct.(ops),
-        diagtype=Propagator(), name=name, operator=operator, factor=factor * sign, weight=weight)
+    if isnothing(orders)
+        return FeynmanGraph(FeynmanGraph[]; topology=[[1, 2]], external_indices=perm, external_legs=[true, true], vertices=OperatorProduct.(ops),
+            diagtype=Propagator(), name=name, operator=operator, factor=factor * sign, weight=weight)
+    else
+        return FeynmanGraph(FeynmanGraph[]; topology=[[1, 2]], external_indices=perm, external_legs=[true, true], vertices=OperatorProduct.(ops),
+            orders=orders, diagtype=Propagator(), name=name, operator=operator, factor=factor * sign, weight=weight)
+    end
 end
 
 """

src/computational_graph/optimize.jl

@@ -191,11 +191,13 @@ end
 
 # Returns:
 - The vector of unique leaf nodes.
+- The vector of unique leaf nodes' index.
 - A mapping dictionary from the id of each unique leaf node to its index in collect(1:length(leafs)).
 """
 function unique_leaves(graphs::AbstractVector{<:AbstractGraph})
     ############### find the unique Leaves #####################
     unique_graphs = []
+    unique_graphs_id = Int[]
     mapping = Dict{Int,Int}()
 
     idx = 1
@@ -210,11 +212,12 @@ function unique_leaves(graphs::AbstractVector{<:AbstractGraph})
         end
         if flag
             push!(unique_graphs, g)
-            mapping[id(g)] = idx
+            push!(unique_graphs_id, g.id)
+            mapping[g.id] = idx
             idx += 1
         end
     end
-    return unique_graphs, mapping
+    return unique_graphs, unique_graphs_id, mapping
 end
 
 """
@@ -245,14 +248,18 @@ function remove_duplicated_leaves!(graphs::Union{Tuple,AbstractVector{<:Abstract
     sort!(leaves, by=x -> id(x)) #sort the id of the leaves in an asscend order
     unique!(x -> id(x), leaves) #filter out the leaves with the same id number
 
-    _unique_leaves, leafmap = unique_leaves(leaves)
+    _unique_leaves, uniqueleaves_id, mapping = unique_leaves(leaves)
     verbose > 0 && length(leaves) > 0 && println("Number of independent Leaves $(length(leaves)) → $(length(_unique_leaves))")
 
+    leafmap = Dict{Int,Int}()
     for g in graphs
         for n in PreOrderDFS(g)
             for (si, sub_g) in enumerate(subgraphs(n))
                 if isleaf(sub_g)
-                    set_subgraph!(n, _unique_leaves[leafmap[id(sub_g)]], si)
+                    set_subgraph!(n, _unique_leaves[mapping[id(sub_g)]], si)
+                    if sub_g.id ∈ uniqueleaves_id
+                        leafmap[sub_g.id] = mapping[sub_g.id]
+                    end
                 end
             end
         end