Bugfix in static compiler

Project.toml

@@ -1,7 +1,7 @@
 name = "FeynmanDiagram"
 uuid = "e424a512-dbd9-41ff-9883-094748823e72"
 authors = ["Kun Chen", "Pengcheng Hou", "Daniel Cerkoney"]
-version = "1.0.0"
+version = "1.0.1"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"

src/backend/compiler_python.jl

@@ -1,68 +1,3 @@
-# ms = pyimport("mindspore")
-
-"""
-    function to_pystatic(operator::Type, subgraphs::AbstractVector{<:AbstractGraph}, subgraph_factors::AbstractVector)
-
-    Returns the static representation of a computational graph node `g` with operator `operator`, subgraphs `subgraphs`, and subgraph factors `subgraph_factors` in python.
-"""
-function to_pystatic(operator::Type, subgraphs::AbstractVector{<:AbstractGraph}, subgraph_factors::AbstractVector)
-    error(
-        "Static representation for computational graph nodes with operator $(operator) not yet implemented! " *
-        "Please define a method `to_static(::Type{$(operator)}, subgraphs::$(typeof(subgraphs)), subgraph_factors::$(typeof(subgraph_factors)))`."
-    )
-end
-
-function to_pystatic(::Type{ComputationalGraphs.Sum}, subgraphs::Vector{Graph{F,W}}, subgraph_factors::Vector{F}) where {F,W}
-    if length(subgraphs) == 1
-        factor_str = subgraph_factors[1] == 1 ? "" : " * $(subgraph_factors[1])"
-        return "(g$(subgraphs[1].id)$factor_str)"
-    else
-        terms = ["g$(g.id)" * (gfactor == 1 ? "" : " * $gfactor") for (g, gfactor) in zip(subgraphs, subgraph_factors)]
-        return "(" * join(terms, " + ") * ")"
-    end
-end
-
-function to_pystatic(::Type{ComputationalGraphs.Prod}, subgraphs::Vector{Graph{F,W}}, subgraph_factors::Vector{F}) where {F,W}
-    if length(subgraphs) == 1
-        factor_str = subgraph_factors[1] == 1 ? "" : " * $(subgraph_factors[1])"
-        return "(g$(subgraphs[1].id)$factor_str)"
-    else
-        terms = ["g$(g.id)" * (gfactor == 1 ? "" : " * $gfactor") for (g, gfactor) in zip(subgraphs, subgraph_factors)]
-        return "(" * join(terms, " * ") * ")"
-        # return "(" * join(["g$(g.id)" for g in subgraphs], " * ") * ")"
-    end
-end
-
-function to_pystatic(::Type{ComputationalGraphs.Power{N}}, subgraphs::Vector{Graph{F,W}}, subgraph_factors::Vector{F}) where {N,F,W}
-    factor_str = subgraph_factors[1] == 1 ? "" : " * $(subgraph_factors[1])"
-    return "((g$(subgraphs[1].id))**$N$factor_str)"
-end
-
-function to_pystatic(::Type{ComputationalGraphs.Sum}, subgraphs::Vector{FeynmanGraph{F,W}}, subgraph_factors::Vector{F}) where {F,W}
-    if length(subgraphs) == 1
-        factor_str = subgraph_factors[1] == 1 ? "" : " * $(subgraph_factors[1])"
-        return "(g$(subgraphs[1].id)$factor_str)"
-    else
-        terms = ["g$(g.id)" * (gfactor == 1 ? "" : " * $gfactor") for (g, gfactor) in zip(subgraphs, subgraph_factors)]
-        return "(" * join(terms, " + ") * ")"
-    end
-end
-
-function to_pystatic(::Type{ComputationalGraphs.Prod}, subgraphs::Vector{FeynmanGraph{F,W}}, subgraph_factors::Vector{F}) where {F,W}
-    if length(subgraphs) == 1
-        factor_str = subgraph_factors[1] == 1 ? "" : " * $(subgraph_factors[1])"
-        return "(g$(subgraphs[1].id)$factor_str)"
-    else
-        terms = ["g$(g.id)" * (gfactor == 1 ? "" : " * $gfactor") for (g, gfactor) in zip(subgraphs, subgraph_factors)]
-        return "(" * join(terms, " * ") * ")"
-    end
-end
-
-function to_pystatic(::Type{ComputationalGraphs.Power{N}}, subgraphs::Vector{FeynmanGraph{F,W}}, subgraph_factors::Vector{F}) where {N,F,W}
-    factor_str = subgraph_factors[1] == 1 ? "" : " * $(subgraph_factors[1])"
-    return "((g$(subgraphs[1].id))**$N$factor_str)"
-end
-
 """
     function to_python_str(graphs::AbstractVector{<:AbstractGraph})
 
@@ -103,7 +38,7 @@ function to_python_str(graphs::AbstractVector{<:AbstractGraph}, framework::Symbo
                 push!(inds_visitedleaf, g_id)
             else
                 g_id in inds_visitednode && continue
-                body *= "    $target = $(to_pystatic(operator(g), subgraphs(g), subgraph_factors(g)))\n"
+                body *= "    $target = $(to_static(operator(g), subgraphs(g), subgraph_factors(g), lang=:python))\n"
                 push!(inds_visitednode, g_id)
             end
             if isroot

src/backend/static.jl

@@ -1,16 +1,16 @@
 """
     function to_static(operator::Type, subgraphs::AbstractVector{<:AbstractGraph}, subgraph_factors::AbstractVector)
 
