Source Code for Module madgraph.iolibs.export_python

  1  ################################################################################ 
  2  # 
  3  # Copyright (c) 2009 The MadGraph Development team and Contributors 
  4  # 
  5  # This file is a part of the MadGraph 5 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 MadGraph license which should accompany this  
 10  # distribution. 
 11  # 
 12  # For more information, please visit: http://madgraph.phys.ucl.ac.be 
 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   
 27  import madgraph.core.color_algebra as color 
 28  import madgraph.core.helas_objects as helas_objects 
 29  import madgraph.iolibs.drawing_eps as draw 
 30  import madgraph.iolibs.files as files 
 31  import madgraph.iolibs.helas_call_writers as helas_call_writers 
 32  import madgraph.iolibs.file_writers as writers 
 33  import madgraph.iolibs.template_files as Template 
 34  import madgraph.iolibs.ufo_expression_parsers as parsers 
 35  import madgraph.iolibs.group_subprocs as group_subprocs 
 36  from madgraph import MadGraph5Error, MG5DIR 
 37   
 38  import madgraph.various.misc as misc 
 39   
 40  import aloha.create_aloha as create_aloha 
 41  import aloha.aloha_writers as aloha_writers 
 42   
 43  _file_path = os.path.split(os.path.dirname(os.path.realpath(__file__)))[0] + '/' 
 44  logger = logging.getLogger('madgraph.export_python') 
 45   
 46   
 47  #=============================================================================== 
 48  # ProcessExporterPython 
 49  #=============================================================================== 
 50 -class ProcessExporterPython(object): 
 51      """Class to take care of exporting a set of matrix elements to 
 52      Python format.""" 
 53   
 54 -    class ProcessExporterPythonError(Exception): 
 55          pass 
 56       
 57 -    def __init__(self, matrix_elements, python_helas_call_writer): 
 58          """Initiate with matrix elements, helas call writer. 
 59          Generate the process matrix element functions as strings.""" 
 60   
 61          self.config_maps = {} 
 62          if isinstance(matrix_elements, helas_objects.HelasMultiProcess): 
 63              self.matrix_elements = matrix_elements.get('matrix_elements') 
 64          elif isinstance(matrix_elements, 
 65                          group_subprocs.SubProcessGroup): 
 66              self.config_maps = matrix_elements.get('diagram_maps') 
 67              self.matrix_elements = matrix_elements.get('matrix_elements') 
 68          elif isinstance(matrix_elements, 
 69                          helas_objects.HelasMatrixElementList): 
 70              self.matrix_elements = matrix_elements 
 71          elif isinstance(matrix_elements, 
 72                          helas_objects.HelasMatrixElement): 
 73              self.matrix_elements = helas_objects.HelasMatrixElementList(\ 
 74                                                       [matrix_elements]) 
 75          if not self.matrix_elements: 
 76              raise MadGraph5Error("No matrix elements to export") 
 77   
 78          self.model = self.matrix_elements[0].get('processes')[0].get('model') 
 79   
 80          self.helas_call_writer = python_helas_call_writer 
 81   
 82          if not isinstance(self.helas_call_writer, helas_call_writers.PythonUFOHelasCallWriter): 
 83              raise ProcessExporterPythonError, \ 
 84                  "helas_call_writer not PythonUFOHelasCallWriter" 
 85   
 86          self.matrix_methods = {} 
 87   
 88      # Methods for generation of process file strings in Python 
 89   
 90      #=========================================================================== 
 91      # write_python_process_cc_file 
 92      #=========================================================================== 
 93 -    def get_python_matrix_methods(self, gauge_check=False): 
 94          """Write the matrix element calculation method for the processes""" 
 95   
 96          replace_dict = {} 
 97   
 98          # Extract version number and date from VERSION file 
 99          info_lines = self.get_mg5_info_lines() 
