From a4506ee662f19d224e3766fda8176bde290144cc Mon Sep 17 00:00:00 2001
From: Tao Wang <fsxbhyy@gmail.com>
Date: Mon, 30 Oct 2023 22:54:23 -0400
Subject: [PATCH 1/3] change taylor expansion api. Add count operation for
 taylor series

---
 example/taylor_expansion.jl                | 75 ++++++++++++-----
 src/computational_graph/tree_properties.jl |  2 +-
 src/utility.jl                             | 95 ++++++++++++++++------
 3 files changed, 123 insertions(+), 49 deletions(-)

diff --git a/example/taylor_expansion.jl b/example/taylor_expansion.jl
index 2cfdb0af..76cd9ae5 100644
--- a/example/taylor_expansion.jl
+++ b/example/taylor_expansion.jl
@@ -1,9 +1,35 @@
 using FeynmanDiagram
 using FeynmanDiagram.Taylor
 using FeynmanDiagram.ComputationalGraphs:
-    eval!, forwardAD, node_derivative, backAD, build_all_leaf_derivative, count_operation
+    eval!, forwardAD, node_derivative, backAD, build_all_leaf_derivative
 using FeynmanDiagram.Utility:
-    taylorexpansion, build_derivative_backAD!
+    taylorexpansion!, taylorexpansion, build_derivative_backAD!, count_operation
+
+function banchmark_AD(glist::Vector{T}) where {T<:Graph}
+    taylormap = Dict{Int,TaylorSeries{T}}()
+    totaloperation = [0, 0]
+    taylorlist = Vector{TaylorSeries{T}}()
+    for g in glist
+        @time taylorexpansion!(g, taylormap)
+
+        t = taylormap[g.id]
+        operation = count_operation(t)
+        totaloperation = totaloperation + operation
+        push!(taylorlist, t)
+        print("operation number: $(operation)\n")
+        t_compare = build_derivative_backAD!(g)
+        for (order, coeff) in (t_compare.coeffs)
+            @assert (eval!(coeff)) == (eval!(Taylor.taylor_factorial(order) * t.coeffs[order]))
+            # gs = Compilers.to_julia_str([coeff,], name="eval_graph!")
+            # println("$(order)  ", gs, "\n")
+            # print("$(order) $(eval!(coeff)) $(eval!(getderivative(t5,order))) $(count_operation(coeff))\n")
+        end
+    end
+
+    total_uniqueoperation = count_operation(taylorlist)
+    print(" total operation number: $(length(taylorlist)) $(totaloperation) $(total_uniqueoperation)\n")
+    return total_uniqueoperation
+end
 g1 = Graph([])
 g2 = Graph([])
 g3 = Graph([]) #, factor=2.0)
@@ -13,6 +39,7 @@ g6 = Graph([])
 G3 = g1
 G4 = 1.0 * g1 * g2
 G5 = 1.0 * (3.0 * G3 + 0.5 * G4)
+#G5 = g5 * g6
 #G6 = (0.5 * g1 * g2 + 0.5 * g2 * g3)
 #G6 = (g1 + g2) * (0.5 * g1 + g3) * g1 #   (0.5 * g1 + g3)
 #G6 = g1 * g2 * g3 * g4 * g5 * g6
@@ -25,28 +52,34 @@ using FeynmanDiagram.Taylor:
 # set_variables("x y", order=3)
 # @time T5 = taylorexpansion(G5)
 # print(T5)
-set_variables("x y", order=[3, 2])
+set_variables("x y", orders=[3, 2])
 #set_variables("x y z a", order=[1, 2, 3, 2])
-@time T5 = taylorexpansion(G6)
-#order = [0, 0, 0, 0, 0, 0]
-#@time print(T5.coeffs[order])
-print("$(count_operation(T5.coeffs))\n")
-# for (order, coeff) in (T5.coeffs)
-#     #gs = Compilers.to_julia_str([coeff,], name="eval_graph!")
-#     #println("$(order)  ", gs, "\n")
-#     print("$(order) $(eval!(coeff)) $(eval!(getcoeff(T5,order))) $(coeff.id) $(count_operation(coeff))\n")
-# end
 
-print("TaylorSeries $(T5)\n")
 
