Changes for benchmark

madgraph_plugin/PyOut_create_aloha.py

@@ -219,7 +219,7 @@ def get_header_txt(self, name=None, couplings=None, **opt):
         arguments_names = [arg[1] for arg in arguments]
 
         # the signature
-        shape_dict = {'list_complex' : '[None,None]',
+        shape_dict = {'list_complex' : '[6,None]',
                       'complex' : '[None]',
                       'double' : '[]'}
         type_dict = {'list_complex' : 'DTYPECOMPLEX',

madgraph_plugin/template_files/matrix_method_python.inc

@@ -43,14 +43,14 @@ def get_model_param(model, param_card_path):
 
 
 smatrix_signature = [
-        tf.TensorSpec(shape=[None,None,4], dtype=DTYPE),
+        tf.TensorSpec(shape=[None,%(nexternal)d,4], dtype=DTYPE),
         %(paramsignature_const)s
         %(paramsignature_func)s
         ]
 
 
 matrix_signature = [
-        tf.TensorSpec(shape=[None,None,4], dtype=DTYPE),
+        tf.TensorSpec(shape=[None,%(nexternal)d,4], dtype=DTYPE),
         tf.TensorSpec(shape=[%(nexternal)d], dtype=DTYPE),
         %(paramsignature_const)s
         %(paramsignature_func)s
@@ -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,13 +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:
-            t = self.matrix(all_ps,hel,%(params)s)
-            ans = ans + t
 
-        return tf.math.real(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):
@@ -118,11 +121,9 @@ class Matrix_%(process_string)s(object):
         # Process parameters
         #  
         ngraphs = %(ngraphs)d
-        nexternal = self.nexternal
         nwavefuncs = %(nwavefuncs)d
         ncolor = %(ncolor)d
         ZERO = float_me(0.)
-        nevts = tf.shape(all_ps, out_type=DTYPEINT)[0]
         #  
         # Color matrix
         #  
@@ -133,22 +134,13 @@ class Matrix_%(process_string)s(object):
         # ----------
         # Begin code
         # ----------
-        amp = [None] * ngraphs
-        w = [None] * nwavefuncs
         %(helas_calls)s
 
         %(jamp_lines)s
 
         ##%(amp2_lines)s
-        matrix = tf.zeros(nevts, dtype=DTYPECOMPLEX)
-        for i in range(ncolor):
-            ztemp = tf.zeros(nevts, dtype=DTYPECOMPLEX)
-            for j in range(ncolor):
-                ztemp = ztemp + cf[i][j]*jamp[j]
-            matrix = matrix + ztemp * tf.math.conj(jamp[i])/denom[i]   
-        #self.jamp.append(jamp)
-
-        return matrix
+        ret = tf.einsum("ie, ij, je -> e", jamp, cf, tf.math.conj(jamp)/tf.reshape(denom, (ncolor, 1)))
+        return tf.math.real(ret)
 
 
 

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

python_package/madflow/scripts/madflow_exec.py

@@ -265,7 +265,7 @@ def luminosity_function(x1, x2, q2array):
     def generate_integrand(lhewriter=None):
         """Generate a cross section with (or without) a LHE parser"""
 
-        def cross_section(xrand, n_dim=ndim, weight=1.0):
+        def cross_section(xrand, weight=1.0):
             """Compute the cross section"""
             # Generate the phase space point
             all_ps, wts, x1, x2, idx = phasespace(xrand)
@@ -306,7 +306,8 @@ def cross_section(xrand, n_dim=ndim, weight=1.0):
                 tf.py_function(func=lhewriter.lhe_parser, inp=[all_ps, ret * weight], Tout=DTYPE)
 
             if args.pt_cut is not None:
-                ret = tf.scatter_nd(idx, ret, shape=xrand.shape[0:1])
+                out_shape = tf.shape(xrand, out_type=DTYPEINT)[0:1]
+                ret = tf.scatter_nd(idx, ret, shape=out_shape)
 
             return ret
 
@@ -325,7 +326,7 @@ def cross_section(xrand, n_dim=ndim, weight=1.0):
 
     vegas = VegasFlow(ndim, events_per_iteration, events_limit=events_limit)
     integrand = generate_integrand()
-    vegas.compile(integrand)
+    vegas.compile(integrand, trace=True)
 
     if args.frozen_iter == 0:
         warmup_iterations = args.iterations // 2