From a8a18ef36c12575a82b7fb0d255782ddde15315a Mon Sep 17 00:00:00 2001 From: Tao Wang Date: Sat, 28 Oct 2023 09:41:50 -0400 Subject: [PATCH] taylor fix api --- example/taylor_expansion.jl | 11 +-- src/TaylorSeries/arithmetic.jl | 47 +++++++++--- src/TaylorSeries/constructors.jl | 10 +-- src/TaylorSeries/parameter.jl | 11 +-- src/TaylorSeries/print.jl | 5 +- src/utility.jl | 123 +++++++++++++++---------------- test/taylor.jl | 8 +- 7 files changed, 123 insertions(+), 92 deletions(-) diff --git a/example/taylor_expansion.jl b/example/taylor_expansion.jl index 97e8fc3f..2cfdb0af 100644 --- a/example/taylor_expansion.jl +++ b/example/taylor_expansion.jl @@ -3,7 +3,7 @@ using FeynmanDiagram.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) @@ -11,7 +11,7 @@ g4 = Graph([]) g5 = Graph([]) g6 = Graph([]) G3 = g1 -G4 = 1.0 * g1 * 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) @@ -23,10 +23,11 @@ using FeynmanDiagram.Taylor: # set_variables("x y", order=3) -# @time T5 = taylorexpansion!(G5) +# @time T5 = taylorexpansion(G5) # print(T5) -set_variables("x y z a", order=7) -@time T5 = taylorexpansion!(G6) +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") diff --git a/src/TaylorSeries/arithmetic.jl b/src/TaylorSeries/arithmetic.jl index 5e73b852..049de226 100644 --- a/src/TaylorSeries/arithmetic.jl +++ b/src/TaylorSeries/arithmetic.jl @@ -121,7 +121,7 @@ function Base.:-(g1::TaylorSeries{T}, c::S) where {T,S<:Number} end """ - function taylor_binomial(o1::Array{Int,1}, o2::Array{Int,1}) + function taylor_binomial(o1::Vector{Int}, o2::Vector{Int}) Return the taylor binomial prefactor when product two high-order derivatives with order o1 and o2. @@ -129,7 +129,7 @@ end - `o1` Order of first derivative - `o2` Order of second derivative """ -function taylor_binomial(o1::Array{Int,1}, o2::Array{Int,1}) +function taylor_binomial(o1::Vector{Int}, o2::Vector{Int}) @assert length(o1) == length(o2) result = 1 for i in eachindex(o1) @@ -143,14 +143,14 @@ end """ - function taylor_factorial(o::Array{Int,1}) + function taylor_factorial(o::Vector{Int}) Return the taylor factorial prefactor with order o. # Arguments: - `o` Order of the taylor coefficient """ -function taylor_factorial(o::Array{Int,1}) +function taylor_factorial(o::Vector{Int}) result = 1 for i in eachindex(o) result *= factorial(o[i]) @@ -172,7 +172,7 @@ function Base.:*(g1::TaylorSeries{T}, g2::TaylorSeries{T}) where {T} for (order1, coeff1) in g1.coeffs for (order2, coeff2) in g2.coeffs order = order1 + order2 - if sum(order) <= _params_Taylor_.order + if all(order .<= get_order()) #combination_coeff = taylor_binomial(order1, order2) if haskey(g.coeffs, order) #g.coeffs[order] += combination_coeff * coeff1 * coeff2 @@ -187,9 +187,38 @@ function Base.:*(g1::TaylorSeries{T}, g2::TaylorSeries{T}) where {T} return g end +# """ +# function Base.:*(g1::TaylorSeries{T}, g2::TaylorSeries{T}) where {T} + +# Returns a taylor series `g1 * g2` representing the product of `g2` with `g1`. + +# # Arguments: +# - `g1` First taylor series +# - `g2` Second taylor series +# """ +# function multi_product(glist::Vector{TaylorSeries{T}}) where {T} +# result = TaylorSeries{T}() +# idxtuple = (keys(glist.coeffs) for g in glist) +# for idx in collect(Iterators.product(idxtuple...)) +# orderidx = collect(orderidx) +# totalorder = sum(glist[i].coeffs[orderidx[i]] for i in eachindex(glist)) +# if all(totalorder .<= get_order()) +# #combination_coeff = taylor_binomial(order1, order2) +# if haskey(g.coeffs, order) +# #g.coeffs[order] += combination_coeff * coeff1 * coeff2 +# g.coeffs[order] += coeff1 * coeff2 +# else +# #g.coeffs[order] = combination_coeff * coeff1 * coeff2 +# g.coeffs[order] = coeff1 * coeff2 +# end + +# end +# end +# return g +# end """ - function getcoeff(g::TaylorSeries, order::Array{Int,1}) + function getcoeff(g::TaylorSeries, order::Vector{Int}) Return the taylor coefficients with given order in taylor series g. @@ -197,7 +226,7 @@ end - `g` Taylor series - `order` Order of target coefficients """ -function getcoeff(g::TaylorSeries, order::Array{Int,1}) +function getcoeff(g::TaylorSeries, order::Vector{Int}) if haskey(g.coeffs, order) return g.coeffs[order] else @@ -206,7 +235,7 @@ function getcoeff(g::TaylorSeries, order::Array{Int,1}) end """ - function getderivative(g::TaylorSeries, order::Array{Int,1}) + function getderivative(g::TaylorSeries, order::Vector{Int}) Return the derivative with given order in taylor series g. @@ -214,7 +243,7 @@ end - `g` Taylor series - `order` Order of derivative """ -function getderivative(g::TaylorSeries, order::Array{Int,1}) +function getderivative(g::TaylorSeries, order::Vector{Int}) if haskey(g.coeffs, order) return taylor_factorial(order) * g.coeffs[order] else diff --git a/src/TaylorSeries/constructors.jl b/src/TaylorSeries/constructors.jl index 7722ff85..185c4ef0 100644 --- a/src/TaylorSeries/constructors.jl +++ b/src/TaylorSeries/constructors.jl @@ -5,17 +5,17 @@ # Members: - `name::Symbol` name of the diagram -- `coeffs::Dict{Array{Int,1},T}` The taylor expansion coefficients. The integer array define the order of corresponding coefficient. +- `coeffs::Dict{Vector{Int},T}` The taylor expansion coefficients. The integer array define the order of corresponding coefficient. """ mutable struct TaylorSeries{T} name::String # "" by default - coeffs::Dict{Array{Int,1},T} + coeffs::Dict{Vector{Int},T} """ - function TaylorSeries{T}(coeffs::Dict{Array{Int,1},T}=Dict{Array{Int,1},T}(), name::String="") where {T} + function TaylorSeries{T}(coeffs::Dict{Vector{Int},T}=Dict{Vector{Int},T}(), name::String="") where {T} Create a TaylorSeries based on given coefficients. """ - function TaylorSeries{T}(coeffs::Dict{Array{Int,1},T}=Dict{Array{Int,1},T}(), name::String="") where {T} + function TaylorSeries{T}(coeffs::Dict{Vector{Int},T}=Dict{Vector{Int},T}(), name::String="") where {T} return new{T}(name, coeffs) end end @@ -33,7 +33,7 @@ function TaylorSeries(::Type{T}, nv::Int) where {T} @assert 0 < nv ≤ get_numvars() v = zeros(Int, get_numvars()) @inbounds v[nv] = 1 - return TaylorSeries{T}(Dict{Array{Int,1},T}(v => one(T))) + return TaylorSeries{T}(Dict{Vector{Int},T}(v => one(T))) end TaylorSeries(nv::Int) = TaylorSeries(Float64, nv) diff --git a/src/TaylorSeries/parameter.jl b/src/TaylorSeries/parameter.jl index cf243872..8f775879 100644 --- a/src/TaylorSeries/parameter.jl +++ b/src/TaylorSeries/parameter.jl @@ -14,7 +14,7 @@ These parameters can be changed using [`set_variables`](@ref) """ mutable struct ParamsTaylor - order::Int + order::Vector{Int} num_vars::Int variable_names::Vector{String} variable_symbols::Vector{Symbol} @@ -23,10 +23,11 @@ end ParamsTaylor(order, num_vars, variable_names) = ParamsTaylor(order, num_vars, variable_names, Symbol.(variable_names)) -const _params_Taylor_ = ParamsTaylor(6, 2, ["x₁", "x₂"]) +const _params_Taylor_ = ParamsTaylor([2, 2], 2, ["x₁", "x₂"]) ## Utilities to get the maximum order, number of variables, their names and symbols get_order() = _params_Taylor_.order +get_order(idx::Int) = _params_Taylor_.order[idx] get_numvars() = _params_Taylor_.num_vars get_variable_names() = _params_Taylor_.variable_names get_variable_symbols() = _params_Taylor_.variable_symbols @@ -69,9 +70,9 @@ julia> set_variables("x", order=6, numvars=2) """ function set_variables(::Type{R}, names::Vector{T}; order=get_order()) where {R,T<:AbstractString} - - order ≥ 1 || error("Order must be at least 1") - + for o in order + o ≥ 1 || error("Order must be at least 1") + end num_vars = length(names) num_vars ≥ 1 || error("Number of variables must be at least 1") diff --git a/src/TaylorSeries/print.jl b/src/TaylorSeries/print.jl index 17abac49..14b827f5 100644 --- a/src/TaylorSeries/print.jl +++ b/src/TaylorSeries/print.jl @@ -124,8 +124,11 @@ end function pretty_print(a::TaylorSeries{T}) where {T} #z = zero(a[0]) space = string("") + # bigO::String = bigOnotation[end] ? + # string(" + 𝒪(‖x‖", superscriptify(_params_Taylor_.order + 1), ")") : + # string("") bigO::String = bigOnotation[end] ? - string(" + 𝒪(‖x‖", superscriptify(_params_Taylor_.order + 1), ")") : + string(" + 𝒪(", [_params_Taylor_.variable_names[i] * superscriptify(_params_Taylor_.order[i] + 1) for i in eachindex(_params_Taylor_.order)]..., ")") : string("") # iszero(a) && return string(space, z, space, bigO) # strout::String = space diff --git a/src/utility.jl b/src/utility.jl index 6eb0e5e2..0ec081e7 100644 --- a/src/utility.jl +++ b/src/utility.jl @@ -8,32 +8,53 @@ using ..Taylor """ - function graph_to_taylor(graph::G, var::Array{Bool,1}=fill(true, get_numvars())) where {G<:Graph} + 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 not provided, by default, assume that g depends on all variables. + 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 -- `graph` Target graph. Must be a leaf. -- `var` The variables graph depends on. +- `graph` Target graph. +- `var_dependence` The variables graph leaves depend on. """ -function graph_to_taylor(graph::G, var::Array{Bool,1}=fill(true, get_numvars())) where {G<:Graph} - @assert isleaf(graph) - maxorder = get_order() - ordtuple = ((var[idx]) ? (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() +function taylorexpansion(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph} + if isleaf(graph) + maxorder = get_order() + 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 + end + ordtuple = ((var[idx]) ? (0:get_order(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 end + 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 + + 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 - return result end """ - graph_to_taylor_withmap!(g::G; coeffmode=true, var::Array{Int,1}=collect(1:get_numvars())) where {G<:Graph} + taylorexpansion_withmap(g::G; coeffmode=true, var::Vector{Int}=collect(1:get_numvars())) where {G<:Graph} Return a taylor series of graph g, together with a map of chain relation ship between generated derivatives. This function is only internally used for constructing high order derivatives by naive nested forward AD. @@ -43,7 +64,7 @@ end - `coeffmode` If true, the generated taylor series saves taylor coefficietnts with the factorial prefactor. If false, the taylor series saves derivatives instead - `var` The index of variables graph depends on """ -function graph_to_taylor_withmap!(g::G; coeffmode=true, var::Array{Bool,1}=fill(true, get_numvars())) where {G<:Graph} +function taylorexpansion_withmap(g::G; coeffmode=true, var::Vector{Bool}=fill(true, get_numvars())) where {G<:Graph} @assert isleaf(g) chainrule_map_leaf = Dict{Int,Dict{Int,G}}() maxorder = get_order() @@ -51,8 +72,8 @@ function graph_to_taylor_withmap!(g::G; coeffmode=true, var::Array{Bool,1}=fill( result = TaylorSeries{G}() result.coeffs[zeros(Int, get_numvars())] = g - for i in 1:get_order() - new_func = Dict{Array{Int,1},G}() + for i in 1:sum(get_order()) + new_func = Dict{Vector{Int},G}() for (order, func) in current_func if !haskey(chainrule_map_leaf, func.id) chainrule_map_leaf[func.id] = Dict{Int,G}() @@ -61,18 +82,20 @@ function graph_to_taylor_withmap!(g::G; coeffmode=true, var::Array{Bool,1}=fill( if var[idx] ordernew = copy(order) ordernew[idx] += 1 - if !haskey(result.coeffs, ordernew) - if coeffmode - funcAD = Graph([]; operator=Sum(), factor=g.factor) + if ordernew[idx] <= get_order(idx) + if !haskey(result.coeffs, ordernew) + if coeffmode + funcAD = Graph([]; operator=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) + end + new_func[ordernew] = funcAD + result.coeffs[ordernew] = funcAD + chainrule_map_leaf[func.id][idx] = funcAD else - #funcAD = taylor_factorial(ordernew) * Graph([]; operator=Sum(), factor=g.factor) - funcAD = Graph([]; operator=Sum(), factor=taylor_factorial(ordernew) * g.factor) + chainrule_map_leaf[func.id][idx] = result.coeffs[ordernew] end - 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 @@ -87,34 +110,6 @@ end @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} - - Return the taylor Series of a graph. If taylor series of the leaves of this graph is not provided, by default we assume the leaves depend on all variables. - -# Arguments: -- `graph` Target graph -- `taylormap` The taylor series corresponding to each node of graphs. The taylor series of leafs can be provided as input -- `varidx` The index of variables graph depends on -""" -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) - 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 #Functions below generate high order derivatives with naive nested forward AD. This part would be significantly refactored later with @@ -124,7 +119,7 @@ function build_derivative_backAD!(g::G, leaftaylor::Dict{Int,TaylorSeries{G}}=Di chainrule_map_leaf = Dict{Int,Dict{Int,G}}() for leaf in Leaves(g) if !haskey(leaftaylor, leaf.id) - leaftaylor[leaf.id], map = graph_to_taylor_withmap!(leaf; coeffmode=false) + leaftaylor[leaf.id], map = taylorexpansion_withmap(leaf; coeffmode=false) chainrule_map_leaf = merge(chainrule_map_leaf, map) end end @@ -134,17 +129,19 @@ function build_derivative_backAD!(g::G, leaftaylor::Dict{Int,TaylorSeries{G}}=Di result = TaylorSeries{G}() result.coeffs[zeros(Int, get_numvars())] = g - for i in 1:get_order() - new_func = Dict{Array{Int,1},G}() + for i in 1:sum(get_order()) + new_func = Dict{Vector{Int},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 + if ordernew[idx] <= get_order(idx) + 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 diff --git a/test/taylor.jl b/test/taylor.jl index f785294d..f7d5e8c7 100644 --- a/test/taylor.jl +++ b/test/taylor.jl @@ -3,7 +3,7 @@ using FeynmanDiagram: Taylor as Taylor @testset verbose = true "TaylorSeries" begin using FeynmanDiagram.Taylor: getcoeff - a, b, c, d, e = set_variables("a b c d e", order=3) + a, b, c, d, e = set_variables("a b c d e", order=[3, 3, 3, 3, 3]) F1 = (a + b) * (a + b) * (a + b) print("$(F1)") @test getcoeff(F1, [2, 1, 0, 0, 0]) == 3.0 @@ -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,9 +34,9 @@ using FeynmanDiagram: Taylor as Taylor using FeynmanDiagram.Taylor: TaylorSeries, getcoeff, set_variables - set_variables("x y z", order=5) + set_variables("x y z", order=[2, 3, 2]) for G in [G3, G4, G5, G6] - T = taylorexpansion!(G) + T = 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])