-@time T5_compare = build_derivative_backAD!(G6)
-print("$(count_operation(T5_compare.coeffs))\n")
-for (order, coeff) in (T5_compare.coeffs)
-    @assert (eval!(coeff)) == (eval!(Taylor.taylor_factorial(order) * T5.coeffs[order]))
-    # gs = Compilers.to_julia_str([coeff,], name="eval_graph!")
-    # println("$(order)  ", gs, "\n")
-    # print("$(order) $(eval!(coeff)) $(eval!(getderivative(T5,order))) $(count_operation(coeff))\n")
-end
+#@time taylormap = taylorexpansion(G6)
+
+
+banchmark_AD([G3, G4, G5, G6])
+# T5 = taylormap[G6.id]
+# #order = [0, 0, 0, 0, 0, 0]
+# #@time print(T5.coeffs[order])
+# print("$(count_operation(T5.coeffs))\n")
+# # for (order, coeff) in (T5.coeffs)
+# #     #gs = Compilers.to_julia_str([coeff,], name="eval_graph!")
+# #     #println("$(order)  ", gs, "\n")
+# #     print("$(order) $(eval!(coeff)) $(eval!(getcoeff(T5,order))) $(coeff.id) $(count_operation(coeff))\n")
+# # end
+
+# print("TaylorSeries $(T5)\n")
+
+# @time T5_compare = build_derivative_backAD!(G6)
+# print("$(count_operation(T5_compare.coeffs))\n")
+# for (order, coeff) in (T5_compare.coeffs)
+#     @assert (eval!(coeff)) == (eval!(Taylor.taylor_factorial(order) * T5.coeffs[order]))
+#     # gs = Compilers.to_julia_str([coeff,], name="eval_graph!")
+#     # println("$(order)  ", gs, "\n")
+#     # print("$(order) $(eval!(coeff)) $(eval!(getderivative(T5,order))) $(count_operation(coeff))\n")
+# end
 
 
 
diff --git a/src/computational_graph/tree_properties.jl b/src/computational_graph/tree_properties.jl
index b315e13d..82f4cce4 100644
--- a/src/computational_graph/tree_properties.jl
+++ b/src/computational_graph/tree_properties.jl
@@ -156,7 +156,7 @@ function count_operation(g::Array{G}) where {G<:Graph}
 end
 
 
-function count_operation(g::Dict{Array{Int,1},G}) where {G<:Graph}
+function count_operation(g::Dict{Vector{Int},G}) where {G<:Graph}
     visited = Set{Int}()
     totalsum = 0
     totalprod = 0
diff --git a/src/utility.jl b/src/utility.jl
index 3deea045..a47311bb 100644
--- a/src/utility.jl
+++ b/src/utility.jl
@@ -2,24 +2,18 @@ module Utility
 using ..ComputationalGraphs
 using ..ComputationalGraphs: Sum, Prod, Power, decrement_power
 using ..ComputationalGraphs: build_all_leaf_derivative, eval!
+import ..ComputationalGraphs: count_operation
 using ..ComputationalGraphs.AbstractTrees
 
 using ..Taylor
 
 
