Source Code for Module aloha.aloha_writers

   1  try: 
   2      import madgraph.iolibs.file_writers as writers  
   3  except: 
   4      import aloha.file_writers as writers 
   5       
   6  import os 
   7  import re  
   8  from numbers import Number 
   9   
  10 -class WriteALOHA:  
  11      """ Generic writing functions """  
  12       
  13      power_symbol = '**' 
  14      change_var_format = str 
  15      change_number_format = str 
  16      extension = '' 
  17      type_to_variable = {2:'F',3:'V',5:'T',1:'S'} 
  18      type_to_size = {'S':3, 'T':18, 'V':6, 'F':6} 
  19       
  20               
  21 -    def __init__(self, abstract_routine, dirpath): 
  22   
  23   
  24          name = get_routine_name(abstract = abstract_routine) 
  25          if dirpath: 
  26              self.dir_out = dirpath 
  27              self.out_path = os.path.join(dirpath, name + self.extension) 
  28          else: 
  29              self.out_path = None 
  30              self.dir_out = None 
  31   
  32          self.obj = abstract_routine.expr 
  33          self.particles =  [self.type_to_variable[spin] for spin in \ 
  34                            abstract_routine.spins] 
  35          self.namestring = name 
  36          self.abstractname = abstract_routine.name 
  37          self.comment = abstract_routine.infostr 
  38          self.offshell = abstract_routine.outgoing 
  39            
  40          self.symmetries = abstract_routine.symmetries 
  41          self.tag = abstract_routine.tag 
  42   
  43          self.outgoing = self.offshell 
  44          if 'C%s' %((self.outgoing + 1) // 2) in self.tag: 
  45              #flip the outgoing tag if in conjugate 
  46              self.outgoing = self.outgoing + self.outgoing % 2 - (self.outgoing +1) % 2 
  47               
  48   
  49   
  50   
  51   
  52          #prepare the necessary object 
  53          self.collect_variables() # Look for the different variables 
  54          self.make_all_lists()    # Compute the expression for the call ordering 
  55                                   #the definition of objects,.. 
  56                                    
  57                                    
  58                                    
  59                                    
  60 -    def pass_to_HELAS(self, indices, start=0): 
  61          """find the Fortran HELAS position for the list of index"""  
  62           
  63           
  64          if len(indices) == 1: 
  65              return indices[0] + start 
  66           
  67          ind_name = self.obj.numerator.lorentz_ind  
  68          if ind_name == ['I3', 'I2']: 
  69              return  4 * indices[1] + indices[0] + start  
  70          elif len(indices) == 2:  
  71              return  4 * indices[0] + indices[1] + start  
  72          else: 
  73              raise Exception, 'WRONG CONTRACTION OF LORENTZ OBJECT for routine %s: %s' \ 
  74                      % (self.namestring, indices)                                  
  75                                    
  76 -    def collect_variables(self): 
  77          """Collects Momenta,Mass,Width into lists""" 
  78            
  79          MomentaList = set() 
  80          OverMList = set() 
  81          for elem in self.obj.tag: 
  82              if elem.startswith('P'): 
  83                  MomentaList.add(elem) 
  84              elif elem.startswith('O'): 
  85                  OverMList.add(elem)  
  86   
  87          MomentaList = list(MomentaList) 
  88          OverMList = list(OverMList) 
  89           
  90          self.collected = {'momenta':MomentaList, 'om':OverMList} 
  91           
  92          return self.collected 
  93   
  94 -    def define_header(self):  
  95          """ Prototype for language specific header"""  
  96          pass 
  97   
  98 -    def define_content(self):  
  99          """Prototype for language specific body"""  
 100          pass 
 101       
 102 -    def define_foot(self): 
 103          """Prototype for language specific footer""" 
 104          pass 
 105       
 106 -    def write_indices_part(self, indices, obj):  
 107          """Routine for making a string out of indices objects""" 
 108           
 109          text = 'output(%s)' % indices 
 110          return text                  
 111           
 112 -    def write_obj(self, obj): 
 113          """Calls the appropriate writing routine""" 
 114           
 115          try: 
 116              vartype = obj.vartype 
 117          except: 
 118              return self.change_number_format(obj) 
 119   
 120          if vartype == 2 : # MultVariable 
 121              return self.write_obj_Mult(obj) 
 122          elif not vartype: # Variable 
 123              return self.write_obj_Var(obj) 
 124          elif vartype == 1 : # AddVariable 
 125              return self.write_obj_Add(obj) 
 126          elif vartype == 5: # ConstantObject 
 127              return self.change_number_format(obj.value) 
 128          else:  
 129              raise Exception('Warning unknown object: %s' % obj.vartype) 
 130   
 131 -    def write_obj_Mult(self, obj): 
 132          """Turn a multvariable into a string"""  
 133          mult_list = [self.write_obj(factor) for factor in obj]  
 134          text = '('  
 135          if obj.prefactor != 1: 
 136              if obj.prefactor != -1: 
 137                  text = self.change_number_format(obj.prefactor) + '*' + text  
 138              else: 
 139                  text = '-' + text 
 140          return text + '*'.join(mult_list) + ')' 
 141       
 142 -    def write_obj_Add(self, obj): 
 143          """Turns addvariable into a string""" 
 144          mult_list = [self.write_obj(factor) for factor in obj] 
 145          prefactor = '' 
 146          if obj.prefactor == 1: 
 147              prefactor = '' 
 148          elif obj.prefactor == -1: 
 149              prefactor = '-' 
 150          else: 
 151              prefactor = '%s*' % self.change_number_format(obj.prefactor) 
 152   
 153          return '(%s %s)' % (prefactor, '+'.join(mult_list)) 
 154   
 155           
 156 -    def write_obj_Var(self, obj): 
 157          text = '' 
 158          if obj.prefactor != 1: 
 159              if obj.prefactor != -1:  
 160                  text = self.change_number_format(obj.prefactor) + '*' + text 
 161              else: 
 162                  text = '-' + text 
 163          text += self.change_var_format(obj.variable) 
 164          if obj.power != 1: 
 165              text = text + self.power_symbol + str(obj.power) 
 166          return text 
 167   
 168 -    def make_all_lists(self): 
 169          """ Make all the list for call ordering, conservation impulsion,  
 170          basic declaration""" 
 171           
 172          DeclareList = self.make_declaration_list() 
 173          CallList = self.make_call_list() 
 174          MomentumConserve = self.make_momentum_conservation() 
 175   
 176          self.calllist =  {'CallList':CallList,'DeclareList':DeclareList, \ 
 177                             'Momentum':MomentumConserve} 
 178   
 179       
 180 -    def make_call_list(self, outgoing=None): 
 181          """find the way to write the call of the functions""" 
 182   
 183          if outgoing is None: 
 184              outgoing = self.offshell 
 185   
 186          call_arg = [] #incoming argument of the routine 
 187   
 188          conjugate = [2*(int(c[1:])-1) for c in self.tag if c[0] == 'C'] 
 189           
 190          for index,spin in enumerate(self.particles): 
 191              if self.offshell == index + 1: 
 192                  continue 
 193               
 194              if index in conjugate: 
 195                  index2, spin2 = index+1, self.particles[index+1] 
 196                  call_arg.append('%s%d' % (spin2, index2 +1))  
 197                  #call_arg.append('%s%d' % (spin, index +1))  
 198              elif index-1 in conjugate: 
 199                  index2, spin2 = index-1, self.particles[index-1] 
 200                  call_arg.append('%s%d' % (spin2, index2 +1))  
 201              else: 
 202                  call_arg.append('%s%d' % (spin, index +1))  
 203                   
 204          return call_arg 
 205   
 206  #    def reorder_call_list(self, call_list, old, new): 
 207  #        """ restore the correct order for symmetries """ 
 208  #        raise 
 209  #        #spins = self.particles 
 210  #        #assert(0 < old < new) 
 211  #        #old, new = old -1, new -1 # pass in real position in particles list 
 212  #        #assert(spins[old] == spins[new]) 
 213  #        #spin =spins[old] 
 214  #         
 215  #        new_call = call_list[:] 
 216  #        #val = new_call.pop(old) 
 217  #        #new_call.insert(new - 1, val) 
 218  #        return new_call 
 219       
 220           
 221 -    def make_momentum_conservation(self): 
 222          """ compute the sign for the momentum conservation """ 
 223           
 224          if not self.offshell: 
 225              return [] 
 226          # How Convert  sign to a string 
 227          sign_dict = {1: '+', -1: '-'} 
 228          # help data  
 229          momentum_conserve = [] 
 230          nb_fermion =0 
 231           
 232          #compute global sign 
 233          if not self.offshell % 2 and self.particles[self.offshell -1] == 'F':  
 234              global_sign = 1 
 235          else: 
 236              global_sign = -1 
 237           
 238          flipped = [2*(int(c[1:])-1) for c in self.tag if c.startswith('C')] 
 239  #        if self.offshell % 2: 
 240  #             not_flip = self.offshell - 1 
 241  #             if not_flip in flipped: 
 242  #                 flipped.remove(not_flip) 
 243  #        else: 
 244  #             not_flip = self.offshell - 2 
 245  #             if not_flip in flipped: 
 246  #                 flipped.remove(not_flip)            
 247           
 248          for index, spin in enumerate(self.particles):  
 249              assert(spin in ['S','F','V','T'])   
 250         
 251              #compute the sign 
 252              if spin != 'F': 
 253                  sign = -1 * global_sign 
 254              elif nb_fermion % 2 == 0: 
 255                  sign = global_sign 
 256                  nb_fermion += 1 
 257                  if index in flipped: 
 258                      sign *= -1 
 259              else:  
 260                  sign = -1 * global_sign 
 261                  nb_fermion += 1 
 262                  if index-1 in flipped: 
 263                      sign *= -1 
 264              # No need to include the outgoing particles in the definitions 
 265              if index == self.outgoing -1: 
 266                  continue  
 267               
 268              # write the 
 269              momentum_conserve.append('%s%s%d' % (sign_dict[sign], spin, \ 
 270                                                                       index + 1)) 
 271           
 272          # Remove the 
 273          if momentum_conserve[0][0] == '+': 
 274              momentum_conserve[0] = momentum_conserve[0][1:] 
 275           
 276          return momentum_conserve 
 277       
 278 -    def make_declaration_list(self): 
 279          """ make the list of declaration nedded by the header """ 
 280           
 281          declare_list = [] 
 282           
 283           
 284          for index, spin in enumerate(self.particles): 
 285              # First define the size of the associate Object  
 286              declare_list.append(self.declare_dict[spin] % (index + 1) )  
 287    
 288          return declare_list 
 289    
 290        
 291 -class ALOHAWriterForFortran(WriteALOHA):  
 292      """routines for writing out Fortran""" 
 293   
 294      extension = '.f' 
 295      declare_dict = {'S':'double complex S%d(*)', 
 296                      'F':'double complex F%d(*)', 
 297                      'V':'double complex V%d(*)', 
 298                      'T':'double complex T%s(*)'} 
 299       
 300 -    def define_header(self, name=None): 
 301          """Define the Header of the fortran file. This include 
 302              - function tag 
 303              - definition of variable 
 304          """ 
 305           
 306          if not name: 
 307              name = self.namestring 
 308                
 309          Momenta = self.collected['momenta'] 
 310          OverM = self.collected['om'] 
 311           
 312          CallList = self.calllist['CallList'] 
 313          declare_list = self.calllist['DeclareList'] 
 314          if 'double complex COUP' in declare_list: 
 315              alredy_update = True 
 316          else: 
 317              alredy_update = False 
 318           
 319          if not alredy_update:     
 320              declare_list.append('double complex COUP') 
 321   
 322          # define the type of function and argument         
 323          if not self.offshell: 
 324              str_out = 'subroutine %(name)s(%(args)s,vertex)\n' % \ 
 325                 {'name': name, 
 326                  'args': ','.join(CallList+ ['COUP']) } 
 327              if not alredy_update:  
 328                  declare_list.append('double complex vertex')  
 329          else: 
 330              if not alredy_update: 
 331                  declare_list.append('double complex denom') 
 332                  declare_list.append('double precision M%(id)d, W%(id)d' %  
 333                                                            {'id': self.outgoing}) 
 334              call_arg = '%(args)s, COUP, M%(id)d, W%(id)d, %(spin)s%(id)d' % \ 
 335                      {'args': ', '.join(CallList),  
 336                       'spin': self.particles[self.offshell -1], 
 337                       'id': self.outgoing} 
 338              str_out = ' subroutine %s(%s)\n' % (name, call_arg)  
 339   
 340          # Forcing implicit None 
 341          str_out += 'implicit none \n' 
 342           
 343          # Declare all the variable 
 344          for elem in declare_list: 
 345              str_out += elem + '\n' 
 346          if len(OverM) > 0:  
 347              str_out += 'double complex ' + ','.join(OverM) + '\n' 
 348          if len(Momenta) > 0: 
 349              str_out += 'double precision ' + '(0:3),'.join(Momenta) + '(0:3)\n' 
 350   
 351          # Add entry for symmetry 
 352          #str_out += '\n' 
 353          #for elem in self.symmetries: 
 354          #    CallList2 = self.reorder_call_list(CallList, self.offshell, elem) 
 355          #    call_arg = '%(args)s, C, M%(id)d, W%(id)d, %(spin)s%(id)d' % \ 
 356          #            {'args': ', '.join(CallList2),  
 357          #             'spin': self.particles[self.offshell -1], 
 358          #             'id': self.offshell} 
 359          #     
 360          #     
 361          #    str_out += ' entry %(name)s(%(args)s)\n' % \ 
 362          #                {'name': get_routine_name(self.abstractname, elem), 
 363          #                 'args': call_arg} 
 364   
 365          return str_out 
 366         
 367 -    def define_momenta(self): 
 368          """Define the Header of the fortran file. This include 
 369              - momentum conservation 
 370              -definition of the impulsion""" 
 371          # Definition of the Momenta 
 372           
 373          momenta = self.collected['momenta'] 
 374          overm = self.collected['om'] 
 375          momentum_conservation = self.calllist['Momentum'] 
 376           
 377          str_out = '' 
 378          # Conservation of Energy Impulsion 
 379          if self.offshell:  
 380              offshelltype = self.particles[self.offshell -1] 
 381              offshell_size = self.type_to_size[offshelltype]             
 382              #Implement the conservation of Energy Impulsion 
 383              for i in range(-1,1): 
 384                  str_out += '%s%d(%d)= ' % (offshelltype, self.outgoing, \ 
 385                                                                offshell_size + i) 
 386                   
 387                  pat=re.compile(r'^[-+]?(?P<spin>\w)') 
 388                  for elem in momentum_conservation: 
 389                      spin = pat.search(elem).group('spin')  
 390                      str_out += '%s(%d)' % (elem, self.type_to_size[spin] + i)   
 391                  str_out += '\n'   
 392                       
 393          # Momentum 
 394          for mom in momenta: 
 395              #Mom is in format PX with X the number of the particle 
 396              index = int(mom[1:]) 
 397              type = self.particles[index - 1] 
 398              energy_pos = self.type_to_size[type] -1 
 399              sign = 1 
 400              if self.offshell == index and type in ['V','S']: 
 401                  sign = -1 
 402              if 'C%s' % ((index +1) // 2)  in self.tag:  
 403                  if index == self.outgoing: 
 404                      pass 
 405                  elif index % 2 and index -1 != self.outgoing: 
 406                      pass 
 407                  elif index % 2 == 1 and index + 1  != self.outgoing: 
 408                      pass 
 409                  else: 
 410                      sign *= -1 
 411               
 412              if sign == -1 : 
 413                  sign = '-' 
 414              else: 
 415                  sign = '' 
 416                               
 417              str_out += '%s(0) = %s dble(%s%d(%d))\n' % (mom, sign, type, index, energy_pos) 
 418              str_out += '%s(1) = %s dble(%s%d(%d))\n' % (mom, sign, type, index, energy_pos + 1) 
 419              str_out += '%s(2) = %s dimag(%s%d(%d))\n' % (mom, sign, type, index, energy_pos + 1) 
 420              str_out += '%s(3) = %s dimag(%s%d(%d))\n' % (mom, sign, type, index, energy_pos)             
 421               
 422                      
 423          # Definition for the One Over Mass**2 terms 
 424          for elem in overm: 
 425              #Mom is in format OMX with X the number of the particle 
 426              index = int(elem[2:]) 
 427              str_out += 'OM%d = 0d0\n' % (index) 
 428              str_out += 'if (M%d .ne. 0d0) OM%d' % (index, index) + '=1d0/M%d**2\n' % (index)  
 429           
 430          # Returning result 
 431          return str_out 
 432           
 433           
 434 -    def change_var_format(self, name):  
 435          """Formatting the variable name to Fortran format""" 
 436           
 437          if '_' in name: 
 438              name = name.replace('_', '(', 1) + ')' 
 439          #name = re.sub('\_(?P<num>\d+)$', '(\g<num>)', name) 
 440          return name 
 441     
 442      zero_pattern = re.compile(r'''0+$''') 
 443 -    def change_number_format(self, number): 
 444          """Formating the number""" 
 445   
 446          def isinteger(x): 
 447              try: 
 448                  return int(x) == x 
 449              except TypeError: 
 450                  return False 
 451   
 452          if isinteger(number): 
 453              out = str(int(number)) 
 454          elif isinstance(number, complex): 
 455              if number.imag: 
 456                  out = '(%s, %s)' % (self.change_number_format(number.real) , \ 
 457                                      self.change_number_format(number.imag)) 
 458              else: 
 459                  out = self.change_number_format(number.real) 
 460          else: 
 461              out = '%.9f' % number 
 462              out = self.zero_pattern.sub('', out) 
 463          return out 
 464           
 465       
 466 -    def define_expression(self): 
 467          OutString = '' 
 468          if not self.offshell: 
 469              for ind in self.obj.listindices(): 
 470                  string = 'Vertex = COUP*' + self.write_obj(self.obj.get_rep(ind)) 
 471                  string = string.replace('+-', '-') 
 472                  string = re.sub('\((?P<num>[+-]*[0-9])(?P<num2>[+-][0-9])[Jj]\)\.', '(\g<num>d0,\g<num2>d0)', string) 
 473                  string = re.sub('(?P<num>[0-9])[Jj]\.', '\g<num>.*(0d0,1d0)', string) 
 474                  OutString = OutString + string + '\n' 
 475          else: 
 476              OffShellParticle = '%s%d' % (self.particles[self.offshell-1],\ 
 477                                                                    self.outgoing) 
 478              numerator = self.obj.numerator 
 479              denominator = self.obj.denominator 
 480              for ind in denominator.listindices(): 
 481                  denom = self.write_obj(denominator.get_rep(ind)) 
 482              string = 'denom =' + '1d0/(' + denom + ')' 
 483              string = string.replace('+-', '-') 
 484              string = re.sub('\((?P<num>[+-]*[0-9])\+(?P<num2>[+-][0-9])[Jj]\)\.', '(\g<num>d0,\g<num2>d0)', string) 
 485              string = re.sub('(?P<num>[0-9])[Jj]\.', '\g<num>*(0d0,1d0)', string) 
 486              OutString = OutString + string + '\n' 
 487              for ind in numerator.listindices(): 
 488                  string = '%s(%d)= COUP*denom*' % (OffShellParticle, self.pass_to_HELAS(ind, start=1)) 
 489                  string += self.write_obj(numerator.get_rep(ind)) 
 490                  string = string.replace('+-', '-') 
 491                  string = re.sub('\((?P<num>[+-][0-9])\+(?P<num2>[+-][0-9])[Jj]\)\.', '(\g<num>d0,\g<num2>d0)', string) 
 492                  string = re.sub('(?P<num>[0-9])[Jj]\.', '\g<num>*(0d0,1d0)', string) 
 493                  OutString = OutString + string + '\n'  
 494          return OutString  
 495       
 496 -    def define_symmetry(self, new_nb): 
 497          number = self.offshell  
 498          calls = self.calllist['CallList'] 
 499                                                                   
 500          Outstring = 'call '+self.namestring+'('+','.join(calls)+',COUP,M%s,W%s,%s%s)\n' \ 
 501                           %(number,number,self.particles[number-1],number) 
 502          return Outstring 
 503       
 504 -    def define_foot(self): 
 505          return 'end\n\n'  
 506   
 507 -    def write(self, mode='self'): 
 508                            
 509          # write head - momenta - body - foot 
 510          text = self.define_header()+'\n' 
 511          text += self.define_momenta()+'\n' 
 512          text += self.define_expression() 
 513          text += self.define_foot() 
 514           
 515          sym_text = [] 
 516          for elem in self.symmetries:  
 517              symmetryhead = self.define_header().replace( \ 
 518                               self.namestring,self.namestring[0:-1]+'%s' %(elem)) 
 519              symmetrybody = self.define_symmetry(elem) 
 520               
 521              sym_text.append(symmetryhead + symmetrybody + self.define_foot()) 
 522               
 523               
 524          if self.out_path: 
 525              writer = writers.FortranWriter(self.out_path) 
 526              writer.downcase = False  
 527              commentstring = 'This File is Automatically generated by ALOHA \n' 
 528              commentstring += 'The process calculated in this file is: \n' 
 529              commentstring += self.comment + '\n' 
 530              writer.write_comments(commentstring) 
 531              writer.writelines(text) 
 532              for text in sym_text: 
 533                  writer.write_comments('\n%s\n' % ('#'*65)) 
 534                  writer.writelines(text) 
 535          else: 
 536              for stext in sym_text: 
 537                  text += '\n\n' + stext 
 538   
 539          return text + '\n' 
 540               
 541 -    def write_combined(self, lor_names, mode='self', offshell=None): 
 542          """Write routine for combine ALOHA call (more than one coupling)""" 
 543           
 544          # Set some usefull command 
 545          if offshell is None: 
 546              sym = 1  
 547              offshell = self.offshell 
 548          else: 
 549              sym = None  # deactivate symetry 
 550               
 551          name = combine_name(self.abstractname, lor_names, offshell, self.tag) 
 552   
 553                    
 554          # write header  
 555          header = self.define_header(name=name) 
 556          new_couplings = ['COUP%s' % (i+1) for i in range(len(lor_names)+1)] 
 557          text = header.replace('COUP', ','.join(new_couplings)) 
 558          # define the TMP for storing output         
 559          if not offshell: 
 560              text += ' double complex TMP\n' 
 561          else: 
 562              spin = self.particles[offshell -1]  
 563              text += ' double complex TMP(%s)\n integer i' % self.type_to_size[spin]          
 564           
 565          # Define which part of the routine should be called 
 566          addon = ''.join(self.tag) + '_%s' % self.offshell 
 567           
 568  #        if 'C' in self.namestring: 
 569  #            short_name, addon = name.split('C',1) 
 570  #            if addon.split('_')[0].isdigit(): 
 571  #                addon = 'C' +self.namestring.split('C',1)[1] 
 572  #            elif all([n.isdigit() for n in addon.split('_')[0].split('C')]): 
 573  #                addon = 'C' +self.namestring.split('C',1)[1] 
 574  #            else: 
 575  #                addon = '_%s' % self.offshell 
 576  #        else: 
 577  #            addon = '_%s' % self.offshell 
 578   
 579          # how to call the routine 
 580          if not offshell: 
 581              main = 'vertex' 
 582              call_arg = '%(args)s, %(COUP)s, %(LAST)s' % \ 
 583                      {'args': ', '.join(self.calllist['CallList']),  
 584                       'COUP':'COUP%d', 
 585                       'spin': self.particles[self.offshell -1], 
 586                       'LAST': '%s'} 
 587          else: 
 588              main = '%(spin)s%(id)d' % \ 
 589                            {'spin': self.particles[offshell -1], 
 590                             'id': self.outgoing} 
 591              call_arg = '%(args)s, %(COUP)s, M%(id)d, W%(id)d, %(LAST)s' % \ 
 592                      {'args': ', '.join(self.calllist['CallList']),  
 593                       'COUP':'COUP%d', 
 594                       'id': self.outgoing, 
 595                       'LAST': '%s'} 
 596   
 597          # make the first call 
 598          line = " CALL %s%s("+call_arg+")\n" 
 599          text += '\n\n' + line % (self.namestring, '', 1, main) 
 600           
 601          # make the other call 
 602          for i,lor in enumerate(lor_names): 
 603              text += line % (lor, addon, i+2, 'TMP') 
 604              if not offshell: 
 605                  text += ' vertex = vertex + tmp\n' 
 606              else: 
 607                  size = self.type_to_size[spin] -2 
 608                  text += """ do i=1,%(id)d 
 609                  %(main)s(i) = %(main)s(i) + tmp(i) 
 610                  enddo\n""" %  {'id': size, 'main':main} 
 611                   
 612   
 613          text += self.define_foot() 
 614           
 615          #ADD SYMETRY 
 616          if sym: 
 617              for elem in self.symmetries: 
 618                  text += self.write_combined(lor_names, mode, elem) 
 619               
 620               
 621          if self.out_path: 
 622              # Prepare a specific file 
 623              path = os.path.join(os.path.dirname(self.out_path), name+'.f') 
 624              writer = writers.FortranWriter(path) 
 625              writer.downcase = False  
 626              commentstring = 'This File is Automatically generated by ALOHA \n' 
 627              writer.write_comments(commentstring) 
 628              writer.writelines(text) 
 629           
 630          return text 
 631       
 632 -def get_routine_name(name=None, outgoing=None, tag=None, abstract=None): 
 633      """ build the name of the aloha function """ 
 634       
 635      assert (name and outgoing) or abstract 
 636   
 637      if tag is None: 
 638          tag = abstract.tag 
 639   
 640      if name is None: 
 641          name = abstract.name + ''.join(tag) 
 642       
 643      if outgoing is None: 
 644          outgoing = abstract.outgoing 
 645   
 646       
 647      return '%s_%s' % (name, outgoing) 
 648   
 649 -def combine_name(name, other_names, outgoing, tag=None): 
 650      """ build the name for combined aloha function """ 
 651       
 652   
 653      # Two possible scheme FFV1C1_2_X or FFV1__FFV2C1_X 
 654      # If they are all in FFVX scheme then use the first 
 655      p=re.compile('^(?P<type>[FSVT]+)(?P<id>\d+)') 
 656      routine = '' 
 657      if p.search(name): 
 658          base, id = p.search(name).groups() 
 659          if tag is not None: 
 660              routine = name + ''.join(tag) 
 661          else: 
 662              routine = name 
 663          for s in other_names: 
 664              try: 
 665                  base2,id2 = p.search(s).groups() 
 666              except: 
 667                  routine = '' 
 668                  break # one matching not good -> other scheme 
 669              if base != base2: 
 670                  routine = '' 
 671                  break  # one matching not good -> other scheme 
 672              else: 
 673                  routine += '_%s' % id2 
 674      if routine: 
 675          return routine +'_%s' % outgoing 
 676   
 677      if tag is not None: 
 678          addon = ''.join(tag) 
 679      else: 
 680          addon = '' 
 681          if 'C' in name: 
 682              short_name, addon = name.split('C',1) 
 683              try: 
 684                  addon = 'C' + str(int(addon)) 
 685              except: 
 686                  addon = '' 
 687              else: 
 688                  name = short_name 
 689   
 690      return '_'.join((name,) + tuple(other_names)) + addon + '_%s' % outgoing 
 691    
 692   
 693 -class ALOHAWriterForCPP(WriteALOHA):  
 694      """Routines for writing out helicity amplitudes as C++ .h and .cc files.""" 
 695       
 696      declare_dict = {'S':'double complex S%d[3]', 
 697                      'F':'double complex F%d[6]', 
 698                      'V':'double complex V%d[6]', 
 699                      'T':'double complex T%s[18]'} 
 700       
 701 -    def define_header(self, name=None): 
 702          """Define the headers for the C++ .h and .cc files. This include 
 703              - function tag 
 704              - definition of variable 
 705              - momentum conservation 
 706          """ 
 707           
 708          if name is None: 
 709              name = self.namestring 
 710               
 711          #Width = self.collected['width'] 
 712          #Mass = self.collected['mass'] 
 713           
 714          CallList = self.calllist['CallList'][:] 
 715           
 716          local_declare = [] 
 717          OffShell = self.offshell 
 718          OffShellParticle = OffShell -1  
 719          # Transform function call variables to C++ format 
 720          for i, call in enumerate(CallList): 
 721              CallList[i] = "complex<double> %s[]" % call 
 722          #if Mass: 
 723          #    Mass[0] = "double %s" % Mass[0] 
 724          #if Width: 
 725          #    Width[0] = "double %s" % Width[0] 
 726           
 727          # define the type of function and argument 
 728          if not OffShell: 
 729              str_out = 'void %(name)s(%(args)s, complex<double>& vertex)' % \ 
 730                 {'name': name, 
 731                  'args': ','.join(CallList + ['complex<double> COUP'])} 
 732          else:  
 733              str_out = 'void %(name)s(%(args)s, double M%(number)d, double W%(number)d, complex<double>%(out)s%(number)d[])' % \ 
 734                {'name': name, 
 735                 'args': ','.join(CallList+ ['complex<double> COUP']), 
 736                 'out': self.particles[self.outgoing - 1], 
 737                 'number': self.outgoing  
 738                 } 
 739   
 740          h_string = str_out + ";\n\n" 
 741          cc_string = str_out + "{\n" 
 742   
 743          return {'h_header': h_string, 'cc_header': cc_string} 
 744               
 745 -    def define_momenta(self): 
 746          """Write the expressions for the momentum of the outgoing 
 747          particle.""" 
 748   
 749          momenta = self.collected['momenta'] 
 750          overm = self.collected['om'] 
 751          momentum_conservation = self.calllist['Momentum'] 
 752           
 753          str_out = '' 
 754          # Declare auxiliary variables 
 755          if self.offshell: 
 756              str_out += 'complex<double> denom;\n' 
 757          if len(overm) > 0:  
 758              str_out += 'complex<double> %s;\n' % ','.join(overm) 
 759          if len(momenta) > 0: 
 760              str_out += 'double %s[4];\n' % '[4],'.join(momenta) 
 761   
 762          # Energy 
 763          if self.offshell:  
 764              offshelltype = self.particles[self.offshell -1] 
 765              offshell_size = self.type_to_size[offshelltype]             
 766              #Implement the conservation of Energy Impulsion 
 767              for i in range(-2,0): 
 768                  str_out += '%s%d[%d]= ' % (offshelltype, self.outgoing, 
 769                                             offshell_size + i) 
 770                   
 771                  pat=re.compile(r'^[-+]?(?P<spin>\w)') 
 772                  for elem in momentum_conservation: 
 773                      spin = pat.search(elem).group('spin')  
 774                      str_out += '%s[%d]' % (elem, self.type_to_size[spin] + i)   
 775                  str_out += ';\n' 
 776           
 777          # Momentum 
 778          for mom in momenta: 
 779              #Mom is in format PX with X the number of the particle 
 780              index = int(mom[1:]) 
 781               
 782              type = self.particles[index - 1] 
 783              energy_pos = self.type_to_size[type] - 2 
 784              sign = 1 
 785              if self.offshell == index and type in ['V', 'S']: 
 786                  sign = -1 
 787              if 'C%s' % ((index +1) // 2)  in self.tag:  
 788                  if index == self.outgoing: 
 789                      pass 
 790                  elif index % 2 and index -1 != self.outgoing: 
 791                      pass 
 792                  elif index %2 == 1 and index + 1  != self.outgoing: 
 793                      pass 
 794                  else: 
 795                      sign *= -1 
 796               
 797              if sign == -1 : 
 798                  sign = '-' 
 799              else: 
 800                  sign = '' 
 801                      
 802              str_out += '%s[0] = %s%s%d[%d].real();\n' % (mom, sign, type, index, energy_pos) 
 803              str_out += '%s[1] = %s%s%d[%d].real();\n' % (mom, sign, type, index, energy_pos + 1) 
 804              str_out += '%s[2] = %s%s%d[%d].imag();\n' % (mom, sign, type, index, energy_pos + 1) 
 805              str_out += '%s[3] = %s%s%d[%d].imag();\n' % (mom, sign, type, index, energy_pos)             
 806               
 807          # Definition for the One Over Mass**2 terms 
 808          for elem in overm: 
 809              #Mom is in format OMX with X the number of the particle 
 810              index = int(elem[2:]) 
 811              str_out += 'OM%d = 0;\n' % (index) 
 812              str_out += 'if (M%d != 0) OM%d' % (index, index) + '= 1./pow(M%d,2);\n' % (index)  
 813           
 814          # Returning result 
 815          return str_out 
 816           
 817           
 818 -    def change_var_format(self, name):  
 819          """Format the variable name to C++ format""" 
 820           
 821          if '_' in name: 
 822              name = name.replace('_','[',1) +']' 
 823          outstring = '' 
 824          counter = 0 
 825          for elem in re.finditer('[FVTSfvts][0-9]\[[0-9]\]',name): 
 826              outstring += name[counter:elem.start()+2]+'['+str(int(name[elem.start()+3:elem.start()+4])-1)+']' 
 827              counter = elem.end() 
 828          outstring += name[counter:] 
 829          #name = re.sub('\_(?P<num>\d+)$', '(\g<num>)', name) 
 830          return outstring 
 831       
 832 -    def change_number_format(self, number): 
 833          """Format numbers into C++ format""" 
 834          if isinstance(number, complex): 
 835              if number.real == int(number.real) and \ 
 836                     number.imag == int(number.imag): 
 837                  out = 'complex<double>(%d., %d.)' % \ 
 838                        (int(number.real), int(number.imag)) 
 839              else: 
 840                  out = 'complex<double>(%.9f, %.9f)' % \ 
 841                        (number.real, number.imag)                 
 842          else: 
 843              if number == int(number): 
 844                  out = '%d.' % int(number) 
 845              else: 
 846                  out = '%.9f' % number 
 847          return out 
 848       
 849 -    def define_expression(self): 
 850          """Write the helicity amplitude in C++ format""" 
 851          OutString = ''  
 852          if not self.offshell: 
 853              for ind in self.obj.listindices(): 
 854                  string = 'vertex = COUP*' + self.write_obj(self.obj.get_rep(ind)) 
 855                  string = string.replace('+-', '-') 
 856                  OutString = OutString + string + ';\n' 
 857          else: 
 858              OffShellParticle = self.particles[self.offshell-1]+'%s'%(self.outgoing) 
 859              numerator = self.obj.numerator 
 860              denominator = self.obj.denominator 
 861              for ind in denominator.listindices(): 
 862                  denom = self.write_obj(denominator.get_rep(ind)) 
 863              string = 'denom =' + '1./(' + denom + ')' 
 864              string = string.replace('+-', '-') 
 865              OutString = OutString + string + ';\n' 
 866              for ind in numerator.listindices(): 
 867                  string = '%s[%d]= COUP*denom*' % (OffShellParticle, self.pass_to_HELAS(ind)) 
 868                  string += self.write_obj(numerator.get_rep(ind)) 
 869                  string = string.replace('+-', '-') 
 870                  OutString = OutString + string + ';\n'  
 871          OutString = re.sub('(?P<variable>[A-Za-z]+[0-9]\[*[0-9]*\]*)\*\*(?P<num>[0-9])','pow(\g<variable>,\g<num>)',OutString) 
 872          return OutString  
 873   
 874      remove_double = re.compile('complex<double> (?P<name>[\w]+)\[\]') 
 875 -    def define_symmetry(self, new_nb): 
 876          """Write the call for symmetric routines""" 
 877          calls = self.calllist['CallList'] 
 878           
 879          for i, call in enumerate(calls): 
 880              if self.remove_double.match(call): 
 881                  calls[i] = self.remove_double.match(call).group('name') 
 882                   
 883          # For the call, need to remove the type specification 
 884          #calls = [self.remove_double.match(call).group('name') for call in \ 
 885          #         calls] 
 886          number = self.offshell  
 887          Outstring = self.namestring+'('+','.join(calls)+',COUP,M%s,W%s,%s%s);\n' \ 
 888                           %(number,number,self.particles[self.offshell-1],number) 
 889          return Outstring 
 890       
 891 -    def define_foot(self): 
 892          """Return the end of the function definition""" 
 893   
 894          return '}\n\n'  
 895   
 896 -    def write_h(self, header, compiler_cmd=True): 
 897          """Return the full contents of the .h file""" 
 898   
 899          h_string = '' 
 900          if compiler_cmd: 
 901              h_string = '#ifndef '+ self.namestring + '_guard\n' 
 902              h_string += '#define ' + self.namestring + '_guard\n' 
 903              h_string += '#include <complex>\n' 
 904              h_string += 'using namespace std;\n\n' 
 905   
 906          h_header = header['h_header'] 
 907   
 908          h_string += h_header 
 909   
 910          for elem in self.symmetries:  
 911              symmetryhead = h_header.replace( \ 
 912                               self.namestring,self.namestring[0:-1]+'%s' %(elem)) 
 913              h_string += symmetryhead 
 914   
 915          if compiler_cmd: 
 916              h_string += '#endif\n\n' 
 917   
 918          return h_string 
 919   
 920 -    def write_combined_h(self, lor_names, offshell=None, compiler_cmd=True): 
 921          """Return the content of the .h file linked to multiple lorentz call.""" 
 922           
 923          name = combine_name(self.abstractname, lor_names, offshell, self.tag) 
 924          text= '' 
 925          if compiler_cmd: 
 926              text = '#ifndef '+ name + '_guard\n' 
 927              text += '#define ' + name + '_guard\n' 
 928              text += '#include <complex>\n' 
 929              text += 'using namespace std;\n\n' 
 930           
 931          # write header  
 932          header = self.define_header(name=name) 
 933          h_header = header['h_header'] 
 934          new_couplings = ['COUP%s' % (i+1) for i in range(len(lor_names)+1)] 
 935          h_string = h_header.replace('COUP', ', complex <double>'.join(new_couplings)) 
 936          text += h_string  
 937           
 938          for elem in self.symmetries:  
 939              text += h_string.replace(name, name[0:-1]+str(elem)) 
 940           
 941          if compiler_cmd: 
 942              text += '#endif' 
 943           
 944          return text 
 945   
 946 -    def write_cc(self, header, compiler_cmd=True): 
 947          """Return the full contents of the .cc file""" 
 948   
 949          cc_string = '' 
 950          if compiler_cmd: 
 951              cc_string = '#include \"%s.h\"\n\n' % self.namestring 
 952          cc_header = header['cc_header'] 
 953          cc_string += cc_header 
 954          cc_string += self.define_momenta() 
 955          cc_string += self.define_expression() 
 956          cc_string += self.define_foot() 
 957   
 958          for elem in self.symmetries:  
 959              symmetryhead = cc_header.replace( \ 
 960                               self.namestring,self.namestring[0:-1]+'%s' %(elem)) 
 961              symmetrybody = self.define_symmetry(elem) 
 962              cc_string += symmetryhead 
 963              cc_string += symmetrybody 
 964              cc_string += self.define_foot() 
 965   
 966          return cc_string 
 967   
 968 -    def write_combined_cc(self, lor_names, offshell=None, compiler_cmd=True, sym=True): 
 969          "Return the content of the .cc file linked to multiple lorentz call." 
 970           
 971          # Set some usefull command 
 972          if offshell is None: 
 973              offshell = self.offshell 
 974               
 975          name = combine_name(self.abstractname, lor_names, offshell, self.tag) 
 976   
 977          text = '' 
 978          if compiler_cmd: 
 979              text += '#include "%s.h"\n\n' % name 
 980              
 981          # write header  
 982          header = self.define_header(name=name)['cc_header'] 
 983          new_couplings = ['COUP%s' % (i+1) for i in range(len(lor_names)+1)] 
 984          text += header.replace('COUP', ', complex<double>'.join(new_couplings)) 
 985          # define the TMP for storing output         
 986          if not offshell: 
 987              text += 'complex<double> tmp;\n' 
 988          else: 
 989              spin = self.particles[offshell -1]  
 990              text += 'complex<double> tmp[%s];\n int i = 0;' % self.type_to_size[spin]          
 991   
 992          # Define which part of the routine should be called 
 993          addon = '' 
 994          if 'C' in self.namestring: 
 995              short_name, addon = name.split('C',1) 
 996              if addon.split('_')[0].isdigit(): 
 997                  addon = 'C' +self.namestring.split('C',1)[1] 
 998              elif all([n.isdigit() for n in addon.split('_')[0].split('C')]): 
 999                  addon = 'C' +self.namestring.split('C',1)[1] 
1000              else: 
1001                  addon = '_%s' % self.offshell 
1002          else: 
1003              addon = '_%s' % self.offshell 
1004   
1005          # how to call the routine 
1006          if not offshell: 
1007              main = 'vertex' 
1008              call_arg = '%(args)s, %(COUP)s, %(LAST)s' % \ 
1009                      {'args': ', '.join(self.calllist['CallList']),  
1010                       'COUP':'COUP%d', 
1011                       'spin': self.particles[self.offshell -1], 
1012                       'LAST': '%s'} 
1013          else: 
1014              main = '%(spin)s%(id)d' % \ 
1015                            {'spin': self.particles[self.offshell -1], 
1016                             'id': self.outgoing} 
1017              call_arg = '%(args)s, %(COUP)s, M%(id)d, W%(id)d, %(LAST)s' % \ 
1018                      {'args': ', '.join(self.calllist['CallList']),  
1019                       'COUP':'COUP%d', 
1020                       'id': self.outgoing, 
1021                       'LAST': '%s'} 
1022   
1023          # make the first call 
1024          line = "%s%s("+call_arg+");\n" 
1025          text += '\n\n' + line % (self.namestring, '', 1, main) 
1026           
1027          # make the other call 
1028          for i,lor in enumerate(lor_names): 
1029              text += line % (lor, addon, i+2, 'tmp') 
1030              if not offshell: 
1031                  text += ' vertex = vertex + tmp;\n' 
1032              else: 
1033                  size = self.type_to_size[spin] -2 
1034                  text += """ while (i < %(id)d) 
1035                  { 
1036                  %(main)s[i] = %(main)s[i] + tmp[i]; 
1037                  i++; 
1038                  }\n""" %  {'id': size, 'main':main} 
1039                   
1040   
1041          text += self.define_foot() 
1042           
1043          if sym: 
1044              for elem in self.symmetries: 
1045                  text += self.write_combined_cc(lor_names, elem,  
1046                                                        compiler_cmd=compiler_cmd, 
1047                                                        sym=False) 
1048               
1049               
1050          if self.out_path: 
1051              # Prepare a specific file 
1052              path = os.path.join(os.path.dirname(self.out_path), name+'.f') 
1053              writer = writers.FortranWriter(path) 
1054              writer.downcase = False  
1055              commentstring = 'This File is Automatically generated by ALOHA \n' 
1056              writer.write_comments(commentstring) 
1057              writer.writelines(text) 
1058           
1059          return text 
1060   
1061   
1062       
1063 -    def write(self, mode='self', **opt): 
1064          """Write the .h and .cc files""" 
1065   
1066   
1067          # write head - momenta - body - foot 
1068          header = self.define_header() 
1069          h_text = self.write_h(header, **opt) 
1070          cc_text = self.write_cc(header, **opt) 
1071           
1072          # write in two file 
1073          if self.out_path: 
1074              writer_h = writers.CPPWriter(self.out_path + ".h") 
1075              writer_cc = writers.CPPWriter(self.out_path + ".cc") 
1076              commentstring = 'This File is Automatically generated by ALOHA \n' 
1077              commentstring += 'The process calculated in this file is: \n' 
1078              commentstring += self.comment + '\n' 
1079              writer_h.write_comments(commentstring) 
1080              writer_cc.write_comments(commentstring) 
1081              writer_h.writelines(h_text) 
1082              writer_cc.writelines(cc_text) 
1083               
1084          return h_text, cc_text 
1085    
1086    
1087    
1088 -    def write_combined(self, lor_names, mode='self', offshell=None, **opt): 
1089          """Write the .h and .cc files associated to the combined file""" 
1090           
1091          # Set some usefull command 
1092          if offshell is None: 
1093              offshell = self.offshell 
1094           
1095          name = combine_name(self.abstractname, lor_names, offshell, self.tag) 
1096           
1097          h_text = self.write_combined_h(lor_names, offshell, **opt) 
1098          cc_text = self.write_combined_cc(lor_names, offshell, sym=True, **opt) 
1099           
1100          if self.out_path: 
1101              # Prepare a specific file 
1102              path = os.path.join(os.path.dirname(self.out_path), name) 
1103              commentstring = 'This File is Automatically generated by ALOHA \n' 
1104               
1105              writer_h = writers.CPPWriter(path + ".h") 
1106              writer_h.write_comments(commentstring) 
1107              writer_h.writelines(h_text) 
1108               
1109              writer_cc = writers.CPPWriter(path + ".cc") 
1110              writer_cc.write_comments(commentstring) 
1111              writer_cc.writelines(cc_text) 
1112          else: 
1113              return h_text, cc_text 
1114           
1115          return h_text, cc_text 
1116           
1117 -class ALOHAWriterForPython(WriteALOHA): 
1118      """ A class for returning a file/a string for python evaluation """ 
1119       
1120      extension = '.py' 
1121       
1122 -    def __init__(self, abstract_routine, dirpath=None): 
1123          """ standard init but if dirpath is None the write routine will 
1124          return a string (which can be evaluated by an exec""" 
1125           
1126          WriteALOHA.__init__(self, abstract_routine, dirpath) 
1127          self.outname = '%s%s' % (self.particles[self.offshell -1], \ 
1128                                                                 self.outgoing) 
1129       
1130      @staticmethod 
1131 -    def change_number_format(obj): 
1132          if obj.real == 0 and obj.imag: 
1133              if int(obj.imag) == obj.imag:  
1134                  return '%ij' % obj.imag 
1135              else: 
1136                  return '%sj' % str(obj.imag) 
1137          else:  
1138              return str(obj) 
1139       
1140       
1141 -    def change_var_format(self, name):  
1142          """Formatting the variable name to Python format 
1143          start to count at zero""" 
1144           
1145          def shift_indices(match): 
1146              """shift the indices for non impulsion object""" 
1147              if match.group('var').startswith('P'): 
1148                  shift = 0 
1149              else:  
1150                  shift = -1 
1151               
1152              return '%s[%s]' % (match.group('var'), \ 
1153                                                  int(match.group('num')) + shift) 
1154               
1155           
1156          name = re.sub('(?P<var>\w*)_(?P<num>\d+)$', shift_indices , name) 
1157          return name 
1158       
1159 -    def define_expression(self): 
1160          """Define the functions in a 100% way """ 
1161           
1162          OutString = '' 
1163          if not self.offshell: 
1164              for ind in self.obj.listindices(): 
1165                  string = 'vertex = COUP*' + self.write_obj(self.obj.get_rep(ind)) 
1166                  string = string.replace('+-', '-') 
1167                  OutString = OutString + string + '\n' 
1168          else: 
1169              OffShellParticle = '%s%d' % (self.particles[self.offshell-1],\ 
1170                                                                    self.offshell) 
1171              numerator = self.obj.numerator 
1172              denominator = self.obj.denominator 
1173              for ind in denominator.listindices(): 
1174                  denom = self.write_obj(denominator.get_rep(ind)) 
1175              string = 'denom =' + '1.0/(' + denom + ')' 
1176              string = string.replace('+-', '-') 
1177              OutString += string + '\n' 
1178              for ind in numerator.listindices(): 
1179                  string = '%s[%d]= COUP*denom*' % (self.outname, self.pass_to_HELAS(ind)) 
1180                  string += self.write_obj(numerator.get_rep(ind)) 
1181                  string = string.replace('+-', '-') 
1182                  OutString += string + '\n'  
1183          return OutString  
1184       
1185 -    def define_foot(self): 
1186          if not self.offshell: 
1187              return 'return vertex\n\n' 
1188          else: 
1189              return 'return %s\n\n' % (self.outname) 
1190               
1191       
1192 -    def define_header(self, name=None): 
1193          """Define the Header of the fortran file. This include 
1194              - function tag 
1195              - definition of variable 
1196          """ 
1197          if name is None: 
1198              name = self.namestring 
1199              
1200          Momenta = self.collected['momenta'] 
1201          OverM = self.collected['om'] 
1202           
1203          CallList = self.calllist['CallList'] 
1204   
1205          str_out = '' 
1206          # define the type of function and argument         
1207          if not self.offshell: 
1208              str_out += 'def %(name)s(%(args)s):\n' % \ 
1209                  {'name': name, 
1210                   'args': ','.join(CallList+ ['COUP']) } 
1211          else: 
1212              str_out += 'def %(name)s(%(args)s, COUP, M%(id)d, W%(id)d):\n' % \ 
1213                  {'name': name, 
1214                   'args': ', '.join(CallList),  
1215                       'id': self.outgoing}             
1216          return str_out      
1217   
1218 -    def make_declaration_list(self): 
1219          """ make the list of declaration nedded by the header """ 
1220          return [] 
1221   
1222 -    def define_momenta(self): 
1223          """Define the Header of the fortran file. This include 
1224              - momentum conservation 
1225              - definition of the impulsion""" 
1226               
1227          # Definition of the Momenta 
1228          momenta = self.collected['momenta'] 
1229          overm = self.collected['om'] 
1230          momentum_conservation = self.calllist['Momentum'] 
1231           
1232          str_out = '' 
1233           
1234          # Definition of the output 
1235          if self.offshell: 
1236              offshelltype = self.particles[self.offshell -1] 
1237              offshell_size = self.type_to_size[offshelltype]             
1238              str_out += '%s = wavefunctions.WaveFunction(size=%s)\n' % \ 
1239                                                  (self.outname, offshell_size) 
1240          # Conservation of Energy Impulsion 
1241          if self.offshell:      
1242              #Implement the conservation of Energy Impulsion 
1243              for i in range(-2,0): 
1244                  str_out += '%s[%d] = ' % (self.outname, offshell_size + i) 
1245                   
1246                  pat=re.compile(r'^[-+]?(?P<spin>\w)') 
1247                  for elem in momentum_conservation: 
1248                      spin = pat.search(elem).group('spin')  
1249                      str_out += '%s[%d]' % (elem, self.type_to_size[spin] + i)   
1250                  str_out += '\n'   
1251                       
1252          # Momentum 
1253          for mom in momenta: 
1254              #Mom is in format PX with X the number of the particle 
1255              index = int(mom[1:]) 
1256              type = self.particles[index - 1] 
1257              energy_pos = self.type_to_size[type] -2 
1258              sign = 1 
1259              if self.offshell == index and type in ['V','S']: 
1260                  sign = -1 
1261   
1262              if 'C%s' % ((index +1) // 2)  in self.tag:  
1263                  if index == self.outgoing: 
1264                      pass 
1265                  elif index % 2 and index -1 != self.outgoing: 
1266                      pass 
1267                  elif index %2 == 1 and index + 1  != self.outgoing: 
1268                      pass 
1269                  else: 
1270                      sign *= -1 
1271   
1272              if sign == -1 : 
1273                  sign = '- ' 
1274              else: 
1275                  sign = '' 
1276   
1277              str_out += '%s = [%scomplex(%s%d[%d]).real, \\\n' % (mom, sign, type, index, energy_pos) 
1278              str_out += '         %scomplex(%s%d[%d]).real, \\\n' % ( sign, type, index, energy_pos + 1) 
1279              str_out += '         %scomplex(%s%d[%d]).imag, \\\n' % ( sign, type, index, energy_pos + 1) 
1280              str_out += '         %scomplex(%s%d[%d]).imag]\n' % ( sign, type, index, energy_pos)  
1281                      
1282          # Definition for the One Over Mass**2 terms 
1283          for elem in overm: 
1284              #Mom is in format OMX with X the number of the particle 
1285              index = int(elem[2:]) 
1286              str_out += 'OM%d = 0.0\n' % (index) 
1287              str_out += 'if (M%d): OM%d' % (index, index) + '=1.0/M%d**2\n' % (index)  
1288           
1289          # Returning result 
1290          return str_out     
1291   
1292 -    def define_symmetry(self, new_nb): 
1293          number = self.offshell  
1294          calls = self.calllist['CallList'] 
1295          Outstring = 'return '+self.namestring+'('+','.join(calls)+',COUP,M%s,W%s)\n' \ 
1296                           %(number,number) 
1297          return Outstring         
1298   
1299 -    def write(self,mode='self'): 
1300                            
1301          # write head - momenta - body - foot 
1302          if mode == 'mg5': 
1303              text = 'import aloha.template_files.wavefunctions as wavefunctions\n' 
1304          else: 
1305              text = 'import wavefunctions\n' 
1306          text += self.define_header() 
1307          content = self.define_momenta() 
1308          content += self.define_expression() 
1309          content += self.define_foot() 
1310           
1311          # correct identation 
1312          text += '    ' +content.replace('\n','\n    ') 
1313           
1314          for elem in self.symmetries: 
1315              text +='\n' + self.define_header().replace( \ 
1316                               self.namestring,self.namestring[0:-1]+'%s' %(elem)) 
1317              text += '    ' +self.define_symmetry(elem) 
1318               
1319               
1320          if self.out_path: 
1321              ff = open(self.out_path,'w').write(text) 
1322   
1323          return text + '\n' 
1324   
1325 -    def write_combined(self, lor_names, mode='self', offshell=None): 
1326          """Write routine for combine ALOHA call (more than one coupling)""" 
1327           
1328          # Set some usefull command 
1329          if offshell is None: 
1330              sym = 1 
1331              offshell = self.offshell   
1332          else: 
1333              sym = None 
1334          name = combine_name(self.abstractname, lor_names, offshell, self.tag) 
1335   
1336          # write head - momenta - body - foot 
1337          text = '' 
1338          #if mode == 'mg5': 
1339          #    text = 'import aloha.template_files.wavefunctions as wavefunctions\n' 
1340          #else: 
1341          #    text = 'import wavefunctions\n' 
1342                       
1343                    
1344          # write header  
1345          header = self.define_header(name=name) 
1346          new_couplings = ['COUP%s' % (i+1) for i in range(len(lor_names)+1)] 
1347          header = header.replace('COUP', ','.join(new_couplings)) 
1348           
1349          text += header 
1350     
1351          # Define which part of the routine should be called 
1352          addon = ''.join(self.tag) + '_%s' % self.offshell 
1353   
1354          # how to call the routine 
1355          if not offshell: 
1356              main = 'vertex' 
1357              call_arg = '%(args)s, %(COUP)s' % \ 
1358                      {'args': ', '.join(self.calllist['CallList']),  
1359                       'COUP':'COUP%d', 
1360                       'spin': self.particles[self.offshell -1]} 
1361          else: 
1362              main = '%(spin)s%(id)d' % \ 
1363                            {'spin': self.particles[self.offshell -1], 
1364                             'id': self.outgoing} 
1365              call_arg = '%(args)s, %(COUP)s, M%(id)d, W%(id)d' % \ 
1366                      {'args': ', '.join(self.calllist['CallList']),  
1367                       'COUP':'COUP%d', 
1368                       'id': self.outgoing} 
1369   
1370          # make the first call 
1371          line = "    %s = %s%s("+call_arg+")\n" 
1372          text += '\n\n' + line % (main, self.namestring, '', 1) 
1373           
1374          # make the other call 
1375          for i,name in enumerate(lor_names): 
1376              text += line % ('tmp',name, addon, i+2) 
1377              if not offshell: 
1378                 text += '    vertex += tmp\n' 
1379              else: 
1380                  size = self.type_to_size[self.particles[offshell -1]] -2 
1381                  text += "    for i in range(%(id)d):\n" % {'id': size} 
1382                  text += "        %(main)s[i] += tmp[i]\n" %{'main': main} 
1383           
1384          text += '    '+self.define_foot() 
1385   
1386          #ADD SYMETRY 
1387          if sym: 
1388              for elem in self.symmetries: 
1389                  text += self.write_combined(lor_names, mode, elem) 
1390               
1391          return text 
1392