100          replace_dict['info_lines'] = info_lines 
101   
102          for ime, matrix_element in enumerate(self.matrix_elements): 
103              process_string = matrix_element.get('processes')[0].shell_string() 
104              if process_string in self.matrix_methods: 
105                  continue 
106   
107              replace_dict['process_string'] = process_string 
108   
109              # Extract number of external particles 
110              (nexternal, ninitial) = matrix_element.get_nexternal_ninitial() 
111              replace_dict['nexternal'] = nexternal 
112   
113              # Extract ncomb 
114              ncomb = matrix_element.get_helicity_combinations() 
115              replace_dict['ncomb'] = ncomb 
116   
117              # Extract helicity lines 
118              helicity_lines = self.get_helicity_matrix(matrix_element) 
119              replace_dict['helicity_lines'] = helicity_lines 
120   
121              # Extract overall denominator 
122              # Averaging initial state color, spin, and identical FS particles 
123              den_factor_line = self.get_den_factor_line(matrix_element) 
124              replace_dict['den_factor_line'] = den_factor_line 
125   
126              # Extract process info lines for all processes 
127              process_lines = self.get_process_info_lines(matrix_element) 
128              replace_dict['process_lines'] = process_lines 
129           
130              # Extract ngraphs 
131              ngraphs = matrix_element.get_number_of_amplitudes() 
132              replace_dict['ngraphs'] = ngraphs 
133   
134              # Extract ndiags 
135              ndiags = len(matrix_element.get('diagrams')) 
136              replace_dict['ndiags'] = ndiags 
137   
138              # Extract nwavefuncs 
139              nwavefuncs = matrix_element.get_number_of_wavefunctions() 
140              replace_dict['nwavefuncs'] = nwavefuncs 
141   
142              # Extract ncolor 
143              ncolor = max(1, len(matrix_element.get('color_basis'))) 
144              replace_dict['ncolor'] = ncolor 
145   
146              # Extract model parameter lines 
147              model_parameter_lines = \ 
148                                   self.get_model_parameter_lines(matrix_element) 
149              replace_dict['model_parameters'] = model_parameter_lines 
150   
151              # Extract color data lines 
152              color_matrix_lines = self.get_color_matrix_lines(matrix_element) 
153              replace_dict['color_matrix_lines'] = \ 
154                                                 "\n        ".join(color_matrix_lines) 
155   
156              # Extract helas calls 
157              helas_calls = self.helas_call_writer.get_matrix_element_calls(\ 
158                                                      matrix_element, gauge_check) 
159              replace_dict['helas_calls'] = "\n        ".join(helas_calls) 
160   
161              # Extract JAMP lines 
162              jamp_lines = self.get_jamp_lines(matrix_element) 
163              replace_dict['jamp_lines'] = "\n        ".join(jamp_lines) 
164   
165              # Extract amp2 lines 
166              amp2_lines = self.get_amp2_lines(matrix_element, 
167                                          self.config_maps.setdefault(ime, [])) 
168              replace_dict['amp2_lines'] = '\n        '.join(amp2_lines) 
169   
170              method_file = open(os.path.join(_file_path, \ 
171                         'iolibs/template_files/matrix_method_python.inc')).read() 
172              method_file = method_file % replace_dict 
173   
174              self.matrix_methods[process_string] = method_file 
175   
176          return self.matrix_methods 
177   
178 -    def get_helicity_matrix(self, matrix_element): 
179          """Return the Helicity matrix definition lines for this matrix element""" 
180   
181          helicity_line = "helicities = [ \\\n        " 
182          helicity_line_list = [] 
183   
184          for helicities in matrix_element.get_helicity_matrix(): 
185              helicity_line_list.append("[" + ",".join(['%d'] * len(helicities)) % \ 
186                                        tuple(helicities) + "]") 
187               
188          return helicity_line + ",\n        ".join(helicity_line_list) + "]" 
189   
190   
191 -    def get_den_factor_line(self, matrix_element): 
192          """Return the denominator factor line for this matrix element""" 
193   
194          return "denominator = %d" % \ 
195                 matrix_element.get_denominator_factor() 
196   
197 -    def get_color_matrix_lines(self, matrix_element): 
198          """Return the color matrix definition lines for this matrix element. Split 
199          rows in chunks of size n.""" 
200   
201          if not matrix_element.get('color_matrix'): 
202              return ["denom = [1.]", "cf = [[1.]];"] 
203          else: 
204              color_denominators = matrix_element.get('color_matrix').\ 
205                                                   get_line_denominators() 
206              denom_string = "denom = [%s];" % \ 
207                             ",".join(["%i" % denom for denom in color_denominators]) 
208   
209              matrix_strings = [] 
210              my_cs = color.ColorString() 
211              for index, denominator in enumerate(color_denominators): 
212                  # Then write the numerators for the matrix elements 
213                  num_list = matrix_element.get('color_matrix').\ 
214                                              get_line_numerators(index, denominator) 
215   
216                  matrix_strings.append("%s" % \ 
217                                       ",".join(["%d" % i for i in num_list])) 
218              matrix_string = "cf = [[" + \ 
219                              "],\n        [".join(matrix_strings) + "]];" 
220              return [denom_string, matrix_string] 
221   
222 -    def get_jamp_lines(self, matrix_element): 
223          """Return the jamp = sum(fermionfactor * amp[i]) lines""" 
224   
225          res_list = [] 
226   
227          for i, coeff_list in enumerate(matrix_element.get_color_amplitudes()): 
228   
229              res = "jamp[%d] = " % i 
230   
231              # Optimization: if all contributions to that color basis element have 
232              # the same coefficient (up to a sign), put it in front 
233              list_fracs = [abs(coefficient[0][1]) for coefficient in coeff_list] 
234              common_factor = False 
235              diff_fracs = list(set(list_fracs)) 
236              if len(diff_fracs) == 1 and abs(diff_fracs[0]) != 1: 
237                  common_factor = True 
238                  global_factor = diff_fracs[0] 
239                  res = res + '%s(' % coeff(1, global_factor, False, 0) 
240   
241              for (coefficient, amp_number) in coeff_list: 
242                  if common_factor: 
243                      res = res + "%samp[%d]" % (coeff(coefficient[0], 
244                                                 coefficient[1] / abs(coefficient[1]), 
245                                                 coefficient[2], 
246                                                 coefficient[3]), 
247                                                 amp_number - 1) 
248                  else: 
249                      res = res + "%samp[%d]" % (coeff(coefficient[0], 
250                                                 coefficient[1], 
251                                                 coefficient[2], 
252                                                 coefficient[3]), 
253                                                 amp_number - 1) 
254   
255              if common_factor: 
256                  res = res + ')' 
257   
258              res_list.append(res) 
259   
260          return res_list 
261   
262 -    def get_amp2_lines(self, matrix_element, config_map = []): 
263          """Return the amp2(i) = sum(amp for diag(i))^2 lines""" 
264   
265          ret_lines = [] 
266          # Get minimum legs in a vertex 
267          minvert = min([max(diag.get_vertex_leg_numbers()) for diag in \ 
268                         matrix_element.get('diagrams')]) 
269          if config_map: 
270              # In this case, we need to sum up all amplitudes that have 
271              # identical topologies, as given by the config_map (which 
272              # gives the topology/config for each of the diagrams 
273              diagrams = matrix_element.get('diagrams') 
274              # Combine the diagrams with identical topologies 
275              config_to_diag_dict = {} 
276              for idiag, diag in enumerate(matrix_element.get('diagrams')): 
277                  if config_map[idiag] == 0: 
278                      continue 
279                  try: 
280                      config_to_diag_dict[config_map[idiag]].append(idiag) 
281                  except KeyError: 
282                      config_to_diag_dict[config_map[idiag]] = [idiag] 
283              # Write out the AMP2s summing squares of amplitudes belonging 
284              # to eiher the same diagram or different diagrams with 
285              # identical propagator properties.  Note that we need to use 
286              # AMP2 number corresponding to the first diagram number used 
287              # for that AMP2. 
288              for config in config_to_diag_dict.keys(): 
289   
290                  line = "self.amp2[%d]+=" % (config_to_diag_dict[config][0]) 
291   
292                  line += "+".join(["abs(amp[%(num)d]*amp[%(num)d].conjugate())" % \ 
293                                    {"num": a.get('number')-1} for a in \ 
294                                    sum([diagrams[idiag].get('amplitudes') for \ 
295                                         idiag in config_to_diag_dict[config]], 
296                                        [])]) 
297                  ret_lines.append(line) 
298              ret_lines.sort() 
299          else: 
300              wf_dict = {} 
301              vx_list = [] 
302              optimization = 0 
303              for idiag, diag in enumerate(matrix_element.get('diagrams')): 
304                  # Ignore any diagrams with 4-particle vertices. 
305                  if max(diag.get_vertex_leg_numbers()) > minvert: 
306                      continue 
307                  # Now write out the expression for AMP2, meaning the sum of 
308                  # squared amplitudes belonging to the same diagram 
309                  line = "self.amp2[%d]+=" % (idiag) 
310                  line += "+".join(["abs(amp[%(num)d]*amp[%(num)d].conjugate())" % \ 
311                                    {"num": a.get('number')-1} for a in \ 
312                                    diag.get('amplitudes')]) 
313                  ret_lines.append(line) 
314   
315          return ret_lines 
316   
317 -    def get_mg5_info_lines(self): 
318          """Return info lines for MG5, suitable to place at beginning of 
319          Python files""" 
320   
321          info = misc.get_pkg_info() 
322          info_lines = "" 
323          if info and info.has_key('version') and  info.has_key('date'): 
324              info_lines = "#  MadGraph 5 v. %s, %s\n" % \ 
325                           (info['version'], info['date']) 
326              info_lines = info_lines + \ 
327                           "        #  By the MadGraph Development Team\n" + \ 
328                           "        #  Please visit us at https://launchpad.net/madgraph5" 
329          else: 
330              info_lines = "        #  by MadGraph 5\n" + \ 
331                           "        #  By the MadGraph Development Team\n" + \ 
332                           "        #  Please visit us at https://launchpad.net/madgraph5"         
333   
334          return info_lines 
335   
336 -    def get_process_info_lines(self, matrix_element): 
337          """Return info lines describing the processes for this matrix element""" 
338   
339          return"\n        ".join([ "# " + process.nice_string().replace('\n', '\n# * ') \ 
340                           for process in matrix_element.get('processes')]) 
341   
342   
343 -    def get_model_parameter_lines(self, matrix_element): 
344          """Return definitions for all model parameters used in this 
345          matrix element""" 
346   
347          # Get all masses and widths used 
348          parameters = [wf.get('mass') for wf in \ 
349                        matrix_element.get_all_wavefunctions()] 
350          parameters += [wf.get('width') for wf in \ 
351                         matrix_element.get_all_wavefunctions()] 
352          parameters = list(set(parameters)) 
353          if 'ZERO' in parameters: 
354              parameters.remove('ZERO') 
355   
356          # Get all couplings used 
357   
358           
359          couplings = list(set([c.replace('-', '') for func \ 
360                                in matrix_element.get_all_wavefunctions() + \ 
361                                matrix_element.get_all_amplitudes() for c in func.get('coupling') 
362                                if func.get('mothers') ])) 
363           
364          return "\n        ".join([\ 
365                           "%(param)s = model.get(\'parameter_dict\')[\"%(param)s\"]"\ 
366                           % {"param": param} for param in parameters]) + \ 
367                 "\n        " + "\n        ".join([\ 
368                           "%(coup)s = model.get(\'coupling_dict\')[\"%(coup)s\"]"\ 
369                                % {"coup": coup} for coup in couplings]) 
370   
371  #=============================================================================== 
372  # Global helper methods 
373  #=============================================================================== 
374   
375 -def coeff(ff_number, frac, is_imaginary, Nc_power, Nc_value=3): 
376      """Returns a nicely formatted string for the coefficients in JAMP lines""" 
377   
378      total_coeff = ff_number * frac * fractions.Fraction(Nc_value) ** Nc_power 
379   
380      if total_coeff == 1: 
381          if is_imaginary: 
382              return '+complex(0,1)*' 
383          else: 
384              return '+' 
385      elif total_coeff == -1: 
386          if is_imaginary: 
387              return '-complex(0,1)*' 
388          else: 
389              return '-' 
390   
391      res_str = '%+i.' % total_coeff.numerator 
392   
393      if total_coeff.denominator != 1: 
394          # Check if total_coeff is an integer 
395          res_str = res_str + '/%i.' % total_coeff.denominator 
396   
397      if is_imaginary: 
398          res_str = res_str + '*complex(0,1)' 
399   
400      return res_str + '*' 
401