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 minvert = min([max(diag.get_vertex_leg_numbers()) for diag in \ 271 matrix_element.get('diagrams')]) 272 if config_map: 273 # In this case, we need to sum up all amplitudes that have 274 # identical topologies, as given by the config_map (which 275 # gives the topology/config for each of the diagrams 276 diagrams = matrix_element.get('diagrams') 277 # Combine the diagrams with identical topologies 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 # Write out the AMP2s summing squares of amplitudes belonging 287 # to eiher the same diagram or different diagrams with 288 # identical propagator properties. Note that we need to use 289 # AMP2 number corresponding to the first diagram number used 290 # for that AMP2. 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 # Ignore any diagrams with 4-particle vertices. 308 if max(diag.get_vertex_leg_numbers()) > minvert: 309 continue 310 # Now write out the expression for AMP2, meaning the sum of 311 # squared amplitudes belonging to the same diagram 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
320 - def get_mg5_info_lines(self):
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
339 - def get_process_info_lines(self, matrix_element):
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
346 - def get_model_parameter_lines(self, matrix_element):
347 """Return definitions for all model parameters used in this 348 matrix element""" 349 350 # Get all masses and widths used 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 # Get all couplings used 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 # Global helper methods 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 # Check if total_coeff is an integer 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