-"""
-    function taylorexpansion(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
-
-    Return a taylor series of graph g. If variable dependence is not specified, by default, assume all leaves of graph depend on all variables. 
-
-#Arguments
+# Internal function that performs taylor expansion on a single graph node recursively.
+function _taylorexpansion!(graph::G, taylormap::Dict{Int,TaylorSeries{G}}, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
+    if haskey(taylormap, graph.id) #If already exist, use taylor series in taylormap.
+        return taylormap[graph.id]
 
-- `graph` Target graph.
-- `var_dependence::Dict{Int,Vector{Bool}}` The variables graph leaves depend on. Should map each leaf ID of g to a Vector{Bool},
-    indicating the taylor variables it depends on. By default, assumes all leaves depend on all variables.
-"""
-function taylorexpansion(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
-    if isleaf(graph)
+    elseif isleaf(graph)
         maxorder = get_orders()
         if haskey(var_dependence, graph.id)
             var = var_dependence[graph.id]
@@ -33,25 +27,54 @@ function taylorexpansion(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{I
             coeff = Graph([]; operator=Sum(), factor=graph.factor)
             result.coeffs[o] = coeff
         end
+        taylormap[graph.id] = result
         return result
     else
-        taylormap = Dict{Int,TaylorSeries{G}}() #Saves the taylor series corresponding to each nodes of the graph
-        for g in Leaves(graph)
-            if !haskey(taylormap, g.id)
-                taylormap[g.id] = taylorexpansion(g, var_dependence)
-            end
-        end
+        taylormap[graph.id] = apply(graph.operator, [taylorexpansion!(sub, taylormap, var_dependence) for sub in graph.subgraphs], graph.subgraph_factors)
+        return taylormap[graph.id]
+    end
+end
 
-        rootid = -1
-        for g in PostOrderDFS(graph) # postorder traversal will visit all subdiagrams of a diagram first
-            rootid = g.id
-            if isleaf(g) || haskey(taylormap, g.id)
-                continue
-            end
-            taylormap[g.id] = apply(g.operator, [taylormap[sub.id] for sub in g.subgraphs], g.subgraph_factors)
-        end
-        return taylormap[rootid]
+
+"""
+    function taylorexpansion!(graph::G, taylormap::Dict{Int,TaylorSeries{G}}(), var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
+
+    Return a taylor series of graph g. If variable dependence is not specified, by default, assume all leaves of graph depend on all variables. The taylor series of all nodes of graph is 
+    saved in the taylormap dictionary.
+
+#Arguments
+
+- `graph` Target graph.
+- `var_dependence::Dict{Int,Vector{Bool}}` The variables graph leaves depend on. Should map each leaf ID of g to a Vector{Bool},
+    indicating the taylor variables it depends on. By default, assumes all leaves depend on all variables.
+- `taylormap::Dict{Int,TaylorSeries{G}}` The taylor series correponding to graph nodes. Should map each graph node ID (does not need to belong to input graph) to a taylor series.
+    All new taylor series generated by taylor expansion will be added into the expansion.
+"""
+function taylorexpansion!(graph::G, taylormap::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}(), var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
+    _taylorexpansion!(graph, taylormap, var_dependence)
+end
+
+
+function taylorexpansion!(graphs::Vector{G}, taylormap::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}(), var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
+    for graph in graphs
+        _taylorexpansion!(graph, taylormap, var_dependence)
+    end
+end
+
+
+function taylorexpansion(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
+    taylormap = Dict{Int,TaylorSeries{G}}()
+    _taylorexpansion!(graph, taylormap, var_dependence)
+    return taylormap
+end
+
+
+function taylorexpansion(graphs::Vector{G}, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
+    taylormap = Dict{Int,TaylorSeries{G}}()
+    for graph in graphs
+        _taylorexpansion!(graph, taylormap, var_dependence)
     end
+    return taylormap
 end
 
 """
@@ -225,4 +248,22 @@ function chainrule!(varidx::Int, dg::Array{G,1}, leaftaylor::Dict{Int,TaylorSeri
     end
 end
 
+function count_operation(g::TaylorSeries{G}) where {G<:Graph}
+    return count_operation(g.coeffs)
+end
+
+function count_operation(graphs::Vector{TaylorSeries{G}}) where {G<:Graph}
+    if length(graphs) == 0
+        return [0, 0]
+    else
+        allcoeffs = Vector{G}()
+        for g in graphs
+            for (order, coeffs) in g.coeffs
+                push!(allcoeffs, coeffs)
+            end
+        end
+        return count_operation(allcoeffs)
+    end
+end
+
 end
\ No newline at end of file

From 04f0a8a5ce2f510df74c6992f0add61aa44eb998 Mon Sep 17 00:00:00 2001
From: Tao Wang <fsxbhyy@gmail.com>
Date: Thu, 2 Nov 2023 15:09:52 -0400
Subject: [PATCH 2/3] clean up api; add diagram->graph conversion

---
 example/taylor_expansion.jl |   6 ++-
 src/frontend/diagtree.jl    |  22 ++++++++
 src/utility.jl              |  24 +++++----
 test/diagram_tree.jl        |   9 ++--
 test/taylor.jl              | 100 +++++++++++++++++++++++++++++++++++-
 5 files changed, 144 insertions(+), 17 deletions(-)

diff --git a/example/taylor_expansion.jl b/example/taylor_expansion.jl
index 76cd9ae5..346bc999 100644
--- a/example/taylor_expansion.jl
+++ b/example/taylor_expansion.jl
@@ -10,9 +10,9 @@ function banchmark_AD(glist::Vector{T}) where {T<:Graph}
     totaloperation = [0, 0]
     taylorlist = Vector{TaylorSeries{T}}()
     for g in glist
-        @time taylorexpansion!(g, taylormap)
+        @time t, taylormap = taylorexpansion!(g; taylormap=taylormap)
+
 
-        t = taylormap[g.id]
         operation = count_operation(t)
         totaloperation = totaloperation + operation
         push!(taylorlist, t)
@@ -60,6 +60,8 @@ set_variables("x y", orders=[3, 2])
 
 
 banchmark_AD([G3, G4, G5, G6])
+
+
 # T5 = taylormap[G6.id]
 # #order = [0, 0, 0, 0, 0, 0]
 # #@time print(T5.coeffs[order])
diff --git a/src/frontend/diagtree.jl b/src/frontend/diagtree.jl
index e0c5c43f..28495675 100644
--- a/src/frontend/diagtree.jl
+++ b/src/frontend/diagtree.jl
@@ -55,4 +55,26 @@ function Graph!(d::DiagTree.Diagram{W}; map=Dict{Int,DiagTree.DiagramId}()) wher
     map[tree.id] = d.id
 
     return root, map
+end
+
+"""
+    function extract_var_dependence(map::Dict{Int,DiagTree.DiagramId}, ::Type{ID}, numvars::Int)
+    
+    Given a map between graph id and DiagramId, extract the variable dependence of all graphs.
+
+# Arguments:
+- `map::Dict{Int,DiagTree.DiagramId}`:  A dictionary mapping graph ids to DiagramIds. DiagramId stores the diagram information of the corresponding graph. 
+- `ID`:  The particular type of ID that has the given variable dependence. 
+- `numvars`: The number of varibles which the diagram depends on.
+"""
+function extract_var_dependence(map::Dict{Int,DiagTree.DiagramId}, ::Type{ID}, numvars::Int) where {ID<:PropagatorId}
+    var_dependence = Dict{Int,Vector{Bool}}()
+    for (id, diagID) in map
+        if diagID isa ID
+            var_dependence[id] = [true for _ in 1:numvars]
+        else
+            var_dependence[id] = [false for _ in 1:numvars]
+        end
+    end
+    return var_dependence
 end
\ No newline at end of file
diff --git a/src/utility.jl b/src/utility.jl
index a47311bb..360c8992 100644
--- a/src/utility.jl
+++ b/src/utility.jl
@@ -9,7 +9,7 @@ using ..Taylor
 
 
 # Internal function that performs taylor expansion on a single graph node recursively.
-function _taylorexpansion!(graph::G, taylormap::Dict{Int,TaylorSeries{G}}, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
+function _taylorexpansion!(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}(); taylormap::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}()) where {G<:Graph}
     if haskey(taylormap, graph.id) #If already exist, use taylor series in taylormap.
         return taylormap[graph.id]
 
@@ -30,7 +30,7 @@ function _taylorexpansion!(graph::G, taylormap::Dict{Int,TaylorSeries{G}}, var_d
         taylormap[graph.id] = result
         return result
     else
-        taylormap[graph.id] = apply(graph.operator, [taylorexpansion!(sub, taylormap, var_dependence) for sub in graph.subgraphs], graph.subgraph_factors)
+        taylormap[graph.id] = apply(graph.operator, [_taylorexpansion!(sub, var_dependence; taylormap=taylormap) for sub in graph.subgraphs], graph.subgraph_factors)
         return taylormap[graph.id]
     end
 end
@@ -50,31 +50,35 @@ end
 - `taylormap::Dict{Int,TaylorSeries{G}}` The taylor series correponding to graph nodes. Should map each graph node ID (does not need to belong to input graph) to a taylor series.
     All new taylor series generated by taylor expansion will be added into the expansion.
 """
-function taylorexpansion!(graph::G, taylormap::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}(), var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
-    _taylorexpansion!(graph, taylormap, var_dependence)
+function taylorexpansion!(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}(); taylormap::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}()) where {G<:Graph}
+    _taylorexpansion!(graph, var_dependence; taylormap=taylormap)
+    return taylormap[graph.id], taylormap
 end
 
 
-function taylorexpansion!(graphs::Vector{G}, taylormap::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}(), var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
+function taylorexpansion!(graphs::Vector{G}, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}(); taylormap::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}()) where {G<:Graph}
+    result = Vector{TaylorSeries{G}}()
     for graph in graphs
-        _taylorexpansion!(graph, taylormap, var_dependence)
+        push!(result, _taylorexpansion!(graph, var_dependence; taylormap=taylormap))
     end
+    return result, taylormap
 end
 
 
 function taylorexpansion(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
     taylormap = Dict{Int,TaylorSeries{G}}()
-    _taylorexpansion!(graph, taylormap, var_dependence)
-    return taylormap
+    _taylorexpansion!(graph, var_dependence; taylormap=taylormap)
+    return taylormap[graph.id], taylormap
 end
 
 
 function taylorexpansion(graphs::Vector{G}, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
     taylormap = Dict{Int,TaylorSeries{G}}()
+    result = Vector{TaylorSeries{G}}()
     for graph in graphs
-        _taylorexpansion!(graph, taylormap, var_dependence)
+        push!(result, _taylorexpansion!(graph, var_dependence; taylormap=taylormap))
     end
-    return taylormap
+    return taylormap[graph.id], taylormap
 end
 
 """
diff --git a/test/diagram_tree.jl b/test/diagram_tree.jl
index 3495b84c..b94f7a70 100644
--- a/test/diagram_tree.jl
+++ b/test/diagram_tree.jl
@@ -91,7 +91,7 @@ end
 
     DiagTree.uidreset()
     # We only consider the direct part of the above diagram
-    spin = 2.0
+    spin = 1.0
     D = 3
     kF, β, mass2 = 1.919, 0.5, 1.0
     Nk, Nt = 4, 2
@@ -107,13 +107,14 @@ end
     # plot_tree(droot_dg)
 
     DiagTree.eval!(root; eval=(x -> 1.0))
-    @test root.weight ≈ -2 + spin
+    factor = 1 / (2π)^D
+    @test root.weight ≈ (-2 + spin) * factor
 
     DiagTree.eval!(droot_dg; eval=(x -> 1.0))
-    @test root.weight ≈ (-2 + spin) * 2
+    @test droot_dg.weight ≈ (-2 + spin) * 2 * factor
 
     DiagTree.eval!(droot_dv; eval=(x -> 1.0))
-    @test root.weight ≈ (-2 + spin) * 2
+    @test droot_dv.weight ≈ (-2 + spin) * 2 * factor
 
     # #more sophisticated test of the weight evaluation
     varK = rand(D, Nk)
diff --git a/test/taylor.jl b/test/taylor.jl
index ce268535..cbb127d8 100644
--- a/test/taylor.jl
+++ b/test/taylor.jl
@@ -34,10 +34,108 @@ using FeynmanDiagram: Taylor as Taylor
 
     set_variables("x y z", orders=[2, 3, 2])
     for G in [G3, G4, G5, G6]
-        T = taylorexpansion(G)
+        T, taylormap = taylorexpansion(G)
         T_compare = build_derivative_backAD!(G)
         for (order, coeff) in T_compare.coeffs
             @test eval!(coeff) == eval!(taylor_factorial(order) * T.coeffs[order])
         end
     end
+
+end
+
+
+function getdiagram(spin=2.0, D=3, Nk=4, Nt=2)
+    """
+        k1-k3                     k2+k3 
+        |                         | 
+    t1.L ↑     t1.L       t2.L     ↑ t2.L
+        |-------------->----------|
+        |       |  k3+k4   |      |
+        |   v   |          |  v   |
+        |       |    k4    |      |
+        |--------------<----------|
+    t1.L ↑    t1.L        t2.L     ↑ t2.L
+        |                         | 
+        k1                        k2
+    """
+
+    DiagTree.uidreset()
+    # We only consider the direct part of the above diagram
+
+    paraG = DiagParaF64(type=GreenDiag,
+        innerLoopNum=0, totalLoopNum=Nk, loopDim=D,
+        hasTau=true, totalTauNum=Nt)
+    paraV = paraG
+
+    # #construct the propagator table
+    gK = [[0.0, 0.0, 1.0, 1.0], [0.0, 0.0, 0.0, 1.0]]
+    gT = [(1, 2), (2, 1)]
+    g = [Diagram{Float64}(BareGreenId(paraG, k=gK[i], t=gT[i]), name=:G) for i in 1:2]
+
+    vdK = [[0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 1.0, 0.0]]
+    # vdT = [[1, 1], [2, 2]]
+    vd = [Diagram{Float64}(BareInteractionId(paraV, ChargeCharge, k=vdK[i], permu=Di), name=:Vd) for i in 1:2]
+
+    veK = [[1, 0, -1, -1], [0, 1, 0, -1]]
+    # veT = [[1, 1], [2, 2]]
+    ve = [Diagram{Float64}(BareInteractionId(paraV, ChargeCharge, k=veK[i], permu=Ex), name=:Ve) for i in 1:2]
+
+    Id = GenericId(paraV)
+    # contruct the tree
+    ggn = Diagram{Float64}(Id, Prod(), [g[1], g[2]])
+    vdd = Diagram{Float64}(Id, Prod(), [vd[1], vd[2]], factor=spin)
+    vde = Diagram{Float64}(Id, Prod(), [vd[1], ve[2]], factor=-1.0)
+    ved = Diagram{Float64}(Id, Prod(), [ve[1], vd[2]], factor=-1.0)
+    vsum = Diagram{Float64}(Id, Sum(), [vdd, vde, ved])
+    root = Diagram{Float64}(Id, Prod(), [vsum, ggn], factor=1 / (2π)^D, name=:root)
+
+    return root, gK, gT, vdK, veK
+end
+
+@testset "Taylor AD of DiagTree" begin
+
+    DiagTree.uidreset()
+    # We only consider the direct part of the above diagram
+    spin = 0.5
+    D = 3
+    kF, β, mass2 = 1.919, 0.5, 1.0
+    Nk, Nt = 4, 2
+
+    root, gK, gT, vdK, veK = getdiagram(spin, D, Nk, Nt)
+
+    #optimize the diagram
+    DiagTree.optimize!([root,])
+
+    # autodiff
+    droot_dg = DiagTree.derivative([root,], BareGreenId)[1]
+    droot_dv = DiagTree.derivative([root,], BareInteractionId)[1]
+    # plot_tree(droot_dg)
+    factor = 1 / (2π)^D
+    DiagTree.eval!(root; eval=(x -> 1.0))
+    @test root.weight ≈ (-2 + spin) * factor
+
+    DiagTree.eval!(droot_dg; eval=(x -> 1.0))
+    @test droot_dg.weight ≈ (-2 + spin) * 2 * factor
+
+    DiagTree.eval!(droot_dv; eval=(x -> 1.0))
+    @test droot_dv.weight ≈ (-2 + spin) * 2 * factor
+
+    set_variables("x"; orders=[2])
+    g, map = FrontEnds.Graph!(root)
+    var_dependence = FrontEnds.extract_var_dependence(map, BareGreenId, 1)
+    t, taylormap = taylorexpansion(g, var_dependence)
+    order = [0]
+    @test eval!(taylormap[g.id].coeffs[order]) ≈ (-2 + spin) * factor
+
+    order = [1]
+    @test eval!(taylormap[g.id].coeffs[order]) ≈ (-2 + spin) * factor * 2 * taylor_factorial(order)
+
+    var_dependence = FrontEnds.extract_var_dependence(map, BareInteractionId, 1)
+
+    t, taylormap = taylorexpansion(g, var_dependence)
+    order = [0]
+    @test eval!(taylormap[g.id].coeffs[order]) ≈ (-2 + spin) * factor
+
+    order = [1]
+    @test eval!(taylormap[g.id].coeffs[order]) ≈ (-2 + spin) * factor * 2 * taylor_factorial(order)
 end
\ No newline at end of file

From d5e3993ef30830053e9c43485d8e1f9aa42dc602 Mon Sep 17 00:00:00 2001
From: Tao Wang <fsxbhyy@gmail.com>
Date: Fri, 3 Nov 2023 18:00:49 -0400
Subject: [PATCH 3/3] fix typo and api

---
 example/taylor_expansion.jl | 45 ++++---------------------------------
 src/frontend/diagtree.jl    |  4 ++--
 src/utility.jl              | 33 ++++++---------------------
 test/taylor.jl              | 12 +++++-----
 4 files changed, 19 insertions(+), 75 deletions(-)

diff --git a/example/taylor_expansion.jl b/example/taylor_expansion.jl
index 346bc999..4a0aa067 100644
--- a/example/taylor_expansion.jl
+++ b/example/taylor_expansion.jl
@@ -3,9 +3,9 @@ using FeynmanDiagram.Taylor
 using FeynmanDiagram.ComputationalGraphs:
     eval!, forwardAD, node_derivative, backAD, build_all_leaf_derivative
 using FeynmanDiagram.Utility:
-    taylorexpansion!, taylorexpansion, build_derivative_backAD!, count_operation
+    taylorexpansion!, build_derivative_backAD!, count_operation
 
-function banchmark_AD(glist::Vector{T}) where {T<:Graph}
+function benchmark_AD(glist::Vector{T}) where {T<:Graph}
     taylormap = Dict{Int,TaylorSeries{T}}()
     totaloperation = [0, 0]
     taylorlist = Vector{TaylorSeries{T}}()
@@ -20,9 +20,6 @@ function banchmark_AD(glist::Vector{T}) where {T<:Graph}
         t_compare = build_derivative_backAD!(g)
         for (order, coeff) in (t_compare.coeffs)
             @assert (eval!(coeff)) == (eval!(Taylor.taylor_factorial(order) * t.coeffs[order]))
-            # gs = Compilers.to_julia_str([coeff,], name="eval_graph!")
-            # println("$(order)  ", gs, "\n")
-            # print("$(order) $(eval!(coeff)) $(eval!(getderivative(t5,order))) $(count_operation(coeff))\n")
         end
     end
 
@@ -39,49 +36,15 @@ g6 = Graph([])
 G3 = g1
 G4 = 1.0 * g1 * g2
 G5 = 1.0 * (3.0 * G3 + 0.5 * G4)
-#G5 = g5 * g6
-#G6 = (0.5 * g1 * g2 + 0.5 * g2 * g3)
-#G6 = (g1 + g2) * (0.5 * g1 + g3) * g1 #   (0.5 * g1 + g3)
-#G6 = g1 * g2 * g3 * g4 * g5 * g6
 G6 = (1.0 * g1 + 2.0 * g2) * (g1 + g3)
-#G6 = 1.5 * g1*g1 + 0.5 * g2 * 1.5 * g1 +    0.5*g2*g3
+
 using FeynmanDiagram.Taylor:
     TaylorSeries, getcoeff, set_variables
 
-
-# set_variables("x y", order=3)
-# @time T5 = taylorexpansion(G5)
-# print(T5)
 set_variables("x y", orders=[3, 2])
-#set_variables("x y z a", order=[1, 2, 3, 2])
-
-
-#@time taylormap = taylorexpansion(G6)
-
-
-banchmark_AD([G3, G4, G5, G6])
-
-
-# T5 = taylormap[G6.id]
-# #order = [0, 0, 0, 0, 0, 0]
-# #@time print(T5.coeffs[order])
-# print("$(count_operation(T5.coeffs))\n")
-# # for (order, coeff) in (T5.coeffs)
-# #     #gs = Compilers.to_julia_str([coeff,], name="eval_graph!")
-# #     #println("$(order)  ", gs, "\n")
-# #     print("$(order) $(eval!(coeff)) $(eval!(getcoeff(T5,order))) $(coeff.id) $(count_operation(coeff))\n")
-# # end
 
-# print("TaylorSeries $(T5)\n")
+benchmark_AD([G3, G4, G5, G6])
 
-# @time T5_compare = build_derivative_backAD!(G6)
-# print("$(count_operation(T5_compare.coeffs))\n")
-# for (order, coeff) in (T5_compare.coeffs)
-#     @assert (eval!(coeff)) == (eval!(Taylor.taylor_factorial(order) * T5.coeffs[order]))
-#     # gs = Compilers.to_julia_str([coeff,], name="eval_graph!")
-#     # println("$(order)  ", gs, "\n")
-#     # print("$(order) $(eval!(coeff)) $(eval!(getderivative(T5,order))) $(count_operation(coeff))\n")
-# end
 
 
 
diff --git a/src/frontend/diagtree.jl b/src/frontend/diagtree.jl
index 28495675..fa246b65 100644
--- a/src/frontend/diagtree.jl
+++ b/src/frontend/diagtree.jl
@@ -65,9 +65,9 @@ end
 # Arguments:
 - `map::Dict{Int,DiagTree.DiagramId}`:  A dictionary mapping graph ids to DiagramIds. DiagramId stores the diagram information of the corresponding graph. 
 - `ID`:  The particular type of ID that has the given variable dependence. 
-- `numvars`: The number of varibles which the diagram depends on.
+- `numvars`: The number of variables which the diagram depends on.
 """
-function extract_var_dependence(map::Dict{Int,DiagTree.DiagramId}, ::Type{ID}, numvars::Int) where {ID<:PropagatorId}
+function extract_var_dependence(map::Dict{Int,DiagTree.DiagramId}, ::Type{ID}; numvars::Int=1) where {ID<:PropagatorId}
     var_dependence = Dict{Int,Vector{Bool}}()
     for (id, diagID) in map
         if diagID isa ID
diff --git a/src/utility.jl b/src/utility.jl
index 360c8992..70ac437f 100644
--- a/src/utility.jl
+++ b/src/utility.jl
@@ -9,9 +9,9 @@ using ..Taylor
 
 
 # Internal function that performs taylor expansion on a single graph node recursively.
-function _taylorexpansion!(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}(); taylormap::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}()) where {G<:Graph}
+function taylorexpansion!(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}(); taylormap::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}()) where {G<:Graph}
     if haskey(taylormap, graph.id) #If already exist, use taylor series in taylormap.