-Returns the static representation of a computational graph node `g` with operator `operator`, subgraphs `subgraphs`, and subgraph factors `subgraph_factors`.
+    Returns the static representation of a computational graph node `g` with operator `operator`, subgraphs `subgraphs`, and subgraph factors `subgraph_factors`.
 """
-function to_static(operator::Type, subgraphs::AbstractVector{<:AbstractGraph}, subgraph_factors::AbstractVector)
+function to_static(operator::Type, subgraphs::AbstractVector{<:AbstractGraph}, subgraph_factors::AbstractVector; lang::Symbol=:julia)
     error(
         "Static representation for computational graph nodes with operator $(operator) not yet implemented! " *
         "Please define a method `to_static(::Type{$(operator)}, subgraphs::$(typeof(subgraphs)), subgraph_factors::$(typeof(subgraph_factors)))`."
     )
 end
 
-function to_static(::Type{ComputationalGraphs.Sum}, subgraphs::Vector{Graph{F,W}}, subgraph_factors::Vector{F}) where {F,W}
+function to_static(::Type{ComputationalGraphs.Sum}, subgraphs::Vector{Graph{F,W}}, subgraph_factors::Vector{F}; lang::Symbol=:julia) where {F,W}
     if length(subgraphs) == 1
         factor_str = subgraph_factors[1] == 1 ? "" : " * $(subgraph_factors[1])"
         return "(g$(subgraphs[1].id)$factor_str)"
@@ -20,7 +20,7 @@ function to_static(::Type{ComputationalGraphs.Sum}, subgraphs::Vector{Graph{F,W}
     end
 end
 
-function to_static(::Type{ComputationalGraphs.Prod}, subgraphs::Vector{Graph{F,W}}, subgraph_factors::Vector{F}) where {F,W}
+function to_static(::Type{ComputationalGraphs.Prod}, subgraphs::Vector{Graph{F,W}}, subgraph_factors::Vector{F}; lang::Symbol=:julia) where {F,W}
     if length(subgraphs) == 1
         factor_str = subgraph_factors[1] == 1 ? "" : " * $(subgraph_factors[1])"
         return "(g$(subgraphs[1].id)$factor_str)"
@@ -31,12 +31,21 @@ function to_static(::Type{ComputationalGraphs.Prod}, subgraphs::Vector{Graph{F,W
     end
 end
 
-function to_static(::Type{ComputationalGraphs.Power{N}}, subgraphs::Vector{Graph{F,W}}, subgraph_factors::Vector{F}) where {N,F,W}
+function to_static(::Type{ComputationalGraphs.Power{N}}, subgraphs::Vector{Graph{F,W}}, subgraph_factors::Vector{F}; lang::Symbol=:julia) where {N,F,W}
     factor_str = subgraph_factors[1] == 1 ? "" : " * $(subgraph_factors[1])"
-    return "((g$(subgraphs[1].id))^$N$factor_str)"
+    if lang == :julia
+        op_str = "^"
+    elseif lang == :c
+        return "pow(g$(subgraphs[1].id), $N)$factor_str"
+    elseif lang == :python
+        op_str = "**"
+    else
+        error("Unsupported language")
+    end
+    return "((g$(subgraphs[1].id))$(op_str)$N$factor_str)"
 end
 
-function to_static(::Type{ComputationalGraphs.Sum}, subgraphs::Vector{FeynmanGraph{F,W}}, subgraph_factors::Vector{F}) where {F,W}
+function to_static(::Type{ComputationalGraphs.Sum}, subgraphs::Vector{FeynmanGraph{F,W}}, subgraph_factors::Vector{F}; lang::Symbol=:julia) where {F,W}
     if length(subgraphs) == 1
         factor_str = subgraph_factors[1] == 1 ? "" : " * $(subgraph_factors[1])"
         return "(g$(subgraphs[1].id)$factor_str)"
@@ -46,7 +55,7 @@ function to_static(::Type{ComputationalGraphs.Sum}, subgraphs::Vector{FeynmanGra
     end
 end
 
-function to_static(::Type{ComputationalGraphs.Prod}, subgraphs::Vector{FeynmanGraph{F,W}}, subgraph_factors::Vector{F}) where {F,W}
+function to_static(::Type{ComputationalGraphs.Prod}, subgraphs::Vector{FeynmanGraph{F,W}}, subgraph_factors::Vector{F}; lang::Symbol=:julia) where {F,W}
     if length(subgraphs) == 1
         factor_str = subgraph_factors[1] == 1 ? "" : " * $(subgraph_factors[1])"
         return "(g$(subgraphs[1].id)$factor_str)"
@@ -56,9 +65,18 @@ function to_static(::Type{ComputationalGraphs.Prod}, subgraphs::Vector{FeynmanGr
     end
 end
 
-function to_static(::Type{ComputationalGraphs.Power{N}}, subgraphs::Vector{FeynmanGraph{F,W}}, subgraph_factors::Vector{F}) where {N,F,W}
+function to_static(::Type{ComputationalGraphs.Power{N}}, subgraphs::Vector{FeynmanGraph{F,W}}, subgraph_factors::Vector{F}; lang::Symbol=:julia) where {N,F,W}
     factor_str = subgraph_factors[1] == 1 ? "" : " * $(subgraph_factors[1])"
-    return "((g$(subgraphs[1].id))^$N$factor_str)"
+    if lang == :julia
+        op_str = "^"
+    elseif lang == :c
+        return "pow(g$(subgraphs[1].id), $N)$factor_str"
+    elseif lang == :python
+        op_str = "**"
+    else
+        error("Unsupported language")
+    end
+    return "((g$(subgraphs[1].id))$(op_str)$N$factor_str)"
 end
 
 """
