Skip to content

Commit

Permalink
taylor wip
Browse files Browse the repository at this point in the history
  • Loading branch information
@fsxbhyy
fsxbhyy committed Oct 22, 2023
1 parent 9f6419d commit 5ac9ac5
Show file tree
Hide file tree
Showing 9 changed files with 294 additions and 205 deletions.
48 changes: 29 additions & 19 deletions example/taylor_expansion.jl
View file
Original file line number Diff line number Diff line change
@@ -1,37 +1,47 @@
using FeynmanDiagram
using FeynmanDiagram.ComputationalGraphs:
eval!, forwardAD, node_derivative, backAD, build_all_leaf_derivative, Leaves, taylorexpansion!, count_operation, build_derivative_backAD!
eval!, forwardAD, node_derivative, backAD, build_all_leaf_derivative, count_operation
using FeynmanDiagram.Utility:
taylorexpansion!, build_derivative_backAD!
g1 = Graph([])
g2 = Graph([])
g3 = Graph([]) #, factor=2.0)
g4 = Graph([])
g5 = Graph([])
g6 = Graph([])
G3 = g1
G4 = 1.0 * g1 * g1
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 = (1.0 * g1 + 0.5 * g2) * (1.5 * g1 + 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=6)

# set_variables("x y", order=3)
# @time T5 = taylorexpansion!(G5)
# print(T5)
set_variables("x", order=3)
@time T5 = taylorexpansion!(G6)
order = [0, 3]
@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)) $(count_operation(coeff))\n")
# end
#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

# @time T5_compare = build_derivative_backAD!(G6)
# print("$(count_operation(T5_compare.coeffs))\n")
# for (order, coeff) in (T5_compare.coeffs)
# #gs = Compilers.to_julia_str([coeff,], name="eval_graph!")
# #println("$(order) ", gs, "\n")
# print("$(order) $(eval!(coeff)) $(eval!(T5.coeffs[order])) $(count_operation(coeff))\n")
# end
@time T5_compare = build_derivative_backAD!(G6)
print("$(count_operation(T5_compare.coeffs))\n")
for (order, coeff) in (T5_compare.coeffs)
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



Expand Down
16 changes: 11 additions & 5 deletions src/FeynmanDiagram.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -98,10 +98,7 @@ export fermionic_annihilation, fermionic_creation, majorana
export bosonic_annihilation, bosonic_creation, real_classic
export correlator_order, normal_order

include("TaylorSeries/TaylorSeries.jl")
using .Taylor
export Taylor
export TaylorSeries, set_variables



include("computational_graph/ComputationalGraph.jl")
Expand All @@ -128,6 +125,12 @@ export relabel, standardize_labels, replace_subgraph, merge_linear_combination,

export optimize!, optimize, merge_all_chains!, merge_all_linear_combinations!, remove_duplicated_leaves!

include("TaylorSeries/TaylorSeries.jl")
using .Taylor
export Taylor
export TaylorSeries, set_variables, taylor_factorial, getcoeff, getderivative


include("backend/compiler.jl")
using .Compilers
export Compilers
Expand Down Expand Up @@ -175,7 +178,10 @@ export addpropagator!, addnode!
export setroot!, addroot!
export evalNaive, showTree


