Skip to content

Commit

Permalink
Merge pull request #131 from numericalEFT/computgraph_pchou
Browse files Browse the repository at this point in the history
Update the leaf mapping in optimize, compiler, and GV.
  • Loading branch information
houpc authored Sep 19, 2023
2 parents 6acd9f7 + 2551137 commit 3985773
Show file tree
Hide file tree
Showing 3 changed files with 70 additions and 109 deletions.
35 changes: 23 additions & 12 deletions src/backend/static.jl
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,8 @@ end
"""
function to_julia_str(graphs::AbstractVector; root::AbstractVector{Int}=[g.id for g in graphs], name::String="eval_graph!")
Compile a list of graphs into a string for a julia static function. The function takes two arguments: `root` and `leaf`. `root` is a vector of the root node ids of the graphs, and `leaf` is a vector of the leaf node ids of the graphs.
Compile a list of graphs into a string for a julia static function. The function takes two arguments: `root` and `leaf`.
`root` is a vector of the root node ids of the graphs, and `leaf` is a vector of the leaf nodes' weights of the graphs.
"""
function to_julia_str(graphs::AbstractVector; root::AbstractVector{Int}=[g.id for g in graphs], name::String="eval_graph!")
head = "function $name(root::AbstractVector, leaf::AbstractVector)\n "
Expand All @@ -42,10 +43,24 @@ function to_julia_str(graphs::AbstractVector; root::AbstractVector{Int}=[g.id fo
return head * body * tail
end

function to_julia_str(graphs::AbstractVector, propagatorMap::Dict{Int,Int}, interactionMap::Dict{Int,Int}; root::AbstractVector{Int}=[g.id for g in graphs], name::String="eval_graph!")
head = "function $name(root::AbstractVector, propagatorVal::AbstractVector, interactionVal::AbstractVector)\n "
"""
function to_julia_str(graphs::AbstractVector, leafMap::Dict{Int,Int}; root::AbstractVector{Int}=[g.id for g in graphs],
leaftypes=[ComputationalGraphs.Propagator, ComputationalGraphs.Interaction], name::String="eval_graph!")
Compile a list of graphs into a string for a julia static function. The complied function takes two arguments: `root` and `leafVal`.
`root` is a vector of the root node ids of the graphs, and `leafVal` is a vector of the leaf nodes' weights of the graphs.
# Arguments:
- `graphs` (AbstractVector): The vector object representing the graphs,
- `leafMap (Dict{Int,Int})`: The mapping dictionary from the id of each leaf to the index of the leaf weight's table `leafVal`.
- `root` (AbstractVector{Int}, optional): The vector of the root node ids of the graphs (defaults to `[g.id for g in graphs]`).
- `leaftypes (optional)`: The set of `GraphType` for all the relevant leaves to the weight table. It defaults to `[Propagator, Interaction]`.
- `name` (String,optional): The name of the complied function (defaults to `"eval_graph!"`).
"""
function to_julia_str(graphs::AbstractVector, leafMap::Dict{Int,Int}; root::AbstractVector{Int}=[g.id for g in graphs],
leaftypes=[ComputationalGraphs.Propagator, ComputationalGraphs.Interaction], name::String="eval_graph!")
head = "function $name(root::AbstractVector, leafVal::AbstractVector)\n "
body = ""
pIdx, iIdx = 1, 1
for graph in graphs
for g in PostOrderDFS(graph) #leaf first search
if g.id in root
Expand All @@ -55,12 +70,8 @@ function to_julia_str(graphs::AbstractVector, propagatorMap::Dict{Int,Int}, inte
end
if isempty(g.subgraphs) #leaf
g.name == "compiled" && continue
if g.type == ComputationalGraphs.Propagator
body *= " $target = propagatorVal[$(propagatorMap[g.id])]\n "
pIdx += 1
elseif g.type == ComputationalGraphs.Interaction
body *= " $target = interactionVal[$(interactionMap[g.id])]\n "
iIdx += 1
if g.type in leaftypes
body *= " $target = leafVal[$(leafMap[g.id])]\n "
end
g.name = "compiled"
else
Expand Down Expand Up @@ -102,10 +113,10 @@ function compile(graphs::AbstractVector;
return @RuntimeGeneratedFunction(func_expr)
end

function compile(graphs::AbstractVector, propagatorMap::Dict{Int,Int}, interactionMap::Dict{Int,Int};
function compile(graphs::AbstractVector, leafMap::Dict{Int,Int};
root::AbstractVector{Int}=[g.id for g in graphs])
# this function return a runtime generated function defined by compile()
func_string = to_julia_str(graphs, propagatorMap, interactionMap; root=root, name="func_name!")
func_string = to_julia_str(graphs, leafMap; root=root, name="func_name!")
func_expr = Meta.parse(func_string)
return @RuntimeGeneratedFunction(func_expr)
end
53 changes: 7 additions & 46 deletions src/computational_graph/optimize.jl
Original file line number Diff line number Diff line change
@@ -1,13 +1,11 @@
function optimize!(graphs::Union{Tuple,AbstractVector}; verbose=0, normalize=nothing)
if isempty(graphs)
# return graphs
return nothing
else
graphs = collect(graphs)
removeOneChildParent!(graphs, verbose=verbose)
mappings = removeDuplicatedLeaves!(graphs, verbose=verbose, normalize=normalize)
return mappings
# return graphs, mappings
end
end

Expand All @@ -22,7 +20,7 @@ function removeOneChildParent!(graphs::AbstractVector{G}; verbose=0) where {G<:G
return graphs
end

function removeDuplicatedLeaves!(graphs::AbstractVector{G}; verbose=0, normalize=nothing, kwargs...) where {G<:Graph}
function removeDuplicatedLeaves!(graphs::AbstractVector{G}; verbose=0, normalize=nothing, leaftypes=[Interaction, Propagator], kwargs...) where {G<:Graph}
verbose > 0 && println("remove duplicated leaves.")
leaves = Vector{G}()
for g in graphs
Expand All @@ -39,7 +37,7 @@ function removeDuplicatedLeaves!(graphs::AbstractVector{G}; verbose=0, normalize

for l in leaves
#make sure all leaves are either propagators or interactions
@assert l.type in [Interaction, Propagator]
@assert l.type in leaftypes
end

function uniqueLeaves(_graphs::Vector{G}) where {G}
Expand All @@ -59,62 +57,25 @@ function removeDuplicatedLeaves!(graphs::AbstractVector{G}; verbose=0, normalize
end
if flag
push!(uniqueGraph, g)
# push!(mapping, length(uniqueGraph))
mapping[g.id] = idx
idx += 1
end
end
return uniqueGraph, mapping
end

green = [l for l in leaves if l.type == Propagator]
interaction = [l for l in leaves if l.type == Interaction]

uniqueGreen, greenMap = uniqueLeaves(green)
uniqueInteraction, interactionMap = uniqueLeaves(interaction)

verbose > 0 && length(green) > 0 && println("Number of independent Propagators $(length(green))$(length(uniqueGreen))")
verbose > 0 && length(interaction) > 0 && println("Number of independent Interactions $(length(interaction))$(length(uniqueInteraction))")
uniqueLeaf, leafMap = uniqueLeaves(leaves)
verbose > 0 && length(leaves) > 0 && println("Number of independent Interactions $(length(leaves))$(length(uniqueLeaf))")

for g in graphs
for n in PreOrderDFS(g)
for (si, sub_g) in enumerate(n.subgraphs)
if sub_g.type == Propagator
n.subgraphs[si] = uniqueGreen[greenMap[sub_g.id]]
elseif sub_g.type == Interaction
n.subgraphs[si] = uniqueInteraction[interactionMap[sub_g.id]]
if sub_g.type in leaftypes
n.subgraphs[si] = uniqueLeaf[leafMap[sub_g.id]]
end
end
end
end

# return uniqueGreen, uniqueInteraction
return greenMap, interactionMap
return leafMap
end

# function removeDuplicatedLeaves!(graphs::AbstractVector{G}; verbose=0, normalize=nothing, kwargs...) where {G<:Graph}
# verbose > 0 && println("remove duplicated leaves.")
# uniqueGreen, uniqueInteraction = Vector{G}(), Vector{G}()
# for g in graphs
# for l in Leaves(g)
# if l.type == Interaction
# loc = findfirst(x -> isequiv(x, l, :id), uniqueInteraction)
# if !isnothing(loc)
# l.id = uniqueInteraction[loc].id
# else
# push!(uniqueInteraction, l)
# end
# elseif l.type == Propagator
# loc = findfirst(x -> isequiv(x, l, :id), uniqueGreen)
# if !isnothing(loc)
# l.id = uniqueGreen[loc].id
# else
# push!(uniqueGreen, l)
# end
# else
# error("the leaf's type cannot be $(l.type)!")
# end
# end
# end
# return uniqueGreen, uniqueInteraction
# end
Loading

0 comments on commit 3985773

Please sign in to comment.