-        return taylormap[graph.id]
+        return taylormap[graph.id], taylormap
 
     elseif isleaf(graph)
         maxorder = get_orders()
@@ -28,10 +28,10 @@ function _taylorexpansion!(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict
             result.coeffs[o] = coeff
         end
         taylormap[graph.id] = result
-        return result
+        return result, taylormap
     else
-        taylormap[graph.id] = apply(graph.operator, [_taylorexpansion!(sub, var_dependence; taylormap=taylormap) for sub in graph.subgraphs], graph.subgraph_factors)
-        return taylormap[graph.id]
+        taylormap[graph.id] = apply(graph.operator, [taylorexpansion!(sub, var_dependence; taylormap=taylormap)[1] for sub in graph.subgraphs], graph.subgraph_factors)
+        return taylormap[graph.id], taylormap
     end
 end
 
@@ -50,37 +50,18 @@ end
 - `taylormap::Dict{Int,TaylorSeries{G}}` The taylor series correponding to graph nodes. Should map each graph node ID (does not need to belong to input graph) to a taylor series.
     All new taylor series generated by taylor expansion will be added into the expansion.
 """
-function taylorexpansion!(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}(); taylormap::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}()) where {G<:Graph}
-    _taylorexpansion!(graph, var_dependence; taylormap=taylormap)
-    return taylormap[graph.id], taylormap
-end
 
 
 function taylorexpansion!(graphs::Vector{G}, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}(); taylormap::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}()) where {G<:Graph}
     result = Vector{TaylorSeries{G}}()
     for graph in graphs
-        push!(result, _taylorexpansion!(graph, var_dependence; taylormap=taylormap))
+        taylor, _ = taylorexpansion!(graph, var_dependence; taylormap=taylormap)
+        push!(result, taylor)
     end
     return result, taylormap
 end
 
 
-function taylorexpansion(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
-    taylormap = Dict{Int,TaylorSeries{G}}()
-    _taylorexpansion!(graph, var_dependence; taylormap=taylormap)
-    return taylormap[graph.id], taylormap
-end
-
-
-function taylorexpansion(graphs::Vector{G}, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
-    taylormap = Dict{Int,TaylorSeries{G}}()
-    result = Vector{TaylorSeries{G}}()
-    for graph in graphs
-        push!(result, _taylorexpansion!(graph, var_dependence; taylormap=taylormap))
-    end
-    return taylormap[graph.id], taylormap
-end
-
 """
     taylorexpansion_withmap(g::G; coeffmode=true, var::Vector{Int}=collect(1:get_numvars())) where {G<:Graph}
     
diff --git a/test/taylor.jl b/test/taylor.jl
index cbb127d8..3726f53f 100644
--- a/test/taylor.jl
+++ b/test/taylor.jl
@@ -23,7 +23,7 @@ using FeynmanDiagram: Taylor as Taylor
     using FeynmanDiagram.ComputationalGraphs:
         eval!, forwardAD, node_derivative, backAD, build_all_leaf_derivative, count_operation
     using FeynmanDiagram.Utility:
-        taylorexpansion, build_derivative_backAD!
+        taylorexpansion!, build_derivative_backAD!
     g1 = Graph([])
     g2 = Graph([])
     g3 = Graph([], factor=2.0)
@@ -34,7 +34,7 @@ using FeynmanDiagram: Taylor as Taylor
 
     set_variables("x y z", orders=[2, 3, 2])
     for G in [G3, G4, G5, G6]
-        T, taylormap = taylorexpansion(G)
+        T, taylormap = taylorexpansion!(G)
         T_compare = build_derivative_backAD!(G)
         for (order, coeff) in T_compare.coeffs
             @test eval!(coeff) == eval!(taylor_factorial(order) * T.coeffs[order])
@@ -122,17 +122,17 @@ end
 
     set_variables("x"; orders=[2])
     g, map = FrontEnds.Graph!(root)
-    var_dependence = FrontEnds.extract_var_dependence(map, BareGreenId, 1)
-    t, taylormap = taylorexpansion(g, var_dependence)
+    var_dependence = FrontEnds.extract_var_dependence(map, BareGreenId)
+    t, taylormap = taylorexpansion!(g, var_dependence)
     order = [0]
     @test eval!(taylormap[g.id].coeffs[order]) ≈ (-2 + spin) * factor
 
     order = [1]
     @test eval!(taylormap[g.id].coeffs[order]) ≈ (-2 + spin) * factor * 2 * taylor_factorial(order)
 
-    var_dependence = FrontEnds.extract_var_dependence(map, BareInteractionId, 1)
+    var_dependence = FrontEnds.extract_var_dependence(map, BareInteractionId)
 
-    t, taylormap = taylorexpansion(g, var_dependence)
+    t, taylormap = taylorexpansion!(g, var_dependence)
     order = [0]
     @test eval!(taylormap[g.id].coeffs[order]) ≈ (-2 + spin) * factor