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 minvert = min([max(diag.get_vertex_leg_numbers()) for diag in \
271 matrix_element.get('diagrams')])
272 if config_map:
273
274
275
276 diagrams = matrix_element.get('diagrams')
277
278 config_to_diag_dict = {}
279 for idiag, diag in enumerate(matrix_element.get('diagrams')):
280 if config_map[idiag] == 0:
281 continue
282 try:
283 config_to_diag_dict[config_map[idiag]].append(idiag)
284 except KeyError:
285 config_to_diag_dict[config_map[idiag]] = [idiag]
286
287
288
289
290
291 for config in config_to_diag_dict.keys():
292
293 line = "self.amp2[%d]+=" % (config_to_diag_dict[config][0])
294
295 line += "+".join(["abs(amp[%(num)d]*amp[%(num)d].conjugate())" % \
296 {"num": a.get('number')-1} for a in \
297 sum([diagrams[idiag].get('amplitudes') for \
298 idiag in config_to_diag_dict[config]],
299 [])])
300 ret_lines.append(line)
301 ret_lines.sort()
302 else:
303 wf_dict = {}
304 vx_list = []
305 optimization = 0
306 for idiag, diag in enumerate(matrix_element.get('diagrams')):
307
308 if max(diag.get_vertex_leg_numbers()) > minvert:
309 continue
310
311
312 line = "self.amp2[%d]+=" % (idiag)
313 line += "+".join(["abs(amp[%(num)d]*amp[%(num)d].conjugate())" % \
314 {"num": a.get('number')-1} for a in \
315 diag.get('amplitudes')])
316 ret_lines.append(line)
317
318 return ret_lines
319
321 """Return info lines for MG5, suitable to place at beginning of
322 Python files"""
323
324 info = misc.get_pkg_info()
325 info_lines = ""
326 if info and info.has_key('version') and info.has_key('date'):
327 info_lines = "# MadGraph5_aMC@NLO v. %s, %s\n" % \
328 (info['version'], info['date'])
329 info_lines = info_lines + \
330 " # By the MadGraph5_aMC@NLO Development Team\n" + \
331 " # Visit launchpad.net/madgraph5 and amcatnlo.web.cern.ch"
332 else:
333 info_lines = " # by MadGraph5_aMC@NLO\n" + \
334 " # By the MadGraph5_aMC@NLO Development Team\n" + \
335 " # Visit launchpad.net/madgraph5 and amcatnlo.web.cern.ch"
336
337 return info_lines
338
340 """Return info lines describing the processes for this matrix element"""
341
342 return"\n ".join([ "# " + process.nice_string().replace('\n', '\n# * ') \
343 for process in matrix_element.get('processes')])
344
345
347 """Return definitions for all model parameters used in this
348 matrix element"""
349
350
351 if aloha.complex_mass:
352 parameters = [(wf.get('mass') == 'ZERO' or wf.get('width')=='ZERO')
353 and wf.get('mass') or 'CMASS_%s' % wf.get('mass')
354 for wf in \
355 matrix_element.get_all_wavefunctions()]
356 parameters += [wf.get('mass') for wf in \
357 matrix_element.get_all_wavefunctions()]
358 else:
359 parameters = [wf.get('mass') for wf in \
360 matrix_element.get_all_wavefunctions()]
361 parameters += [wf.get('width') for wf in \
362 matrix_element.get_all_wavefunctions()]
363 parameters = list(set(parameters))
364 if 'ZERO' in parameters:
365 parameters.remove('ZERO')
366
367
368
369
370 couplings = list(set([c.replace('-', '') for func \
371 in matrix_element.get_all_wavefunctions() + \
372 matrix_element.get_all_amplitudes() for c in func.get('coupling')
373 if func.get('mothers') ]))
374
375 return "\n ".join([\
376 "%(param)s = model.get(\'parameter_dict\')[\"%(param)s\"]"\
377 % {"param": param} for param in parameters]) + \
378 "\n " + "\n ".join([\
379 "%(coup)s = model.get(\'coupling_dict\')[\"%(coup)s\"]"\
380 % {"coup": coup} for coup in couplings])
381
382
383
384
385
386 -def coeff(ff_number, frac, is_imaginary, Nc_power, Nc_value=3):
387 """Returns a nicely formatted string for the coefficients in JAMP lines"""
388
389 total_coeff = ff_number * frac * fractions.Fraction(Nc_value) ** Nc_power
390
391 if total_coeff == 1:
392 if is_imaginary:
393 return '+complex(0,1)*'
394 else:
395 return '+'
396 elif total_coeff == -1:
397 if is_imaginary:
398 return '-complex(0,1)*'
399 else:
400 return '-'
401
402 res_str = '%+i.' % total_coeff.numerator
403
404 if total_coeff.denominator != 1:
405
406 res_str = res_str + '/%i.' % total_coeff.denominator
407
408 if is_imaginary:
409 res_str = res_str + '*complex(0,1)'
410
411 return res_str + '*'
412