@@ -165,7 +183,7 @@ function to_Cstr(graphs::AbstractVector{<:AbstractGraph}; root::AbstractVector{I
             else
                 g_id in inds_visitednode && continue
                 declare *= " g$g_id,"
-                body *= "    $target = $(to_static(operator(g), subgraphs(g), subgraph_factors(g)));\n"
+                body *= "    $target = $(to_static(operator(g), subgraphs(g), subgraph_factors(g), lang=:c));\n"
                 push!(inds_visitednode, g_id)
             end
             if isroot
@@ -252,6 +270,9 @@ function compile_C(graphs::AbstractVector{<:AbstractGraph}, filename::String;
     datatype::DataType=_dtype.weight, root::AbstractVector{Int}=[id(g) for g in graphs], func_name="eval_graph")
     func_string, leafmap = to_Cstr(graphs; datatype=datatype, root=root, name=func_name)
     open(filename, "a") do f
+        if position(f) == 0
+            write(f, "#include <math.h>\n")
+        end
         write(f, func_string)
     end
     return leafmap

src/frontend/frontends.jl

@@ -1,6 +1,6 @@
 module FrontEnds
 
-import ..ComputationalGraphs
+import ..ComputationalGraphs 
 import ..ComputationalGraphs: Graph, FeynmanGraph, _dtype
 import ..QuantumOperators
 import ..Taylor
@@ -134,11 +134,11 @@ function leafstates(leaf_maps::Vector{Dict{Int,G}}, labelProd::LabelProduct) whe
         for idx in 1:len_leaves
             g = leafmap[idx]
             vertices = g.properties.vertices
-            if IR.diagram_type(g) == IR.Interaction
+            if ComputationalGraphs.diagram_type(g) == ComputationalGraphs.Interaction
                 In = Out = vertices[1][1].label
                 push!(leafType[ikey], 0)
                 push!(leafLoopIndex[ikey], 1)
-            elseif IR.diagram_type(g) == IR.Propagator
+            elseif ComputationalGraphs.diagram_type(g) == ComputationalGraphs.Propagator
                 if (Op.isfermionic(vertices[1]))
                     In, Out = vertices[2][1].label, vertices[1][1].label
                     # push!(leafType[ikey], g.orders[1] * 2 + 1)
@@ -197,7 +197,7 @@ function leafstates(leaf_maps::Vector{Dict{Int,G}}, maxloopNum::Int) where {G<:G
 
         for idx in 1:len_leaves
             leaf = leafmap[idx]
-            @assert IR.isleaf(leaf)
+            @assert ComputationalGraphs.isleaf(leaf)
             diagId, leaf_orders = leaf.properties, leaf.orders
             loopmom = copy(diagId.extK)
             len = length(loopmom)