Package aloha :: Module aloha_writers
[hide private]
[frames] | no frames]

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