Merge pull request #155 from numericalEFT/tao_AD

Refactor taylor expansion

example/benchmark.jl

@@ -0,0 +1,41 @@
+using FeynmanDiagram
+using FeynmanDiagram.Taylor
+using FeynmanDiagram.ComputationalGraphs:
+    eval!, Leaves
+using FeynmanDiagram.Utility:
+    taylorexpansion!, count_operation
+
+
+function assign_leaves(g::FeynmanGraph, taylormap)
+    leafmap = Dict{Int,Int}()
+    leafvec = Vector{Float64}()
+    idx = 0
+    for leaf in Leaves(g)
+        taylor = taylormap[leaf.id]
+        for (order, coeff) in taylor.coeffs
+            idx += 1
+            push!(leafvec, 1.0 / taylor_factorial(order))
+            leafmap[coeff.id] = idx
+            print("assign $(order) $(coeff.id)  $(taylor_factorial(order)) $(leafvec[idx])\n")
+        end
+    end
+    return leafmap, leafvec
+end
+
+dict_g, fl, bl, leafmap = diagdictGV(:sigma, [(2, 0, 0), (2, 0, 1), (2, 0, 2), (2, 1, 0), (2, 1, 1), (2, 2, 0), (2, 1, 2), (2, 2, 2)], 3)
+
+g = dict_g[(2, 0, 0)]
+
+set_variables("x y", orders=[2, 2])
+propagator_var = ([true, false], [false, true]) # Specify variable dependence of fermi (first element) and bose (second element) particles.
+t, taylormap = taylorexpansion!(g[1][1], propagator_var, (fl, bl))
+
+for (order, graph) in dict_g
+    if graph[2][1] == g[2][1]
+        idx = 1
+    else
+        idx = 2
+    end
+    print("$(order) $(eval!(graph[1][idx])) $(eval!(t.coeffs[[order[2],order[3]]]))\n")
+end
+

src/FeynmanDiagram.jl

@@ -147,6 +147,7 @@ export PropagatorId, BareGreenId, BareInteractionId
 export BareGreenNId, BareHoppingId, GreenNId, ConnectedGreenNId
 export uidreset, toDataFrame, mergeby, plot_tree
 
+
 include("parquet_builder/parquet.jl")
 using .Parquet
 export Parquet
@@ -166,10 +167,7 @@ export addpropagator!, addnode!
 export setroot!, addroot!
 export evalNaive, showTree
 
-include("utility.jl")
-using .Utility
-export Utility
-export taylorexpansion!
+
 
 include("frontend/frontends.jl")
 using .FrontEnds
@@ -180,6 +178,11 @@ include("frontend/GV.jl")
 using .GV
 export GV
 export diagdictGV, leafstates
+
+include("utility.jl")
+using .Utility
+export Utility
+export taylorexpansion!
 # export read_onediagram, read_diagrams
 
 ##################### precompile #######################

src/computational_graph/eval.jl

