Merge pull request #154 from numericalEFT/tao_AD

API clean up

example/taylor_expansion.jl

@@ -1,9 +1,32 @@
 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!, build_derivative_backAD!, count_operation
+
+function benchmark_AD(glist::Vector{T}) where {T<:Graph}
+    taylormap = Dict{Int,TaylorSeries{T}}()
+    totaloperation = [0, 0]
+    taylorlist = Vector{TaylorSeries{T}}()
+    for g in glist
+        @time t, taylormap = taylorexpansion!(g; taylormap=taylormap)
+
+
+        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]))
+        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,40 +36,15 @@ g6 = Graph([])
 G3 = g1
 G4 = 1.0 * g1 * g2
 G5 = 1.0 * (3.0 * G3 + 0.5 * G4)
-#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", orders=[3, 2])
+
+benchmark_AD([G3, G4, G5, G6])
 
-# set_variables("x y", order=3)
-# @time T5 = taylorexpansion(G5)
-# print(T5)
-set_variables("x y", order=[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
 
 
 

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

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 variables which the diagram depends on.
+"""
+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
+            var_dependence[id] = [true for _ in 1:numvars]
+        else
+            var_dependence[id] = [false for _ in 1:numvars]
+        end
+    end
+    return var_dependence
 end

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. 
+# 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}
+    if haskey(taylormap, graph.id) #If already exist, use taylor series in taylormap.
+        return taylormap[graph.id], taylormap
 
-#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.
-"""
-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,27 +27,41 @@ function taylorexpansion(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{I
             coeff = Graph([]; operator=Sum(), factor=graph.factor)
             result.coeffs[o] = coeff
         end
-        return result
+        taylormap[graph.id] = result
+        return result, taylormap
     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, var_dependence; taylormap=taylormap)[1] for sub in graph.subgraphs], graph.subgraph_factors)
+        return taylormap[graph.id], taylormap
+    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!(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
+        taylor, _ = taylorexpansion!(graph, var_dependence; taylormap=taylormap)
+        push!(result, taylor)
     end
+    return result, taylormap
 end
 
+
 """
     taylorexpansion_withmap(g::G; coeffmode=true, var::Vector{Int}=collect(1:get_numvars())) where {G<:Graph}
     
@@ -225,4 +233,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

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)

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,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)
+    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)
+
+    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