1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16 """Methods and classes to export matrix elements to Python format."""
17
18 import fractions
19 import glob
20 import itertools
21 import logging
22 import os
23 import re
24 import shutil
25 import subprocess
26 import aloha
27
28 import madgraph.core.color_algebra as color
29 import madgraph.core.helas_objects as helas_objects
30 import madgraph.iolibs.drawing_eps as draw
31 import madgraph.iolibs.files as files
32 import madgraph.iolibs.helas_call_writers as helas_call_writers
33 import madgraph.iolibs.file_writers as writers
34 import madgraph.iolibs.template_files as Template
35 import madgraph.iolibs.ufo_expression_parsers as parsers
36 import madgraph.iolibs.group_subprocs as group_subprocs
37 from madgraph import MadGraph5Error, MG5DIR
38
39 import madgraph.various.misc as misc
40
41 import aloha.create_aloha as create_aloha
42 import aloha.aloha_writers as aloha_writers
43
44 _file_path = os.path.split(os.path.dirname(os.path.realpath(__file__)))[0] + '/'
45 logger = logging.getLogger('madgraph.export_python')
46
47
48
49
50
52 """Class to take care of exporting a set of matrix elements to
53 Python format."""
54
57
58 - def __init__(self, matrix_elements, python_helas_call_writer):
59 """Initiate with matrix elements, helas call writer.
60 Generate the process matrix element functions as strings."""
61
62 self.config_maps = {}
63 if isinstance(matrix_elements, helas_objects.HelasMultiProcess):
64 self.matrix_elements = matrix_elements.get('matrix_elements')
65 elif isinstance(matrix_elements,
66 group_subprocs.SubProcessGroup):
67 self.config_maps = matrix_elements.get('diagram_maps')
68 self.matrix_elements = matrix_elements.get('matrix_elements')
69 elif isinstance(matrix_elements,
70 helas_objects.HelasMatrixElementList):
71 self.matrix_elements = matrix_elements
72 elif isinstance(matrix_elements,
73 helas_objects.HelasMatrixElement):
74 self.matrix_elements = helas_objects.HelasMatrixElementList(\
75 [matrix_elements])
76 if not self.matrix_elements:
77 raise MadGraph5Error("No matrix elements to export")
78
79 self.model = self.matrix_elements[0].get('processes')[0].get('model')
80
81 self.helas_call_writer = python_helas_call_writer
82
83 if not isinstance(self.helas_call_writer, helas_call_writers.PythonUFOHelasCallWriter):
84 raise Exception, \
85 "helas_call_writer not PythonUFOHelasCallWriter"
86
87 self.matrix_methods = {}
88
89
90
91
92
93
95 """Write the matrix element calculation method for the processes"""
96
97 replace_dict = {}
98
99
100 info_lines = self.get_mg5_info_lines()
101 replace_dict['info_lines'] = info_lines
102
103 for ime, matrix_element in enumerate(self.matrix_elements):
104 process_string = matrix_element.get('processes')[0].shell_string()
105 if process_string in self.matrix_methods:
106 continue
107
108 replace_dict['process_string'] = process_string
109
110
111 (nexternal, ninitial) = matrix_element.get_nexternal_ninitial()
112 replace_dict['nexternal'] = nexternal
113
114
115 ncomb = matrix_element.get_helicity_combinations()
116 replace_dict['ncomb'] = ncomb
117
118
119 helicity_lines = self.get_helicity_matrix(matrix_element)
120 replace_dict['helicity_lines'] = helicity_lines
121
122
123
124 den_factor_line = self.get_den_factor_line(matrix_element)
125 replace_dict['den_factor_line'] = den_factor_line
126
127
128 process_lines = self.get_process_info_lines(matrix_element)
129 replace_dict['process_lines'] = process_lines
130
131
132 ngraphs = matrix_element.get_number_of_amplitudes()
133 replace_dict['ngraphs'] = ngraphs
134
135
136 ndiags = len(matrix_element.get('diagrams'))
137 replace_dict['ndiags'] = ndiags
138
139
140 helas_calls = self.helas_call_writer.get_matrix_element_calls(\
141 matrix_element, gauge_check)
142 replace_dict['helas_calls'] = "\n ".join(helas_calls)
143
144
145 nwavefuncs = matrix_element.get_number_of_wavefunctions()
146 replace_dict['nwavefuncs'] = nwavefuncs
147
148
149 ncolor = max(1, len(matrix_element.get('color_basis')))
150 replace_dict['ncolor'] = ncolor
151
152
153 model_parameter_lines = \
154 self.get_model_parameter_lines(matrix_element)
155 replace_dict['model_parameters'] = model_parameter_lines
156
157
158 color_matrix_lines = self.get_color_matrix_lines(matrix_element)
159 replace_dict['color_matrix_lines'] = \
160 "\n ".join(color_matrix_lines)
161
162
163
164
165 jamp_lines = self.get_jamp_lines(matrix_element)
166 replace_dict['jamp_lines'] = "\n ".join(jamp_lines)
167
168
169 amp2_lines = self.get_amp2_lines(matrix_element,
170 self.config_maps.setdefault(ime, []))
171 replace_dict['amp2_lines'] = '\n '.join(amp2_lines)
172
173 method_file = open(os.path.join(_file_path, \
174 'iolibs/template_files/matrix_method_python.inc')).read()
175 method_file = method_file % replace_dict
176
177 self.matrix_methods[process_string] = method_file
178
179 return self.matrix_methods
180
182 """Return the Helicity matrix definition lines for this matrix element"""
183
184 helicity_line = "helicities = [ \\\n "
185 helicity_line_list = []
186
187 for helicities in matrix_element.get_helicity_matrix():
188 helicity_line_list.append("[" + ",".join(['%d'] * len(helicities)) % \
189 tuple(helicities) + "]")
190
191 return helicity_line + ",\n ".join(helicity_line_list) + "]"
192
193
195 """Return the denominator factor line for this matrix element"""
196
197 return "denominator = %d" % \
198 matrix_element.get_denominator_factor()
199
201 """Return the color matrix definition lines for this matrix element. Split
202 rows in chunks of size n."""
203
204 if not matrix_element.get('color_matrix'):
205 return ["denom = [1.]", "cf = [[1.]];"]
206 else:
207 color_denominators = matrix_element.get('color_matrix').\
208 get_line_denominators()
209 denom_string = "denom = [%s];" % \
210 ",".join(["%i" % denom for denom in color_denominators])
211
212 matrix_strings = []
213 my_cs = color.ColorString()
214 for index, denominator in enumerate(color_denominators):
215
216 num_list = matrix_element.get('color_matrix').\
217 get_line_numerators(index, denominator)
218
219 matrix_strings.append("%s" % \
220 ",".join(["%d" % i for i in num_list]))
221 matrix_string = "cf = [[" + \
222 "],\n [".join(matrix_strings) + "]];"
223 return [denom_string, matrix_string]
224
226 """Return the jamp = sum(fermionfactor * amp[i]) lines"""
227
228 res_list = []
229
230 for i, coeff_list in enumerate(matrix_element.get_color_amplitudes()):
231
232 res = "jamp[%d] = " % i
233
234
235
236 list_fracs = [abs(coefficient[0][1]) for coefficient in coeff_list]
237 common_factor = False
238 diff_fracs = list(set(list_fracs))
239 if len(diff_fracs) == 1 and abs(diff_fracs[0]) != 1:
240 common_factor = True
241 global_factor = diff_fracs[0]
242 res = res + '%s(' % coeff(1, global_factor, False, 0)
243
244 for (coefficient, amp_number) in coeff_list:
245 if common_factor:
246 res = res + "%samp[%d]" % (coeff(coefficient[0],
247 coefficient[1] / abs(coefficient[1]),
248 coefficient[2],
249 coefficient[3]),
250 amp_number - 1)
251 else:
252 res = res + "%samp[%d]" % (coeff(coefficient[0],
253 coefficient[1],
254 coefficient[2],
255 coefficient[3]),
256 amp_number - 1)
257
258 if common_factor:
259 res = res + ')'
260
261 res_list.append(res)
262
263 return res_list
264
266 """Return the amp2(i) = sum(amp for diag(i))^2 lines"""
267
268 ret_lines = []
269
270
271 vert_list = [max(diag.get_vertex_leg_numbers()) for diag in \
272 matrix_element.get('diagrams') if diag.get_vertex_leg_numbers()!=[]]
273 minvert = min(vert_list) if vert_list!=[] else 0
274
275 if config_map:
276
277
278
279 diagrams = matrix_element.get('diagrams')
280
281 config_to_diag_dict = {}
282 for idiag, diag in enumerate(matrix_element.get('diagrams')):
283 if config_map[idiag] == 0:
284 continue
285 try:
286 config_to_diag_dict[config_map[idiag]].append(idiag)
287 except KeyError:
288 config_to_diag_dict[config_map[idiag]] = [idiag]
289
290
291
292
293
294 for config in config_to_diag_dict.keys():
295
296 line = "self.amp2[%d]+=" % (config_to_diag_dict[config][0])
297
298 line += "+".join(["abs(amp[%(num)d]*amp[%(num)d].conjugate())" % \
299 {"num": a.get('number')-1} for a in \
300 sum([diagrams[idiag].get('amplitudes') for \
301 idiag in config_to_diag_dict[config]],
302 [])])
303 ret_lines.append(line)
304 ret_lines.sort()
305 else:
306 wf_dict = {}
307 vx_list = []
308 optimization = 0
309 for idiag, diag in enumerate(matrix_element.get('diagrams')):
310
311 if diag.get_vertex_leg_numbers()!=[] and \
312 max(diag.get_vertex_leg_numbers()) > minvert:
313 continue
314
315
316 line = "self.amp2[%d]+=" % (idiag)
317 line += "+".join(["abs(amp[%(num)d]*amp[%(num)d].conjugate())" % \
318 {"num": a.get('number')-1} for a in \
319 diag.get('amplitudes')])
320 ret_lines.append(line)
321
322 return ret_lines
323
325 """Return info lines for MG5, suitable to place at beginning of
326 Python files"""
327
328 info = misc.get_pkg_info()
329 info_lines = ""
330 if info and info.has_key('version') and info.has_key('date'):
331 info_lines = "# MadGraph5_aMC@NLO v. %s, %s\n" % \
332 (info['version'], info['date'])
333 info_lines = info_lines + \
334 " # By the MadGraph5_aMC@NLO Development Team\n" + \
335 " # Visit launchpad.net/madgraph5 and amcatnlo.web.cern.ch"
336 else:
337 info_lines = " # by MadGraph5_aMC@NLO\n" + \
338 " # By the MadGraph5_aMC@NLO Development Team\n" + \
339 " # Visit launchpad.net/madgraph5 and amcatnlo.web.cern.ch"
340
341 return info_lines
342
344 """Return info lines describing the processes for this matrix element"""
345
346 return"\n ".join([ "# " + process.nice_string().replace('\n', '\n# * ') \
347 for process in matrix_element.get('processes')])
348
349
351 """Return definitions for all model parameters used in this
352 matrix element"""
353
354
355 if aloha.complex_mass:
356 parameters = [(wf.get('mass') == 'ZERO' or wf.get('width')=='ZERO')
357 and wf.get('mass') or 'CMASS_%s' % wf.get('mass')
358 for wf in \
359 matrix_element.get_all_wavefunctions()]
360 parameters += [wf.get('mass') for wf in \
361 matrix_element.get_all_wavefunctions()]
362 else:
363 parameters = [wf.get('mass') for wf in \
364 matrix_element.get_all_wavefunctions()]
365 parameters += [wf.get('width') for wf in \
366 matrix_element.get_all_wavefunctions()]
367 parameters = list(set(parameters))
368 if 'ZERO' in parameters:
369 parameters.remove('ZERO')
370
371
372
373
374 couplings = list(set([c.replace('-', '') for func \
375 in matrix_element.get_all_wavefunctions() + \
376 matrix_element.get_all_amplitudes() for c in func.get('coupling')
377 if func.get('mothers') ]))
378
379 return "\n ".join([\
380 "%(param)s = model.get(\'parameter_dict\')[\"%(param)s\"]"\
381 % {"param": param} for param in parameters]) + \
382 "\n " + "\n ".join([\
383 "%(coup)s = model.get(\'coupling_dict\')[\"%(coup)s\"]"\
384 % {"coup": coup} for coup in couplings])
385
386
387
388
389
390 -def coeff(ff_number, frac, is_imaginary, Nc_power, Nc_value=3):
391 """Returns a nicely formatted string for the coefficients in JAMP lines"""
392
393 total_coeff = ff_number * frac * fractions.Fraction(Nc_value) ** Nc_power
394
395 if total_coeff == 1:
396 if is_imaginary:
397 return '+complex(0,1)*'
398 else:
399 return '+'
400 elif total_coeff == -1:
401 if is_imaginary:
402 return '-complex(0,1)*'
403 else:
404 return '-'
405
406 res_str = '%+i.' % total_coeff.numerator
407
408 if total_coeff.denominator != 1:
409
410 res_str = res_str + '/%i.' % total_coeff.denominator
411
412 if is_imaginary:
413 res_str = res_str + '*complex(0,1)'
414
415 return res_str + '*'
416