@@ -30,12 +30,17 @@ function eval!(g::Graph{F,W}, leafmap::Dict{Int,Int}=Dict{Int,Int}(), leaf::Vect
     return result
 end
 
+
 function eval!(g::FeynmanGraph{F,W}, leafmap::Dict{Int,Int}=Dict{Int,Int}(), leaf::Vector{W}=Vector{W}()) where {F,W}
     result = nothing
 
     for node in PostOrderDFS(g)
         if isleaf(node)
-            node.weight = leaf[leafmap[node.id]]
+            if isempty(leafmap)
+                node.weight = 1.0
+            else
+                node.weight = leaf[leafmap[node.id]]
+            end
         else
             node.weight = apply(node.operator, node.subgraphs, node.subgraph_factors)
         end

src/diagram_tree/DiagTree.jl

@@ -41,5 +41,6 @@ export DiagramId, GenericId, Ver4Id, Ver3Id, GreenId, SigmaId, PolarId
 export PropagatorId, BareGreenId, BareInteractionId
 export BareGreenNId, BareHoppingId, GreenNId, ConnectedGreenNId
 export uidreset, toDataFrame, mergeby, plot_tree
+export isleaf
 
 end

src/utility.jl

@@ -1,56 +1,181 @@
 module Utility
 using ..ComputationalGraphs
-using ..ComputationalGraphs: Sum, Prod, Power, decrement_power
-using ..ComputationalGraphs: build_all_leaf_derivative, eval!
+#using ..ComputationalGraphs: Sum, Prod, Power, decrement_power
+using ..ComputationalGraphs: decrement_power
+using ..ComputationalGraphs: build_all_leaf_derivative, eval!, isfermionic
 import ..ComputationalGraphs: count_operation
 using ..ComputationalGraphs.AbstractTrees
-
+using ..DiagTree
+using ..DiagTree: Diagram, PropagatorId, BareGreenId, BareInteractionId
+using ..FrontEnds: LabelProduct
 using ..Taylor
 
+@inline apply(::Type{ComputationalGraphs.Sum}, diags::Vector{T}, factors::Vector{F}) where {T<:TaylorSeries,F<:Number} = sum(d * f for (d, f) in zip(diags, factors))
+@inline apply(::Type{ComputationalGraphs.Prod}, diags::Vector{T}, factors::Vector{F}) where {T<:TaylorSeries,F<:Number} = prod(d * f for (d, f) in zip(diags, factors))
+@inline apply(::Type{ComputationalGraphs.Power{N}}, diags::Vector{T}, factors::Vector{F}) where {N,T<:TaylorSeries,F<:Number} = (diags[1])^N * factors[1]
+
+@inline apply(::Type{DiagTree.Sum}, diags::Vector{T}, factors::Vector{F}) where {T<:TaylorSeries,F<:Number} = sum(d * f for (d, f) in zip(diags, factors))
+@inline apply(::Type{DiagTree.Prod}, diags::Vector{T}, factors::Vector{F}) where {T<:TaylorSeries,F<:Number} = prod(d * f for (d, f) in zip(diags, factors))
+
 
-# 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}
+    
+    Return a taylor series of graph g, together with a map of between nodes of g and correponding taylor series.
+# Arguments:
+- `graph`  Target graph 
+- `var_dependence::Dict{Int,Vector{Bool}}` A dictionary that specifies the variable dependence of target graph leaves. Should map the id of each leaf to a Bool vector. 
+    The length of the vector should be the same as number of variables.
+- `taylormap::Dict{Int,TaylorSeries}` A dicitonary that maps id of each node of target graph to its correponding taylor series.
+"""
 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], taylormap
 
     elseif isleaf(graph)
-        maxorder = get_orders()
         if haskey(var_dependence, graph.id)
             var = var_dependence[graph.id]
         else
-            var = fill(true, get_numvars()) #if dependence not provided, assume the graph depends on all variables
+            var = fill(false, get_numvars()) #if dependence not provided, assume the graph depends on no variables
         end
         ordtuple = ((var[idx]) ? (0:get_orders(idx)) : (0:0) for idx in 1:get_numvars())
         result = TaylorSeries{G}()
         for order in collect(Iterators.product(ordtuple...)) #varidx specifies the variables graph depends on. Iterate over all taylor coefficients of those variables.
             o = collect(order)
-            coeff = Graph([]; operator=Sum(), factor=graph.factor)
-            result.coeffs[o] = coeff
+            if sum(o) == 0      # For a graph the zero order taylor coefficient is just itself.
+                result.coeffs[o] = graph
+            else
+                coeff = Graph([]; operator=ComputationalGraphs.Sum(), factor=graph.factor)
+                result.coeffs[o] = coeff
+            end
         end
         taylormap[graph.id] = result
         return result, taylormap
     else
-        taylormap[graph.id] = apply(graph.operator, [taylorexpansion!(sub, var_dependence; taylormap=taylormap)[1] for sub in graph.subgraphs], graph.subgraph_factors)
+        taylormap[graph.id] = graph.factor * 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
 
-
 """
-    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::FeynmanGraph{F,W}, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}(); taylormap::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}()) where {F,W}
+    
+    Return a taylor series of FeynmanGraph g, together with a map of between nodes of g and correponding taylor series.
+# Arguments:
+- `graph`  Target FeynmanGraph 
+- `var_dependence::Dict{Int,Vector{Bool}}` A dictionary that specifies the variable dependence of target graph leaves. Should map the id of each leaf to a Bool vector. 
+    The length of the vector should be the same as number of variables.
+- `taylormap::Dict{Int,TaylorSeries}` A dicitonary that maps id of each node of target graph to its correponding taylor series.
+"""
+function taylorexpansion!(graph::FeynmanGraph{F,W}, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}(); taylormap::Dict{Int,TaylorSeries{Graph{F,W}}}=Dict{Int,TaylorSeries{Graph{F,W}}}()) where {F,W}
+    if haskey(taylormap, graph.id) #If already exist, use taylor series in taylormap.
+        return taylormap[graph.id], taylormap
+
+    elseif isleaf(graph)
+        if haskey(var_dependence, graph.id)
+            var = var_dependence[graph.id]
+        else
+            var = fill(false, get_numvars()) #if dependence not provided, assume the graph depends on no variables
+        end
+        ordtuple = ((var[idx]) ? (0:get_orders(idx)) : (0:0) for idx in 1:get_numvars())
+        result = TaylorSeries{Graph{F,W}}()
+        for order in collect(Iterators.product(ordtuple...)) #varidx specifies the variables graph depends on. Iterate over all taylor coefficients of those variables.
+            o = collect(order)
+            coeff = Graph([]; operator=ComputationalGraphs.Sum(), factor=graph.factor)
+            result.coeffs[o] = coeff
+        end
+        taylormap[graph.id] = result
+        return result, taylormap
+    else
+        taylormap[graph.id] = graph.factor * 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
 
-    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.
+"""
+    function taylorexpansion!(graph::Diagram{W}, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}(); taylormap::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}()) where {W}
+    
+    Return a taylor series of Diagram g, together with a map of between nodes of g and correponding taylor series.
+# Arguments:
+- `graph`  Target diagram 
+- `var_dependence::Dict{Int,Vector{Bool}}` A dictionary that specifies the variable dependence of target diagram leaves. Should map the id of each leaf to a Bool vector. 
+    The length of the vector should be the same as number of variables.
+- `taylormap::Dict{Int,TaylorSeries}` A dicitonary that maps id of each node of target diagram to its correponding taylor series.
+"""
+function taylorexpansion!(graph::Diagram{W}, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}(); taylormap::Dict{Int,TaylorSeries{Graph{W,W}}}=Dict{Int,TaylorSeries{Graph{W,W}}}()) where {W}
+    if haskey(taylormap, graph.hash) #If already exist, use taylor series in taylormap.
+        return taylormap[graph.hash], taylormap
 
-#Arguments
+    elseif isempty(graph.subdiagram)
+        if haskey(var_dependence, graph.hash)
+            var = var_dependence[graph.hash]
+        else
+            var = fill(false, get_numvars()) #if dependence not provhashed, assume the graph depends on no variables
+        end
+        ordtuple = ((var[idx]) ? (0:get_orders(idx)) : (0:0) for idx in 1:get_numvars())
+        result = TaylorSeries{Graph{W,W}}()
+        for order in collect(Iterators.product(ordtuple...)) #varidx specifies the variables graph depends on. Iterate over all taylor coefficients of those variables.
+            o = collect(order)
+            coeff = Graph([]; operator=ComputationalGraphs.Sum(), factor=graph.factor)
+            result.coeffs[o] = coeff
+        end
+        taylormap[graph.hash] = result
+        return result, taylormap
+    else
+        taylormap[graph.hash] = graph.factor * apply(typeof(graph.operator), [taylorexpansion!(sub, var_dependence; taylormap=taylormap)[1] for sub in graph.subdiagram], ones(W, length(graph.subdiagram)))
+        return taylormap[graph.hash], taylormap
+    end
+end
 
-- `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::FeynmanGraph{F,W}, propagator_var::Tuple{Vector{Bool},Vector{Bool}}, label::Tuple{LabelProduct,LabelProduct}; taylormap::Dict{Int,TaylorSeries{Graph{F,W}}}=Dict{Int,TaylorSeries{Graph{F,W}}}()) where {F,W}
+    
+    Return a taylor series of FeynmanGraph g, together with a map of between nodes of g and correponding taylor series. In this set up, the leaves that are the same type of propagators (fermi or bose) depend on the same set of variables, 
+    whereas other types of Feynman diagrams (such as vertices) depends on no variables that need to be differentiated (for AD purpose, they are just constants).
+# Arguments:
+- `graph`  Target FeynmanGraph
+- `propagator_var::Tuple{Vector{Bool},Vector{Bool}}` A Tuple that specifies the variable dependence of fermi (first element) and bose propagator (second element).
+    The dependence is given by a vector of the length same as the number of variables.
+- `label::Tuple{LabelProduct,LabelProduct}` A Tuple fermi (first element) and bose LabelProduct (second element).
+- `taylormap::Dict{Int,TaylorSeries}` A dicitonary that maps id of each node of target diagram to its correponding taylor series.
+"""
+function taylorexpansion!(graph::FeynmanGraph{F,W}, propagator_var::Tuple{Vector{Bool},Vector{Bool}}, label::Tuple{LabelProduct,LabelProduct}; taylormap::Dict{Int,TaylorSeries{Graph{F,W}}}=Dict{Int,TaylorSeries{Graph{F,W}}}()) where {F,W}
+    var_dependence = Dict{Int,Vector{Bool}}()
+    for leaf in Leaves(graph)
+        if ComputationalGraphs.diagram_type(leaf) == ComputationalGraphs.Propagator
+            In = leaf.properties.vertices[2][1].label
+            if isfermionic(leaf.properties.vertices[1])
+                if label[1][In][2] >= 0 #For fake propagator, this label is smaller than zero, and those propagators should not be differentiated.
+                    var_dependence[leaf.id] = [propagator_var[1][idx] ? true : false for idx in 1:get_numvars()]
+                end
+            else
+                var_dependence[leaf.id] = [propagator_var[2][idx] ? true : false for idx in 1:get_numvars()]
+            end
+        end
+    end
+    return taylorexpansion!(graph, var_dependence; taylormap=taylormap)
+end
 