include("utility.jl")
using .Utility
export Utility
export taylorexpansion!, build_derivative_backAD!
##################### precompile #######################
# precompile as the final step of the module definition:
if ccall(:jl_generating_output, Cint, ()) == 1 # if we're precompiling the package
Expand Down
4 changes: 2 additions & 2 deletions src/TaylorSeries/TaylorSeries.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ see also [`HomogeneousPolynomial`](@ref).
"""
module Taylor


using ..ComputationalGraphs
#using Markdown


Expand All @@ -26,7 +26,7 @@ export get_order, get_numvars,
get_variable_names, get_variable_symbols,
# jacobian, hessian, jacobian!, hessian!,
displayBigO, use_show_default,
getcoeff
getcoeff, getderivative, taylor_factorial
# function __init__()
# @static if !isdefined(Base, :get_extension)
# @require IntervalArithmetic = "d1acc4aa-44c8-5952-acd4-ba5d80a2a253" begin
Expand Down
20 changes: 16 additions & 4 deletions src/TaylorSeries/arithmetic.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -145,11 +145,13 @@ function Base.:*(g1::TaylorSeries{T}, g2::TaylorSeries{T}) where {T}
for (order2, coeff2) in g2.coeffs
order = order1 + order2
if sum(order) <= _params_Taylor_.order
combination_coeff = taylor_binomial(order1, order2)
#combination_coeff = taylor_binomial(order1, order2)
if haskey(g.coeffs, order)
g.coeffs[order] += combination_coeff * coeff1 * coeff2
#g.coeffs[order] += combination_coeff * coeff1 * coeff2
g.coeffs[order] += coeff1 * coeff2
else
g.coeffs[order] = combination_coeff * coeff1 * coeff2
#g.coeffs[order] = combination_coeff * coeff1 * coeff2
g.coeffs[order] = coeff1 * coeff2
end
end
end
Expand All @@ -164,7 +166,17 @@ end

function getcoeff(g::TaylorSeries, order::Array{Int,1})
if haskey(g.coeffs, order)
return 1 / taylor_factorial(order) * g.coeffs[order]
return g.coeffs[order]
#return 1 / taylor_factorial(order) * g.coeffs[order]
else
return nothing
end
end

function getderivative(g::TaylorSeries, order::Array{Int,1})
if haskey(g.coeffs, order)
#return g.coeffs[order]
return taylor_factorial(order) * g.coeffs[order]
else
return nothing
end
Expand Down
34 changes: 33 additions & 1 deletion src/TaylorSeries/print.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -164,7 +164,39 @@ function homogPol2str(a::TaylorSeries{T}) where {T<:Number}
end
@inbounds c = coeff
iszero(c) && continue
cadena = numbr2str(c / factor, ifirst)
#cadena = numbr2str(c / factor, ifirst)
cadena = numbr2str(c, ifirst)
strout = string(strout, cadena, monom, space)
ifirst = false
end
return strout[1:prevind(strout, end)]
end

function homogPol2str(a::TaylorSeries{T}) where {T<:AbstractGraph}
numVars = get_numvars()
#z = zero(a.coeffs[1])
space = string(" ")
strout::String = space
ifirst = true
for (order, coeff) in a.coeffs
monom::String = string("")
factor = 1
for ivar = 1:numVars
powivar = order[ivar]
if powivar == 1
monom = string(monom, name_taylorvar(ivar))
elseif powivar > 1
monom = string(monom, name_taylorvar(ivar), superscriptify(powivar))
factor *= factorial(powivar)
end
end
c = coeff
if ifirst
cadena = "g$(coeff.id)"
ifirst = false
else
cadena = "+ g$(coeff.id)"
end
strout = string(strout, cadena, monom, space)
ifirst = false
end
Expand Down
2 changes: 1 addition & 1 deletion src/computational_graph/ComputationalGraph.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ module ComputationalGraphs
using AbstractTrees
using StaticArrays
using Printf, PyCall, DataFrames
using ..Taylor
#using ..Taylor
macro todo()
return :(error("Not yet implemented!"))
end
Expand Down
172 changes: 0 additions & 172 deletions src/computational_graph/operation.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -325,109 +325,8 @@ function build_all_leaf_derivative(diag::Graph{F,W}, maxorder=Inf) where {F,W}
return result, chainrule_map
end

function build_derivative_backAD!(g::G, leaftaylor::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}()) where {G<:Graph}
chainrule_map_leaf = Dict{Int,Dict{Int,G}}()
for leaf in Leaves(g)
if !haskey(leaftaylor, leaf.id)
leaftaylor[leaf.id] = graph_to_taylor_withmap!(leaf; chainrule_map_leaf=chainrule_map_leaf)
end
end
leafAD, chainrule_map = build_all_leaf_derivative(g)
current_func = Dict(zeros(Int, get_numvars()) => g)

result = TaylorSeries{G}()
result.coeffs[zeros(Int, get_numvars())] = g
for i in 1:get_order()
new_func = Dict{Array{Int,1},G}()
for (order, func) in current_func
for idx in 1:get_numvars()
ordernew = copy(order)
ordernew[idx] += 1
if !haskey(result.coeffs, ordernew)
funcAD = forwardAD_taylor(func, idx, chainrule_map, chainrule_map_leaf, leaftaylor)
if !isnothing(funcAD)
new_func[ordernew] = funcAD
result.coeffs[ordernew] = funcAD
end
end
end
end
current_func = new_func
end
return result
end

function forwardAD_taylor(g::G, varidx::Int, chainrule_map::Dict{Int,Array{G,1}}, chainrule_map_leaf::Dict{Int,Dict{Int,G}}, leaftaylor::Dict{Int,TaylorSeries{G}}) where {G<:Graph}
if haskey(chainrule_map, g.id)
return chainrule!(varidx, chainrule_map[g.id], leaftaylor)
elseif haskey(chainrule_map_leaf, g.id)
#if haskey(chainrule_map_leaf, g.id)
map = chainrule_map_leaf[g.id]
if haskey(map, varidx)
return map[varidx]
else
return nothing
end
elseif g.operator == Sum
children = Array{G,1}()
for graph in g.subgraphs
dgraph = forwardAD_taylor(graph, varidx, chainrule_map, chainrule_map_leaf, leaftaylor)
if !isnothing(dgraph)
push!(children, dgraph)
end
end
if isempty(children)
return nothing
else
return linear_combination(children, g.subgraph_factors)
end
elseif g.operator == 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))
end
end
if isempty(children)
return nothing
else
return linear_combination(children)
end
elseif g.operator <: Power

dgraph = forwardAD_taylor(g.subgraphs[1], varidx, chainrule_map, chainrule_map_leaf, leaftaylor)
if isnothing(dgraph)
return nothing
else
power = eltype(g.operator)
if power == 1
return dgraph
else
return dgraph * Graph(g.subgraphs; subgraph_factors=power * g.subgraph_factors, operator=decrement_power(g.operator))
end
end
end
end

function chainrule!(varidx::Int, dg::Array{G,1}, leaftaylor::Dict{Int,TaylorSeries{G}}) where {G<:Graph}
children = Array{G,1}()
order = zeros(Int, get_numvars())
order[varidx] += 1
for i in 1:length(dg)÷2
taylor = leaftaylor[dg[2*i-1].id]
if haskey(taylor.coeffs, order)
coeff = taylor.coeffs[order]
push!(children, coeff * dg[2*i])
end
end
if isempty(children)
return nothing
else
return linear_combination(children)
end
end


function insert_dualDict!(dict_kv::Dict{Tk,Tv}, dict_vk::Dict{Tv,Tk}, key::Tk, value::Tv) where {Tk,Tv}
Expand Down Expand Up @@ -644,74 +543,3 @@ function build_derivative_graph(graph::G, orders::NTuple{N,Int};
return build_derivative_graph([graph], orders; nodes_id=nodes_id)
end

function graph_to_taylor(graph::G, varidx::Array{Int,1}=collect(1:get_numvars())) where {G<:Graph}
@assert isleaf(graph)
maxorder = get_order()
ordtuple = ((idx in varidx) ? (0:maxorder) : (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)
if sum(o) <= get_order()
coeff = Graph([]; operator=Sum(), factor=graph.factor)
result.coeffs[o] = coeff
end
end
return result
end


function graph_to_taylor_withmap!(g::G; varidx::Array{Int,1}=collect(1:get_numvars()), chainrule_map_leaf::Dict{Int,Dict{Int,G}}=Dict{Int,Dict{Int,G}}()) where {G<:Graph}
@assert isleaf(g)
maxorder = get_order()
current_func = Dict(zeros(Int, get_numvars()) => g)
result = TaylorSeries{G}()
result.coeffs[zeros(Int, get_numvars())] = g

for i in 1:get_order()
new_func = Dict{Array{Int,1},G}()
for (order, func) in current_func
if !haskey(chainrule_map_leaf, func.id)
chainrule_map_leaf[func.id] = Dict{Int,G}()
end
for idx in varidx
ordernew = copy(order)
ordernew[idx] += 1
if !haskey(result.coeffs, ordernew)
funcAD = Graph([]; operator=Sum(), factor=g.factor)
new_func[ordernew] = funcAD
result.coeffs[ordernew] = funcAD
chainrule_map_leaf[func.id][idx] = funcAD
else
chainrule_map_leaf[func.id][idx] = result.coeffs[ordernew]
end
end
end
current_func = new_func
end

return result
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]

function taylorexpansion!(graph::G, taylormap::Dict{Int,T}=Dict{Int,TaylorSeries{G}}()) where {G<:Graph,T<:TaylorSeries}
if isempty(taylormap)
for g in Leaves(graph)
if !haskey(taylormap, g.id)
#taylormap[g.id] = graph_to_taylor(g)
taylormap[g.id] = graph_to_taylor_withmap!(g)
end
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]
end
Loading

0 comments on commit 5ac9ac5

Please sign in to comment.