Package madgraph :: Package iolibs :: Module export_python
[hide private]
[frames] | no frames]

Source Code for Module madgraph.iolibs.export_python

  1  ################################################################################ 
  2  # 
  3  # Copyright (c) 2009 The MadGraph5_aMC@NLO Development team and Contributors 
  4  # 
  5  # This file is a part of the MadGraph5_aMC@NLO project, an application which  
  6  # automatically generates Feynman diagrams and matrix elements for arbitrary 
  7  # high-energy processes in the Standard Model and beyond. 
  8  # 
  9  # It is subject to the MadGraph5_aMC@NLO license which should accompany this  
 10  # distribution. 
 11  # 
 12  # For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch 
 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  # ProcessExporterPython 
 50  #=============================================================================== 
51 -class ProcessExporterPython(object):
52 """Class to take care of exporting a set of matrix elements to 53 Python format.""" 54
55 - class ProcessExporterPythonError(Exception):
56 pass
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 # Methods for generation of process file strings in Python 90 91 #=========================================================================== 92 # write_python_process_cc_file 93 #===========================================================================
94 - def get_python_matrix_methods(self, gauge_check=False):
95 """Write the matrix element calculation method for the processes""" 96 97 replace_dict = {} 98 99 # Extract version number and date from VERSION file 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 # Extract number of external particles 111 (nexternal, ninitial) = matrix_element.get_nexternal_ninitial() 112 replace_dict['nexternal'] = nexternal 113 114 # Extract ncomb 115 ncomb = matrix_element.get_helicity_combinations() 116 replace_dict['ncomb'] = ncomb 117 118 # Extract helicity lines 119 helicity_lines = self.get_helicity_matrix(matrix_element) 120 replace_dict['helicity_lines'] = helicity_lines 121 122 # Extract overall denominator 123 # Averaging initial state color, spin, and identical FS particles 124 den_factor_line = self.get_den_factor_line(matrix_element) 125 replace_dict['den_factor_line'] = den_factor_line 126 127 # Extract process info lines for all processes 128 process_lines = self.get_process_info_lines(matrix_element) 129 replace_dict['process_lines'] = process_lines 130 131 # Extract ngraphs 132 ngraphs = matrix_element.get_number_of_amplitudes() 133 replace_dict['ngraphs'] = ngraphs 134 135 # Extract ndiags 136 ndiags = len(matrix_element.get('diagrams')) 137 replace_dict['ndiags'] = ndiags 138 139 # Extract helas calls 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 # Extract nwavefuncs 145 nwavefuncs = matrix_element.get_number_of_wavefunctions() 146 replace_dict['nwavefuncs'] = nwavefuncs 147 148 # Extract ncolor 149 ncolor = max(1, len(matrix_element.get('color_basis'))) 150 replace_dict['ncolor'] = ncolor 151 152 # Extract model parameter lines 153 model_parameter_lines = \ 154 self.get_model_parameter_lines(matrix_element) 155 replace_dict['model_parameters'] = model_parameter_lines 156 157 # Extract color data lines 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 # Extract JAMP lines 165 jamp_lines = self.get_jamp_lines(matrix_element) 166 replace_dict['jamp_lines'] = "\n ".join(jamp_lines) 167 168 # Extract amp2 lines 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
181 - def get_helicity_matrix(self, matrix_element):
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
194 - def get_den_factor_line(self, matrix_element):
195 """Return the denominator factor line for this matrix element""" 196 197 return "denominator = %d" % \ 198 matrix_element.get_denominator_factor()
199
200 - def get_color_matrix_lines(self, matrix_element):
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 # Then write the numerators for the matrix elements 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
225 - def get_jamp_lines(self, matrix_element):
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 # Optimization: if all contributions to that color basis element have 235 # the same coefficient (up to a sign), put it in front 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
265 - def get_amp2_lines(self, matrix_element, config_map = []):
266 """Return the amp2(i) = sum(amp for diag(i))^2 lines""" 267 268 ret_lines = [] 269 # Get minimum legs in a vertex 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 # In this case, we need to sum up all amplitudes that have 277 # identical topologies, as given by the config_map (which 278 # gives the topology/config for each of the diagrams 279 diagrams = matrix_element.get('diagrams') 280 # Combine the diagrams with identical topologies 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 # Write out the AMP2s summing squares of amplitudes belonging 290 # to eiher the same diagram or different diagrams with 291 # identical propagator properties. Note that we need to use 292 # AMP2 number corresponding to the first diagram number used 293 # for that AMP2. 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 # Ignore any diagrams with 4-particle vertices. 311 if diag.get_vertex_leg_numbers()!=[] and \ 312 max(diag.get_vertex_leg_numbers()) > minvert: 313 continue 314 # Now write out the expression for AMP2, meaning the sum of 315 # squared amplitudes belonging to the same diagram 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
324 - def get_mg5_info_lines(self):
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
343 - def get_process_info_lines(self, matrix_element):
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
350 - def get_model_parameter_lines(self, matrix_element):
351 """Return definitions for all model parameters used in this 352 matrix element""" 353 354 # Get all masses and widths used 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 # Get all couplings used 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 # Global helper methods 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 # Check if total_coeff is an integer 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