+"""
+    function taylorexpansion!(graph::Diagram{W}, propagator_var::Dict{DataType,Vector{Bool}}; taylormap::Dict{Int,TaylorSeries{Graph{W,W}}}=Dict{Int,TaylorSeries{Graph{W,W}}}()) where {W}
+    
+    Return a taylor series of Diagram g, together with a map of between nodes of g and correponding taylor series. In this set up, the leaves that are the same type of diagrams (such as Green functions) depend on the same set of variables.
+    
+# Arguments:
+- `graph`  Target Diagram
+- `propagator_var::Dict{DataType,Vector{Bool}}` A dictionary that specifies the variable dependence of different types of diagrams. Should be a map between DataTypes in DiagramID and Bool vectors.
+    The dependence is given by a vector of the length same as the number of variables.
+- `taylormap::Dict{Int,TaylorSeries}` A dicitonary that maps id of each node of target diagram to its correponding taylor series.
+"""
+function taylorexpansion!(graph::Diagram{W}, propagator_var::Dict{DataType,Vector{Bool}}; taylormap::Dict{Int,TaylorSeries{Graph{W,W}}}=Dict{Int,TaylorSeries{Graph{W,W}}}()) where {W}
+    var_dependence = Dict{Int,Vector{Bool}}()
+    for leaf in Leaves(graph)
+        if haskey(propagator_var, typeof(leaf.id))
+            var_dependence[leaf.hash] = [propagator_var[typeof(leaf.id)][idx] ? true : false for idx in 1:get_numvars()]
+        end
+    end
+    return taylorexpansion!(graph, var_dependence; taylormap=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}}()
@@ -94,10 +219,10 @@ function taylorexpansion_withmap(g::G; coeffmode=true, var::Vector{Bool}=fill(tr
                     if ordernew[idx] <= get_orders(idx)
                         if !haskey(result.coeffs, ordernew)
                             if coeffmode
-                                funcAD = Graph([]; operator=Sum(), factor=g.factor)
+                                funcAD = Graph([]; operator=ComputationalGraphs.Sum(), factor=g.factor)
                             else
-                                #funcAD = taylor_factorial(ordernew) * Graph([]; operator=Sum(), factor=g.factor)
-                                funcAD = Graph([]; operator=Sum(), factor=taylor_factorial(ordernew) * g.factor)
+                                #funcAD = taylor_factorial(ordernew) * Graph([]; operator=ComputationalGraphs.Sum(), factor=g.factor)
+                                funcAD = Graph([]; operator=ComputationalGraphs.Sum(), factor=taylor_factorial(ordernew) * g.factor)
                             end
                             new_func[ordernew] = funcAD
                             result.coeffs[ordernew] = funcAD
@@ -115,10 +240,6 @@ function taylorexpansion_withmap(g::G; coeffmode=true, var::Vector{Bool}=fill(tr
     return result, chainrule_map_leaf
 end
 
-@inline apply(::Type{Sum}, diags::Vector{T}, factors::Vector{F}) where {T<:TaylorSeries,F<:Number} = sum(d * f for (d, f) in zip(diags, factors))
-@inline apply(::Type{Prod}, diags::Vector{T}, factors::Vector{F}) where {T<:TaylorSeries,F<:Number} = prod(d * f for (d, f) in zip(diags, factors))
-@inline apply(::Type{Power{N}}, diags::Vector{T}, factors::Vector{F}) where {N,T<:TaylorSeries,F<:Number} = (diags[1])^N * factors[1]
-
 
 
 #Functions below generate high order derivatives with naive nested forward AD. This part would be significantly refactored later with 
@@ -157,7 +278,7 @@ function build_derivative_backAD!(g::G, leaftaylor::Dict{Int,TaylorSeries{G}}=Di
         end
         current_func = new_func
     end
-    return result
+    return result, leaftaylor
 end
 
 
@@ -172,7 +293,7 @@ function forwardAD_taylor(g::G, varidx::Int, chainrule_map::Dict{Int,Array{G,1}}
         else
             return nothing
         end
-    elseif g.operator == Sum
+    elseif g.operator == ComputationalGraphs.Sum
         children = Array{G,1}()
         for graph in g.subgraphs
             dgraph = forwardAD_taylor(graph, varidx, chainrule_map, chainrule_map_leaf, leaftaylor)
@@ -185,21 +306,21 @@ function forwardAD_taylor(g::G, varidx::Int, chainrule_map::Dict{Int,Array{G,1}}
         else
             return linear_combination(children, g.subgraph_factors)
         end
-    elseif g.operator == Prod
+    elseif g.operator == ComputationalGraphs.Prod
         children = Array{G,1}()
         for (i, graph) in enumerate(g.subgraphs)
             dgraph = forwardAD_taylor(graph, varidx, chainrule_map, chainrule_map_leaf, leaftaylor)
             if !isnothing(dgraph)
                 subgraphs = [j == i ? dgraph : subg for (j, subg) in enumerate(g.subgraphs)]
-                push!(children, Graph(subgraphs; operator=Prod(), subgraph_factors=g.subgraph_factors))
+                push!(children, Graph(subgraphs; operator=ComputationalGraphs.Prod(), subgraph_factors=g.subgraph_factors))
             end
         end
         if isempty(children)
             return nothing
         else
             return linear_combination(children)
         end
-    elseif g.operator <: Power
+    elseif g.operator <: ComputationalGraphs.Power
 
         dgraph = forwardAD_taylor(g.subgraphs[1], varidx, chainrule_map, chainrule_map_leaf, leaftaylor)
         if isnothing(dgraph)