final

Author: scarlehoff

madgraph_plugin/template_files/matrix_method_python.inc

@@ -77,6 +77,14 @@ class Matrix_%(process_string)s(object):
     def __str__(self):
         return "%(process_string)s"
 
+    @tf.function(input_signature=[tf.TensorSpec(shape=[None, %(nexternal)d], dtype=DTYPE)]+smatrix_signature)
+    def wrapper(self, all_hel, all_ps, %(params)s):
+        nevts = tf.shape(all_ps, out_type=DTYPEINT)[0]
+        ans = tf.zeros(nevts, dtype=DTYPE)
+        for hel in self.helicities:
+            ans += self.matrix(all_ps,hel,%(params)s)
+        return ans
+
     @tf.function(input_signature=smatrix_signature)
     def smatrix(self,all_ps,%(params)s):
         #  
@@ -96,12 +104,8 @@ class Matrix_%(process_string)s(object):
         # ----------
         # BEGIN CODE
         # ----------
-        nevts = tf.shape(all_ps, out_type=DTYPEINT)[0]
-        ans = tf.zeros(nevts, dtype=DTYPECOMPLEX)
-        for hel in self.helicities:
-            ans += self.matrix(all_ps,hel,%(params)s)
 
-        return ans/self.denominator
+        return self.wrapper(self.helicities, all_ps, %(params)s)/self.denominator
 
     @tf.function(input_signature=matrix_signature)
     def matrix(self,all_ps,hel,%(params)s):

python_package/madflow/phasespace.py

@@ -19,6 +19,7 @@
 
 # Helpers for rambo
 events_signature = tf.TensorSpec(shape=[None, 1], dtype=DTYPE)
+events_signature_clean = tf.TensorSpec(shape=[None], dtype=DTYPE)
 p_signature = tf.TensorSpec(shape=[None, 4], dtype=DTYPE)
 ps_signature = tf.TensorSpec(shape=[None, None, 4], dtype=DTYPE)
 
@@ -118,6 +119,7 @@ def _conformal_transformation(input_q, bquad):
     return tf.concat([pnrg, pvec], axis=1)  # (n_events, 4)
 
 
+@tf.function(input_signature=[p_signature])
 def _gen_unconstrained_momenta(xrand):
     """
     Generates unconstrained 4-momenta
@@ -212,6 +214,7 @@ def rambo(xrand, n_particles, sqrts, masses=None, check_physical=False):
     return massive_p, wt
 
 
+@tf.function(input_signature=[tf.TensorSpec(shape=[None, 2], dtype=DTYPE)] + 2*[tf.TensorSpec(shape=[], dtype=DTYPE)])
 def _get_x1x2(xarr, shat_min, s_in):
     """Receives two random numbers and return the
     value of the invariant mass of the center of mass
@@ -291,6 +294,7 @@ def ramboflow(xrand, nparticles, com_sqrts, masses=None):
     return final_p, wgt, x1, x2
 
 
+@tf.function(input_signature=[ps_signature] + [events_signature_clean]*2)
 def _boost_to_lab(p_com, x1, x2):
     """Boost the momenta back from the COM frame of the initial partons
     to the lab frame