From 2afb8205d5ca7a61397798126c5a7cfb3921cc91 Mon Sep 17 00:00:00 2001 From: juacrumar Date: Fri, 11 Jun 2021 12:17:36 +0200 Subject: [PATCH 1/5] change color loop to einsum --- .../template_files/matrix_method_python.inc | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/madgraph_plugin/template_files/matrix_method_python.inc b/madgraph_plugin/template_files/matrix_method_python.inc index 976b54d..2a7edb4 100644 --- a/madgraph_plugin/template_files/matrix_method_python.inc +++ b/madgraph_plugin/template_files/matrix_method_python.inc @@ -118,11 +118,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 +131,12 @@ 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 + return tf.einsum("ie, ij, je -> e", jamp, cf, tf.math.conj(jamp)/tf.reshape(denom, (ncolor, 1))) From e4c6a60d90d8b33196df9de9b0de171fe11eaaad Mon Sep 17 00:00:00 2001 From: juacrumar Date: Fri, 11 Jun 2021 17:22:47 +0200 Subject: [PATCH 2/5] add exact information for the signatures --- madgraph_plugin/PyOut_create_aloha.py | 2 +- madgraph_plugin/template_files/matrix_method_python.inc | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/madgraph_plugin/PyOut_create_aloha.py b/madgraph_plugin/PyOut_create_aloha.py index fdfdcef..fdf1550 100644 --- a/madgraph_plugin/PyOut_create_aloha.py +++ b/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', diff --git a/madgraph_plugin/template_files/matrix_method_python.inc b/madgraph_plugin/template_files/matrix_method_python.inc index 2a7edb4..881e516 100644 --- a/madgraph_plugin/template_files/matrix_method_python.inc +++ b/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 From e080199ad434b5839fae01f57a103f4899093733 Mon Sep 17 00:00:00 2001 From: juacrumar Date: Sun, 13 Jun 2021 12:56:22 +0200 Subject: [PATCH 3/5] some changes --- python_package/madflow/scripts/madflow_exec.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python_package/madflow/scripts/madflow_exec.py b/python_package/madflow/scripts/madflow_exec.py index c32210e..c7ad5fc 100755 --- a/python_package/madflow/scripts/madflow_exec.py +++ b/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 From c599323ffc1332913334b6f9db5adfeae6553338 Mon Sep 17 00:00:00 2001 From: juacrumar Date: Sun, 13 Jun 2021 13:11:38 +0200 Subject: [PATCH 4/5] more fixes --- madgraph_plugin/template_files/matrix_method_python.inc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/madgraph_plugin/template_files/matrix_method_python.inc b/madgraph_plugin/template_files/matrix_method_python.inc index 881e516..bbe5534 100644 --- a/madgraph_plugin/template_files/matrix_method_python.inc +++ b/madgraph_plugin/template_files/matrix_method_python.inc @@ -99,10 +99,9 @@ class Matrix_%(process_string)s(object): 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 + ans += self.matrix(all_ps,hel,%(params)s) - return tf.math.real(ans)/self.denominator + return ans/self.denominator @tf.function(input_signature=matrix_signature) def matrix(self,all_ps,hel,%(params)s): @@ -136,7 +135,8 @@ class Matrix_%(process_string)s(object): %(jamp_lines)s ##%(amp2_lines)s - return tf.einsum("ie, ij, je -> e", jamp, cf, tf.math.conj(jamp)/tf.reshape(denom, (ncolor, 1))) + ret = tf.einsum("ie, ij, je -> e", jamp, cf, tf.math.conj(jamp)/tf.reshape(denom, (ncolor, 1))) + return tf.math.real(ret) From 2e689d0a30ad68c8e78a63506ee6b28da75b0743 Mon Sep 17 00:00:00 2001 From: juacrumar Date: Sun, 13 Jun 2021 14:50:29 +0200 Subject: [PATCH 5/5] final --- .../template_files/matrix_method_python.inc | 14 +++++++++----- python_package/madflow/phasespace.py | 4 ++++ 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/madgraph_plugin/template_files/matrix_method_python.inc b/madgraph_plugin/template_files/matrix_method_python.inc index bbe5534..e5846f7 100644 --- a/madgraph_plugin/template_files/matrix_method_python.inc +++ b/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): diff --git a/python_package/madflow/phasespace.py b/python_package/madflow/phasespace.py index 55c4ca0..d90fd0a 100644 --- a/python_package/madflow/phasespace.py +++ b/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