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      import madgraph.various.q_polynomial as q_polynomial 
   4      import madgraph.various.misc as misc 
   5  except Exception: 
   6      import aloha.file_writers as writers 
   7      import aloha.q_polynomial as q_polynomial 
   8      import aloha.misc as misc 
   9   
  10  import aloha 
  11  import aloha.aloha_lib as aloha_lib 
  12  import cmath 
  13  import os 
  14  import re  
  15  from numbers import Number 
  16  from collections import defaultdict 
  17  from fractions import Fraction 
  18  # fast way to deal with string 
  19  from cStringIO import StringIO 
  20  # Look at http://www.skymind.com/~ocrow/python_string/  
  21  # For knowing how to deal with long strings efficiently. 
  22  import itertools 
  23   
  24  KERNEL = aloha_lib.KERNEL 
  25  pjoin = os.path.join 
26 27 -class WriteALOHA:
28 """ Generic writing functions """ 29 30 power_symbol = '**' 31 change_number_format = str 32 extension = '' 33 type_to_variable = {2:'F',3:'V',5:'T',1:'S',4:'R', -1:'S'} 34 type_to_size = {'S':3, 'T':18, 'V':6, 'F':6,'R':18} 35 36
37 - def __init__(self, abstract_routine, dirpath):
38 if aloha.loop_mode: 39 self.momentum_size = 4 40 else: 41 self.momentum_size = 2 42 43 self.has_model_parameter = False 44 45 name = get_routine_name(abstract = abstract_routine) 46 47 if dirpath: 48 self.dir_out = dirpath 49 self.out_path = os.path.join(dirpath, name + self.extension) 50 else: 51 self.out_path = None 52 self.dir_out = None 53 54 self.routine = abstract_routine 55 self.tag = self.routine.tag 56 self.name = name 57 58 self.particles = [self.type_to_variable[spin] for spin in \ 59 abstract_routine.spins] 60 61 self.offshell = abstract_routine.outgoing # position of the outgoing in particle list 62 self.outgoing = self.offshell # expected position for the argument list 63 if 'C%s' %((self.outgoing + 1) // 2) in self.tag: 64 #flip the outgoing tag if in conjugate 65 self.outgoing = self.outgoing + self.outgoing % 2 - (self.outgoing +1) % 2 66 self.outname = '%s%s' % (self.particles[self.outgoing -1], \ 67 self.outgoing) 68 #initialize global helper routine 69 self.declaration = Declaration_list()
70 71
72 - def pass_to_HELAS(self, indices, start=0):
73 """find the Fortran HELAS position for the list of index""" 74 75 76 if len(indices) == 1: 77 return indices[0] + start + self.momentum_size 78 79 try: 80 # When the expr is not a SplitCoefficient 81 ind_name = self.routine.expr.lorentz_ind 82 except: 83 # When the expr is a loop one, i.e. with SplitCoefficient 84 if len(set([tuple(expr.lorentz_ind) for expr in self.routine.expr.values()]))!=1: 85 raise Exception('All SplitCoefficients do not share the same indices names.') 86 for expr in self.routine.expr.values(): 87 ind_name = expr.lorentz_ind 88 break 89 90 if ind_name == ['I3', 'I2']: 91 return 4 * indices[1] + indices[0] + start + self.momentum_size 92 elif len(indices) == 2: 93 return 4 * indices[0] + indices[1] + start + self.momentum_size 94 else: 95 raise Exception, 'WRONG CONTRACTION OF LORENTZ OBJECT for routine %s: %s' \ 96 % (self.name, ind_name)
97
98 - def get_header_txt(self,mode=''):
99 """ Prototype for language specific header""" 100 raise Exception, 'THis function should be overwritten' 101 return '' 102
103 - def get_declaration_txt(self):
104 """ Prototype for how to write the declaration of variable""" 105 return ''
106
107 - def define_content(self):
108 """Prototype for language specific body""" 109 pass 110
111 - def get_momenta_txt(self):
112 """ Prototype for the definition of the momenta""" 113 raise Exception, 'THis function should be overwritten'
114
115 - def get_momentum_conservation_sign(self):
116 """find the sign associated to the momentum conservation""" 117 118 # help data 119 signs = [] 120 nb_fermion =0 121 122 #compute global sign 123 124 global_sign = -1 125 126 flipped = [2*(int(c[1:])-1) for c in self.tag if c.startswith('C')] 127 for index, spin in enumerate(self.particles): 128 assert(spin in ['S','F','V','T', 'R']) 129 130 #compute the sign 131 if 1:#spin != 'F': 132 sign = -1 * global_sign 133 elif nb_fermion % 2 == 0: 134 sign = global_sign 135 nb_fermion += 1 136 if index in flipped: 137 sign *= -1 138 else: 139 sign = -1 * global_sign 140 nb_fermion += 1 141 if index-1 in flipped: 142 sign *= -1 143 144 # No need to include the outgoing particles in the definitions 145 if index == self.outgoing -1: 146 signs.append('0*') 147 continue 148 149 if sign == 1: 150 signs.append('+') 151 else: 152 signs.append('-') 153 return signs
154 155
156 - def get_P_sign(self, index):
157 158 type = self.particles[index - 1] 159 energy_pos = self.type_to_size[type] -1 160 sign = 1 161 if self.outgoing == index: 162 sign = -1 163 #if 'C%s' % ((index +1) // 2) in self.tag: 164 # if index == self.outgoing: 165 # pass 166 # elif index % 2 and index -1 != self.outgoing: 167 # pass 168 # elif index % 2 == 1 and index + 1 != self.outgoing: 169 # pass 170 # else: 171 # sign *= -1 172 173 if sign == -1 : 174 return '-' 175 else: 176 return ''
177 178 179 180 181
182 - def get_foot_txt(self):
183 """Prototype for language specific footer""" 184 return ''
185
186 - def define_argument_list(self, couplings=None):
187 """define a list with the string of object required as incoming argument""" 188 189 call_arg = [] #incoming argument of the routine 190 191 conjugate = [2*(int(c[1:])-1) for c in self.tag if c[0] == 'C'] 192 193 194 for index,spin in enumerate(self.particles): 195 if self.offshell == index + 1: 196 continue 197 198 if index in conjugate: 199 index2, spin2 = index+1, self.particles[index+1] 200 call_arg.append(('list_complex','%s%d' % (spin2, index2 +1))) 201 #call_arg.append('%s%d' % (spin, index +1)) 202 elif index-1 in conjugate: 203 index2, spin2 = index-1, self.particles[index-1] 204 call_arg.append(('list_complex','%s%d' % (spin2, index2 +1))) 205 else: 206 call_arg.append(('list_complex','%s%d' % (spin, index +1))) 207 208 # couplings 209 if couplings is None: 210 detected_couplings = [name for type, name in self.declaration if name.startswith('COUP')] 211 coup_sort = lambda x,y: int(x[4:])-int(y[4:]) 212 detected_couplings.sort(coup_sort) 213 if detected_couplings: 214 couplings = detected_couplings 215 else: 216 couplings = ['COUP'] 217 218 for coup in couplings: 219 call_arg.append(('complex', coup)) 220 self.declaration.add(('complex',coup)) 221 222 if self.offshell: 223 if aloha.complex_mass: 224 call_arg.append(('complex','M%s' % self.outgoing)) 225 self.declaration.add(('complex','M%s' % self.outgoing)) 226 else: 227 call_arg.append(('double','M%s' % self.outgoing)) 228 self.declaration.add(('double','M%s' % self.outgoing)) 229 call_arg.append(('double','W%s' % self.outgoing)) 230 self.declaration.add(('double','W%s' % self.outgoing)) 231 232 assert len(call_arg) == len(set([a[1] for a in call_arg])) 233 assert len(self.declaration) == len(set([a[1] for a in self.declaration])), self.declaration 234 self.call_arg = call_arg 235 return call_arg
236
237 - def write(self, mode=None):
238 239 self.mode = mode 240 241 core_text = self.define_expression() 242 self.define_argument_list() 243 out = StringIO() 244 out.write(self.get_header_txt(mode=self.mode)) 245 out.write(self.get_declaration_txt()) 246 out.write(self.get_momenta_txt()) 247 out.write(core_text) 248 out.write(self.get_foot_txt()) 249 250 for elem in self.routine.symmetries: 251 out.write('\n') 252 out.write(self.define_symmetry(elem)) 253 254 text = out.getvalue() 255 256 if self.out_path: 257 writer = self.writer(self.out_path) 258 commentstring = 'This File is Automatically generated by ALOHA \n' 259 commentstring += 'The process calculated in this file is: \n' 260 commentstring += self.routine.infostr + '\n' 261 writer.write_comments(commentstring) 262 writer.writelines(text) 263 264 return text + '\n'
265 266
267 - def write_indices_part(self, indices, obj):
268 """Routine for making a string out of indices objects""" 269 270 text = 'output(%s)' % indices 271 return text 272
273 - def write_obj(self, obj, prefactor=True):
274 """Calls the appropriate writing routine""" 275 276 try: 277 vartype = obj.vartype 278 except Exception: 279 return self.change_number_format(obj) 280 281 # The order is from the most current one to the les probable one 282 if vartype == 1 : # AddVariable 283 return self.write_obj_Add(obj, prefactor) 284 elif vartype == 2 : # MultVariable 285 return self.write_MultVariable(obj, prefactor) 286 elif vartype == 6 : # MultContainer 287 return self.write_MultContainer(obj, prefactor) 288 elif vartype == 0 : # MultContainer 289 return self.write_variable(obj) 290 else: 291 raise Exception('Warning unknown object: %s' % obj.vartype)
292
293 - def write_MultVariable(self, obj, prefactor=True):
294 """Turn a multvariable into a string""" 295 296 mult_list = [self.write_variable_id(id) for id in obj] 297 data = {'factors': '*'.join(mult_list)} 298 if prefactor and obj.prefactor != 1: 299 if obj.prefactor != -1: 300 text = '%(prefactor)s * %(factors)s' 301 data['prefactor'] = self.change_number_format(obj.prefactor) 302 else: 303 text = '-%(factors)s' 304 else: 305 text = '%(factors)s' 306 return text % data
307
308 - def write_MultContainer(self, obj, prefactor=True):
309 """Turn a multvariable into a string""" 310 311 mult_list = [self.write_obj(id) for id in obj] 312 data = {'factors': '*'.join(mult_list)} 313 if prefactor and obj.prefactor != 1: 314 if obj.prefactor != -1: 315 text = '%(prefactor)s * %(factors)s' 316 data['prefactor'] = self.change_number_format(obj.prefactor) 317 else: 318 text = '-%(factors)s' 319 else: 320 text = '%(factors)s' 321 return text % data
322 323
324 - def write_obj_Add(self, obj, prefactor=True):
325 """Turns addvariable into a string""" 326 327 data = defaultdict(list) 328 number = [] 329 [data[p.prefactor].append(p) if hasattr(p, 'prefactor') else number.append(p) 330 for p in obj] 331 332 file_str = StringIO() 333 334 if prefactor and obj.prefactor != 1: 335 formatted = self.change_number_format(obj.prefactor) 336 if formatted.startswith(('+','-')): 337 file_str.write('(%s)' % formatted) 338 else: 339 file_str.write(formatted) 340 file_str.write('*(') 341 else: 342 file_str.write('(') 343 first=True 344 for value, obj_list in data.items(): 345 add= '+' 346 if value not in [-1,1]: 347 nb_str = self.change_number_format(value) 348 if nb_str[0] in ['+','-']: 349 file_str.write(nb_str) 350 else: 351 file_str.write('+') 352 file_str.write(nb_str) 353 file_str.write('*(') 354 elif value == -1: 355 add = '-' 356 file_str.write('-') 357 elif not first: 358 file_str.write('+') 359 else: 360 file_str.write('') 361 first = False 362 file_str.write(add.join([self.write_obj(obj, prefactor=False) 363 for obj in obj_list])) 364 if value not in [1,-1]: 365 file_str.write(')') 366 if number: 367 total = sum(number) 368 file_str.write('+ %s' % self.change_number_format(total)) 369 370 file_str.write(')') 371 return file_str.getvalue()
372
373 - def write_variable(self, obj):
374 return self.change_var_format(obj)
375
376 - def write_variable_id(self, id):
377 378 obj = aloha_lib.KERNEL.objs[id] 379 return self.write_variable(obj)
380
381 - def change_var_format(self, obj):
382 """format the way to write the variable and add it to the declaration list 383 """ 384 385 str_var = str(obj) 386 self.declaration.add((obj.type, str_var)) 387 return str_var
388 389 390
391 - def make_call_list(self, outgoing=None):
392 """find the way to write the call of the functions""" 393 394 if outgoing is None: 395 outgoing = self.offshell 396 397 call_arg = [] #incoming argument of the routine 398 399 conjugate = [2*(int(c[1:])-1) for c in self.tag if c[0] == 'C'] 400 401 for index,spin in enumerate(self.particles): 402 if self.offshell == index + 1: 403 continue 404 405 if index in conjugate: 406 index2, spin2 = index+1, self.particles[index+1] 407 call_arg.append('%s%d' % (spin2, index2 +1)) 408 #call_arg.append('%s%d' % (spin, index +1)) 409 elif index-1 in conjugate: 410 index2, spin2 = index-1, self.particles[index-1] 411 call_arg.append('%s%d' % (spin2, index2 +1)) 412 else: 413 call_arg.append('%s%d' % (spin, index +1)) 414 415 416 return call_arg
417 418
419 - def make_declaration_list(self):
420 """ make the list of declaration nedded by the header """ 421 422 declare_list = [] 423 424 425 for index, spin in enumerate(self.particles): 426 # First define the size of the associate Object 427 declare_list.append(self.declare_dict[spin] % (index + 1) ) 428 429 return declare_list
430
431 432 433 434 435 -class ALOHAWriterForFortran(WriteALOHA):
436 """routines for writing out Fortran""" 437 438 extension = '.f' 439 writer = writers.FortranWriter 440 441 type2def = {} 442 type2def['int'] = 'integer*4' 443 if aloha.mp_precision: 444 type2def['double'] = 'real*16' 445 type2def['complex'] = 'complex*32' 446 format = 'q0' 447 else: 448 type2def['double'] = 'real*8' 449 type2def['complex'] = 'complex*16' 450 451 format = 'd0' 452
453 - def get_fct_format(self, fct):
454 """Put the function in the correct format""" 455 if not hasattr(self, 'fct_format'): 456 one = self.change_number_format(1) 457 self.fct_format = {'csc' : '{0}/cos(dble(%s))'.format(one), 458 'sec': '{0}/sin(dble(%s))'.format(one), 459 'acsc': 'asin({0}/(dble(%s)))'.format(one), 460 'asec': 'acos({0}/(%s))'.format(one), 461 're': ' dble(%s)', 462 'im': 'imag(%s)', 463 'cmath.sqrt':'sqrt(dble(%s))', 464 'sqrt': 'sqrt(dble(%s))', 465 'complexconjugate': 'conjg(dcmplx(%s))', 466 '/' : '{0}/(%s)'.format(one), 467 'pow': '(%s)**(%s)', 468 'log': 'log(dble(%s))', 469 'asin': 'asin(dble(%s))', 470 'acos': 'acos(dble(%s))', 471 'abs': 'std::abs(%s)', 472 'fabs': 'std::abs(%s)', 473 'math.abs': 'std::abs(%s)', 474 'cmath.abs': 'std::abs(%s)', 475 '':'(%s)' 476 } 477 478 if fct in self.fct_format: 479 return self.fct_format[fct] 480 else: 481 self.declaration.add(('fct', fct)) 482 return '{0}(%s)'.format(fct)
483 484 485
486 - def get_header_txt(self, name=None, couplings=None, **opt):
487 """Define the Header of the fortran file. 488 """ 489 if name is None: 490 name = self.name 491 492 out = StringIO() 493 # define the type of function and argument 494 495 arguments = [arg for format, arg in self.define_argument_list(couplings)] 496 if not self.offshell: 497 output = 'vertex' 498 self.declaration.add(('complex','vertex')) 499 else: 500 output = '%(spin)s%(id)d' % { 501 'spin': self.particles[self.outgoing -1], 502 'id': self.outgoing} 503 self.declaration.add(('list_complex', output)) 504 505 out.write('subroutine %(name)s(%(args)s,%(output)s)\n' % \ 506 {'output':output, 'name': name, 'args': ', '.join(arguments)}) 507 508 return out.getvalue()
509
510 - def get_declaration_txt(self):
511 """ Prototype for how to write the declaration of variable 512 Include the symmetry line (entry FFV_2) 513 """ 514 515 out = StringIO() 516 out.write('implicit none\n') 517 # Check if we are in formfactor mode 518 if self.has_model_parameter: 519 out.write(' include "../MODEL/input.inc"\n') 520 out.write(' include "../MODEL/coupl.inc"\n') 521 argument_var = [name for type,name in self.call_arg] 522 # define the complex number CI = 0+1j 523 if 'MP' in self.tag: 524 out.write(' complex*32 CI\n') 525 if KERNEL.has_pi: 526 out.write(' REAL ( KIND = 16 ) PI\n') 527 else: 528 out.write(' complex*16 CI\n') 529 if KERNEL.has_pi: 530 out.write(' double precision PI\n') 531 out.write(' parameter (CI=(%s,%s))\n' % 532 (self.change_number_format(0),self.change_number_format(1))) 533 if KERNEL.has_pi: 534 out.write(' parameter (PI=%s)\n' % self.change_number_format(cmath.pi)) 535 for type, name in self.declaration: 536 if type.startswith('list'): 537 type = type[5:] 538 #determine the size of the list 539 if name in argument_var: 540 size ='*' 541 elif name.startswith('P'): 542 size='0:3' 543 elif name[0] in ['F','V']: 544 if aloha.loop_mode: 545 size = 8 546 else: 547 size = 6 548 elif name[0] == 'S': 549 if aloha.loop_mode: 550 size = 5 551 else: 552 size = 3 553 elif name[0] in ['R','T']: 554 if aloha.loop_mode: 555 size = 20 556 else: 557 size = 18 558 else: 559 size = '*' 560 561 out.write(' %s %s(%s)\n' % (self.type2def[type], name, size)) 562 elif type == 'fct': 563 if name.upper() in ['EXP','LOG','SIN','COS','ASIN','ACOS']: 564 continue 565 out.write(' %s %s\n' % (self.type2def['complex'], name)) 566 out.write(' external %s\n' % (name)) 567 else: 568 out.write(' %s %s\n' % (self.type2def[type], name)) 569 570 # Add the lines corresponding to the symmetry 571 572 #number = self.offshell 573 #arguments = [name for format, name in self.define_argument_list()] 574 #new_name = self.name.rsplit('_')[0] + '_%s' % new_nb 575 #return '%s\n call %s(%s)' % \ 576 # (self.get_header_txt(new_name, couplings), self.name, ','.join(arguments)) 577 couplings = [name for type, name in self.declaration if name.startswith('COUP') ] 578 couplings.sort() 579 for elem in self.routine.symmetries: 580 new_name = self.name.rsplit('_',1)[0] + '_%s' % elem 581 out.write('%s\n' % self.get_header_txt(new_name, couplings).replace('subroutine','entry')) 582 583 584 return out.getvalue()
585
586 - def get_momenta_txt(self):
587 """Define the Header of the fortran file. This include 588 - momentum conservation 589 - definition of the impulsion""" 590 591 out = StringIO() 592 593 # Define all the required momenta 594 p = [] # a list for keeping track how to write the momentum 595 596 signs = self.get_momentum_conservation_sign() 597 for i,type in enumerate(self.particles): 598 if self.declaration.is_used('OM%s' % (i+1)): 599 out.write(" OM{0} = {1}\n if (M{0}.ne.{1}) OM{0}={2}/M{0}**2\n".format( 600 i+1, self.change_number_format(0), self.change_number_format(1))) 601 602 603 if i+1 == self.outgoing: 604 out_type = type 605 out_size = self.type_to_size[type] 606 continue 607 elif self.offshell: 608 p.append('{0}{1}{2}(%(i)s)'.format(signs[i],type,i+1,type)) 609 610 if self.declaration.is_used('P%s' % (i+1)): 611 self.get_one_momenta_def(i+1, out) 612 613 # define the resulting momenta 614 if self.offshell: 615 616 617 energy_pos = out_size -2 618 type = self.particles[self.outgoing-1] 619 620 for i in range(self.momentum_size): 621 dict_energy = {'i':1+i} 622 out.write(' %s%s(%s) = %s\n' % (type,self.outgoing, 1+i, 623 ''.join(p) % dict_energy)) 624 if self.declaration.is_used('P%s' % self.outgoing): 625 self.get_one_momenta_def(self.outgoing, out) 626 627 if "P1T" in self.tag or "P1L" in self.tag: 628 for i in range(1,4): 629 P = "P%s" % (self.outgoing) 630 value = ["1d-30", "0d0", "1d-15"] 631 out.write(" IF (DABS(%(P)s(0))*1e-10.gt.DABS(%(P)s(%(i)s))) %(P)s(%(i)s)=%(val)s\n" 632 % {"P": P, "i":i, 'val':value[i-1]}) 633 i = self.outgoing -1 634 if self.declaration.is_used('Tnorm%s' % (i+1)): 635 out.write(" TNORM{0} = DSQRT(P{0}(1)*P{0}(1)+P{0}(2)*P{0}(2)+P{0}(3)*P{0}(3))\n".format( 636 i+1)) 637 if self.declaration.is_used('TnormZ%s' % (i+1)): 638 out.write(" TNORMZ{0} = TNORM{0} - P{0}(3)\n".format( 639 i+1)) 640 641 if self.declaration.is_used('FWP%s' % (i+1)): 642 out.write(" FWP{0} = DSQRT(-P{0}(0) + TNORM{0})\n"\ 643 .format(i+1)) 644 if self.declaration.is_used('FWM%s' % (i+1)): 645 out.write(" FWM{0} = DSQRT(-P{0}(0) - TNORM{0})\n"\ 646 .format(i+1)) 647 #out.write(" FWM{0} = M{0}/FWP{0}\n".format(i+1)) 648 649 # Returning result 650 return out.getvalue()
651
652 - def get_one_momenta_def(self, i, strfile):
653 654 type = self.particles[i-1] 655 656 if aloha.loop_mode: 657 template ='P%(i)d(%(j)d) = %(sign)s%(type)s%(i)d(%(nb)d)\n' 658 else: 659 template ='P%(i)d(%(j)d) = %(sign)s%(operator)s(%(type)s%(i)d(%(nb2)d))\n' 660 661 nb2 = 1 662 for j in range(4): 663 if not aloha.loop_mode: 664 nb = j + 1 665 if j == 0: 666 assert not aloha.mp_precision 667 operator = 'dble' # not suppose to pass here in mp 668 elif j == 1: 669 nb2 += 1 670 elif j == 2: 671 assert not aloha.mp_precision 672 operator = 'dimag' # not suppose to pass here in mp 673 elif j ==3: 674 nb2 -= 1 675 else: 676 operator ='' 677 nb = 1+ j 678 nb2 = 1 + j 679 strfile.write(template % {'j':j,'type': type, 'i': i, 680 'nb': nb, 'nb2': nb2, 'operator':operator, 681 'sign': self.get_P_sign(i)})
682
683 - def shift_indices(self, match):
684 """shift the indices for non impulsion object""" 685 if match.group('var').startswith('P'): 686 shift = 0 687 else: 688 shift = self.momentum_size 689 return '%s(%s)' % (match.group('var'), int(match.group('num')) + shift)
690
691 - def change_var_format(self, name):
692 """Formatting the variable name to Fortran format""" 693 694 if isinstance(name, aloha_lib.ExtVariable): 695 # external parameter nothing to do but handling model prefix 696 self.has_model_parameter = True 697 if name.lower() in ['pi', 'as', 'mu_r', 'aewm1','g']: 698 return name 699 if name.startswith(aloha.aloha_prefix): 700 return name 701 return '%s%s' % (aloha.aloha_prefix, name) 702 703 if '_' in name: 704 vtype = name.type 705 decla = name.split('_',1)[0] 706 self.declaration.add(('list_%s' % vtype, decla)) 707 else: 708 self.declaration.add((name.type, name)) 709 name = re.sub('(?P<var>\w*)_(?P<num>\d+)$', self.shift_indices , name) 710 return name 711
712 - def change_number_format(self, number):
713 """Formating the number""" 714 715 def isinteger(x): 716 try: 717 return int(x) == x 718 except TypeError: 719 return False
720 721 if isinteger(number): 722 out = '%s%s' % (str(int(number)),self.format) 723 elif isinstance(number, complex): 724 if number.imag: 725 if number.real: 726 out = '(%s + %s*CI)' % (self.change_number_format(number.real), \ 727 self.change_number_format(number.imag)) 728 else: 729 if number.imag == 1: 730 out = 'CI' 731 elif number.imag == -1: 732 out = '-CI' 733 else: 734 out = '%s * CI' % self.change_number_format(number.imag) 735 else: 736 out = '%s' % (self.change_number_format(number.real)) 737 else: 738 tmp = Fraction(str(number)) 739 tmp = tmp.limit_denominator(100) 740 if not abs(tmp - number) / abs(tmp + number) < 1e-8: 741 if 'e' in str(number): 742 out = str(number).replace('e','d') 743 else: 744 out = '%s%s' % (number, self.format) 745 else: 746 out = '%s%s/%s%s' % (tmp.numerator, self.format, tmp.denominator, self.format) 747 return out 748
749 - def define_expression(self):
750 """Define the functions in a 100% way """ 751 752 out = StringIO() 753 754 if self.routine.contracted: 755 for name,obj in self.routine.contracted.items(): 756 out.write(' %s = %s\n' % (name, self.write_obj(obj))) 757 self.declaration.add(('complex', name)) 758 759 760 def sort_fct(a, b): 761 if len(a) < len(b): 762 return -1 763 elif len(a) > len(b): 764 return 1 765 elif a < b: 766 return -1 767 else: 768 return +1
769 770 keys = self.routine.fct.keys() 771 keys.sort(sort_fct) 772 for name in keys: 773 fct, objs = self.routine.fct[name] 774 format = ' %s = %s\n' % (name, self.get_fct_format(fct)) 775 try: 776 text = format % ','.join([self.write_obj(obj) for obj in objs]) 777 except TypeError: 778 text = format % tuple([self.write_obj(obj) for obj in objs]) 779 finally: 780 out.write(text) 781 782 783 numerator = self.routine.expr 784 if not 'Coup(1)' in self.routine.infostr: 785 coup_name = 'COUP' 786 else: 787 coup_name = '%s' % self.change_number_format(1) 788 789 790 if not self.offshell: 791 if coup_name == 'COUP': 792 formatted = self.write_obj(numerator.get_rep([0])) 793 if formatted.startswith(('+','-')): 794 out.write(' vertex = COUP*(%s)\n' % formatted) 795 else: 796 out.write(' vertex = COUP*%s\n' % formatted) 797 else: 798 out.write(' vertex = %s\n' % self.write_obj(numerator.get_rep([0]))) 799 else: 800 OffShellParticle = '%s%d' % (self.particles[self.offshell-1],\ 801 self.offshell) 802 is_loop = False 803 if 'L' in self.tag: 804 if self.tag.count('L') == 1 and 'PL' in self.tag: 805 is_loop = False 806 else: 807 is_loop = True 808 809 if not is_loop: 810 coeff = 'denom*' 811 if not aloha.complex_mass: 812 if self.routine.denominator: 813 out.write(' denom = %(COUP)s/(%(denom)s)\n' % {'COUP': coup_name,\ 814 'denom':self.write_obj(self.routine.denominator)}) 815 else: 816 out.write(' denom = %(COUP)s/(P%(i)s(0)**2-P%(i)s(1)**2-P%(i)s(2)**2-P%(i)s(3)**2 - M%(i)s * (M%(i)s -CI* W%(i)s))\n' % \ 817 {'i': self.outgoing, 'COUP': coup_name}) 818 else: 819 if self.routine.denominator: 820 raise Exception, 'modify denominator are not compatible with complex mass scheme' 821 822 out.write(' denom = %(COUP)s/(P%(i)s(0)**2-P%(i)s(1)**2-P%(i)s(2)**2-P%(i)s(3)**2 - M%(i)s**2)\n' % \ 823 {'i': self.outgoing, 'COUP': coup_name}) 824 self.declaration.add(('complex','denom')) 825 if aloha.loop_mode: 826 ptype = 'list_complex' 827 else: 828 ptype = 'list_double' 829 self.declaration.add((ptype,'P%s' % self.outgoing)) 830 else: 831 if coup_name == 'COUP': 832 coeff = 'COUP*' 833 else: 834 coeff = '' 835 to_order = {} 836 for ind in numerator.listindices(): 837 formatted = self.write_obj(numerator.get_rep(ind)) 838 if formatted.startswith(('+','-')): 839 if '*' in formatted: 840 formatted = '(%s)*%s' % tuple(formatted.split('*',1)) 841 else: 842 if formatted.startswith('+'): 843 formatted = formatted[1:] 844 else: 845 formatted = '(-1)*%s' % formatted[1:] 846 to_order[self.pass_to_HELAS(ind)] = \ 847 ' %s(%d)= %s%s\n' % (self.outname, self.pass_to_HELAS(ind)+1, 848 coeff, formatted) 849 key = to_order.keys() 850 key.sort() 851 for i in key: 852 out.write(to_order[i]) 853 return out.getvalue() 854
855 - def define_symmetry(self, new_nb, couplings=None):
856 return ''
857 #number = self.offshell 858 #arguments = [name for format, name in self.define_argument_list()] 859 #new_name = self.name.rsplit('_')[0] + '_%s' % new_nb 860 #return '%s\n call %s(%s)' % \ 861 # (self.get_header_txt(new_name, couplings), self.name, ','.join(arguments)) 862
863 - def get_foot_txt(self):
864 return 'end\n\n'
865
866 - def write_combined(self, lor_names, mode='self', offshell=None):
867 """Write routine for combine ALOHA call (more than one coupling)""" 868 869 # Set some usefull command 870 if offshell is None: 871 sym = 1 872 offshell = self.offshell 873 else: 874 sym = None 875 name = combine_name(self.routine.name, lor_names, offshell, self.tag) 876 self.name = name 877 # write head - momenta - body - foot 878 text = StringIO() 879 routine = StringIO() 880 data = {} # for the formating of the line 881 882 # write header 883 new_couplings = ['COUP%s' % (i+1) for i in range(len(lor_names)+1)] 884 text.write(self.get_header_txt(name=name, couplings=new_couplings)) 885 886 # Define which part of the routine should be called 887 data['addon'] = ''.join(self.tag) + '_%s' % self.offshell 888 889 # how to call the routine 890 argument = [name for format, name in self.define_argument_list(new_couplings)] 891 index= argument.index('COUP1') 892 data['before_coup'] = ','.join(argument[:index]) 893 data['after_coup'] = ','.join(argument[index+len(lor_names)+1:]) 894 if data['after_coup']: 895 data['after_coup'] = ',' + data['after_coup'] 896 897 lor_list = (self.routine.name,) + lor_names 898 line = " call %(name)s%(addon)s(%(before_coup)s,%(coup)s%(after_coup)s,%(out)s)\n" 899 main = '%(spin)s%(id)d' % {'spin': self.particles[self.outgoing -1], 900 'id': self.outgoing} 901 for i, name in enumerate(lor_list): 902 data['name'] = name 903 data['coup'] = 'COUP%d' % (i+1) 904 if i == 0: 905 if not offshell: 906 data['out'] = 'vertex' 907 else: 908 data['out'] = main 909 elif i==1: 910 if self.offshell: 911 type = self.particles[self.outgoing-1] 912 self.declaration.add(('list_complex','%stmp' % type)) 913 else: 914 type = '' 915 self.declaration.add(('complex','%stmp' % type)) 916 data['out'] = '%stmp' % type 917 routine.write(line % data) 918 if i: 919 if not offshell: 920 routine.write( ' vertex = vertex + tmp\n') 921 else: 922 size = self.type_to_size[self.particles[self.outgoing -1]] -2 923 routine.write(" do i = %s, %s\n" % (self.momentum_size+1, self.momentum_size+size)) 924 routine.write(" %(main)s(i) = %(main)s(i) + %(tmp)s(i)\n" %\ 925 {'main': main, 'tmp': data['out']}) 926 routine.write(' enddo\n') 927 self.declaration.add(('int','i')) 928 929 self.declaration.discard(('complex','COUP')) 930 for name in aloha_lib.KERNEL.reduced_expr2: 931 self.declaration.discard(('complex', name)) 932 933 #clean pointless declaration 934 #self.declaration.discard 935 936 937 text.write(self.get_declaration_txt()) 938 text.write(routine.getvalue()) 939 text.write(self.get_foot_txt()) 940 941 942 text = text.getvalue() 943 if self.out_path: 944 writer = self.writer(self.out_path,'a') 945 commentstring = 'This File is Automatically generated by ALOHA \n' 946 commentstring += 'The process calculated in this file is: \n' 947 commentstring += self.routine.infostr + '\n' 948 writer.write_comments(commentstring) 949 writer.writelines(text) 950 return text
951
952 953 -class QP(object):
954 """routines for writing out Fortran""" 955 956 type2def = {} 957 type2def['int'] = 'integer*4' 958 type2def['double'] = 'real*16' 959 type2def['complex'] = 'complex*32' 960 format = 'q0' 961
962 -class ALOHAWriterForFortranQP(QP, ALOHAWriterForFortran):
963
964 - def __init__(self, *arg):
965 return ALOHAWriterForFortran.__init__(self, *arg)
966
967 -class ALOHAWriterForFortranLoop(ALOHAWriterForFortran):
968 """routines for writing out Fortran""" 969
970 - def __init__(self, abstract_routine, dirpath):
971 972 ALOHAWriterForFortran.__init__(self, abstract_routine, dirpath) 973 # position of the outgoing in particle list 974 self.l_id = [int(c[1:]) for c in abstract_routine.tag if c[0] == 'L'][0] 975 self.l_helas_id = self.l_id # expected position for the argument list 976 if 'C%s' %((self.l_id + 1) // 2) in abstract_routine.tag: 977 #flip the outgoing tag if in conjugate 978 self.l_helas_id += self.l_id % 2 - (self.l_id +1) % 2
979 980
981 - def define_expression(self):
982 """Define the functions in a 100% way """ 983 984 out = StringIO() 985 986 if self.routine.contracted: 987 for name,obj in self.routine.contracted.items(): 988 out.write(' %s = %s\n' % (name, self.write_obj(obj))) 989 self.declaration.add(('complex', name)) 990 991 if not 'Coup(1)' in self.routine.infostr: 992 coup = True 993 else: 994 coup = False 995 996 rank = self.routine.expr.get_max_rank() 997 poly_object = q_polynomial.Polynomial(rank) 998 nb_coeff = q_polynomial.get_number_of_coefs_for_rank(rank) 999 size = self.type_to_size[self.particles[self.l_id-1]] - 2 1000 for K in range(size): 1001 for J in range(nb_coeff): 1002 data = poly_object.get_coef_at_position(J) 1003 arg = [data.count(i) for i in range(4)] # momentum 1004 arg += [0] * (K) + [1] + [0] * (size-1-K) 1005 try: 1006 expr = self.routine.expr[tuple(arg)] 1007 except KeyError: 1008 expr = None 1009 for ind in self.routine.expr.values()[0].listindices(): 1010 if expr: 1011 data = expr.get_rep(ind) 1012 else: 1013 data = 0 1014 if data and coup: 1015 out.write(' COEFF(%s,%s,%s)= coup*%s\n' % ( 1016 self.pass_to_HELAS(ind)+1-self.momentum_size, 1017 J, K+1, self.write_obj(data))) 1018 else: 1019 out.write(' COEFF(%s,%s,%s)= %s\n' % ( 1020 self.pass_to_HELAS(ind)+1-self.momentum_size, 1021 J, K+1, self.write_obj(data))) 1022 1023 1024 return out.getvalue()
1025
1026 - def get_declaration_txt(self):
1027 """ Prototype for how to write the declaration of variable""" 1028 1029 out = StringIO() 1030 out.write('implicit none\n') 1031 # define the complex number CI = 0+1j 1032 if 'MP' in self.tag: 1033 out.write(' complex*32 CI\n') 1034 else: 1035 out.write(' complex*16 CI\n') 1036 out.write(' parameter (CI=(%s,%s))\n' % 1037 (self.change_number_format(0),self.change_number_format(1))) 1038 argument_var = [name for type,name in self.call_arg] 1039 for type, name in self.declaration: 1040 if type.startswith('list'): 1041 type = type[5:] 1042 #determine the size of the list 1043 if name.startswith('P'): 1044 size='0:3' 1045 elif name in argument_var: 1046 size ='*' 1047 elif name[0] in ['F','V']: 1048 if aloha.loop_mode: 1049 size = 8 1050 else: 1051 size = 6 1052 elif name[0] == 'S': 1053 if aloha.loop_mode: 1054 size = 5 1055 else: 1056 size = 3 1057 elif name[0] in ['R','T']: 1058 if aloha.loop_mode: 1059 size = 20 1060 else: 1061 size = 18 1062 elif name == 'coeff': 1063 out.write("include 'coef_specs.inc'\n") 1064 size = 'MAXLWFSIZE,0:VERTEXMAXCOEFS-1,MAXLWFSIZE' 1065 1066 out.write(' %s %s(%s)\n' % (self.type2def[type], name, size)) 1067 elif type == 'fct': 1068 if name.upper() in ['EXP','LOG','SIN','COS','ASIN','ACOS']: 1069 continue 1070 out.write(' %s %s\n' % (self.type2def['complex'], name)) 1071 out.write(' external %s\n' % (name)) 1072 else: 1073 out.write(' %s %s\n' % (self.type2def[type], name)) 1074 1075 return out.getvalue()
1076 1077
1078 - def define_argument_list(self, couplings=None):
1079 """define a list with the string of object required as incoming argument""" 1080 1081 conjugate = [2*(int(c[1:])-1) for c in self.tag if c[0] == 'C'] 1082 call_arg = [] 1083 #incoming argument of the routine 1084 call_arg.append( ('list_complex', 'P%s'% self.l_helas_id) ) 1085 1086 self.declaration.add(call_arg[0]) 1087 1088 for index,spin in enumerate(self.particles): 1089 if self.outgoing == index + 1: 1090 continue 1091 if self.l_helas_id == index + 1: 1092 continue 1093 call_arg.append(('complex','%s%d' % (spin, index +1))) 1094 self.declaration.add(('list_complex', call_arg[-1][-1])) 1095 1096 # couplings 1097 if couplings is None: 1098 detected_couplings = [name for type, name in self.declaration if name.startswith('COUP')] 1099 coup_sort = lambda x,y: int(x[4:])-int(y[4:]) 1100 detected_couplings.sort(coup_sort) 1101 if detected_couplings: 1102 couplings = detected_couplings 1103 else: 1104 couplings = ['COUP'] 1105 1106 for coup in couplings: 1107 call_arg.append(('complex', coup)) 1108 self.declaration.add(('complex',coup)) 1109 1110 if self.offshell: 1111 if aloha.complex_mass: 1112 call_arg.append(('complex','M%s' % self.outgoing)) 1113 self.declaration.add(('complex','M%s' % self.outgoing)) 1114 else: 1115 call_arg.append(('double','M%s' % self.outgoing)) 1116 self.declaration.add(('double','M%s' % self.outgoing)) 1117 call_arg.append(('double','W%s' % self.outgoing)) 1118 self.declaration.add(('double','W%s' % self.outgoing)) 1119 1120 self.call_arg = call_arg 1121 1122 return call_arg
1123
1124 - def get_momenta_txt(self):
1125 """Define the Header of the ortran file. This include 1126 - momentum conservation 1127 - definition of the impulsion""" 1128 1129 out = StringIO() 1130 1131 # Define all the required momenta 1132 p = [] # a list for keeping track how to write the momentum 1133 size = [] 1134 1135 signs = self.get_momentum_conservation_sign() 1136 1137 for i,type in enumerate(self.particles): 1138 if self.declaration.is_used('OM%s' % (i+1)): 1139 out.write(" OM{0} = {1}\n if (M{0}.ne.{1}) OM{0}={2}/M{0}**2\n".format( 1140 i+1, self.change_number_format(0), self.change_number_format(1))) 1141 1142 if i+1 == self.outgoing: 1143 out_type = 'P' 1144 continue 1145 elif i+1 == self.l_helas_id: 1146 p.append('%sP%s({%s})' % (signs[i],i+1,len(size))) 1147 size.append(0) 1148 continue 1149 elif self.offshell: 1150 p.append('%s%s%s({%s})' % (signs[i],type,i+1,len(size))) 1151 size.append(1) 1152 1153 if self.declaration.is_used('P%s' % (i+1)): 1154 self.get_one_momenta_def(i+1, out) 1155 1156 # define the resulting momenta 1157 if self.offshell: 1158 if aloha.loop_mode: 1159 size_p = 4 1160 else: 1161 size_p = 2 1162 for i in range(size_p): 1163 out.write(' P%s(%s) = %s\n' % (self.outgoing, i, 1164 ''.join(p).format(*[s+i for s in size]))) 1165 1166 1167 # Returning result 1168 return out.getvalue()
1169 1170
1171 - def get_loop_argument(self, key):
1172 """return the position for the argument in the HELAS convention""" 1173 1174 loop_momentum = key[:4] 1175 basis = key[4:] 1176 1177 loop_pos = sum([loop_momentum[i] * (i+1) for i in range(4)]) 1178 basis_pos = sum([basis[i] * (i+1) for i in range(len(basis))]) 1179 return (str(loop_pos), str(basis_pos))
1180 1181 1182 1183 1184 1185
1186 - def get_header_txt(self, name=None, couplings=None, **opt):
1187 """Define the Header of the fortran file. This include 1188 - function tag 1189 - definition of variable 1190 """ 1191 if name is None: 1192 name = self.name 1193 1194 out = StringIO() 1195 # define the type of function and argument 1196 1197 arguments = [arg for format, arg in self.define_argument_list(couplings)] 1198 self.declaration.add(('list_complex', 'P%s'% self.outgoing)) 1199 self.declaration.add(('list_complex', 'P%s'% self.l_helas_id)) 1200 self.declaration.add(('list_complex', 'coeff')) 1201 out.write('subroutine %(name)s(%(args)s, P%(out)s, COEFF)\n' % \ 1202 {'name': name, 'args': ', '.join(arguments), 1203 'out':self.outgoing}) 1204 1205 return out.getvalue()
1206
1207 -class ALOHAWriterForFortranLoopQP(QP, ALOHAWriterForFortranLoop):
1208 """routines for writing out Fortran""" 1209
1210 - def __init__(self, *arg):
1211 return ALOHAWriterForFortranLoop.__init__(self, *arg)
1212
1213 -def get_routine_name(name=None, outgoing=None, tag=None, abstract=None):
1214 """ build the name of the aloha function """ 1215 1216 assert (name and outgoing is not None) or abstract 1217 1218 if tag is None: 1219 tag = list(abstract.tag) 1220 else: 1221 tag=list(tag) 1222 1223 if name is None: 1224 prefix='' 1225 if 'MP' in tag: 1226 prefix = 'MP_' 1227 tag.remove('MP') 1228 if any(t.startswith('P') for t in tag): 1229 #put the propagator tag at the end 1230 propa = [t for t in tag if t.startswith('P')][0] 1231 tag.remove(propa) 1232 tag.append(propa) 1233 name = prefix + abstract.name + ''.join(tag) 1234 1235 if outgoing is None: 1236 outgoing = abstract.outgoing 1237 1238 return '%s_%s' % (name, outgoing)
1239
1240 -def combine_name(name, other_names, outgoing, tag=None, unknown_propa=False):
1241 """ build the name for combined aloha function """ 1242 1243 def myHash(target_string): 1244 if len(target_string)<50: 1245 return target_string 1246 if '%(propa)s' in target_string: 1247 return 'ALOHA_'+(str(hash(target_string.lower()))).replace('-','m')+'%(propa)s' 1248 else: 1249 return 'ALOHA_'+(str(hash(target_string.lower()))).replace('-','m')
1250 1251 if tag and any(t.startswith('P') for t in tag[:-1]): 1252 # propagator need to be the last entry for the tag 1253 for i,t in enumerate(tag): 1254 if t.startswith('P'): 1255 tag.pop(i) 1256 tag.append(t) 1257 break 1258 1259 # Two possible scheme FFV1C1_2_X or FFV1__FFV2C1_X 1260 # If they are all in FFVX scheme then use the first 1261 p=re.compile('^(?P<type>[RFSVT]{2,})(?P<id>\d+)$') 1262 routine = '' 1263 if p.search(name): 1264 base, id = p.search(name).groups() 1265 routine = name 1266 for s in other_names: 1267 try: 1268 base2,id2 = p.search(s).groups() 1269 except Exception: 1270 routine = '' 1271 break # one matching not good -> other scheme 1272 if base != base2: 1273 routine = '' 1274 break # one matching not good -> other scheme 1275 else: 1276 routine += '_%s' % id2 1277 1278 if routine: 1279 if tag is not None: 1280 routine += ''.join(tag) 1281 if unknown_propa and outgoing: 1282 routine += '%(propa)s' 1283 if outgoing is not None: 1284 return myHash(routine)+'_%s' % outgoing 1285 # return routine +'_%s' % outgoing 1286 else: 1287 return myHash(routine) 1288 # return routine 1289 1290 if tag is not None: 1291 addon = ''.join(tag) 1292 else: 1293 addon = '' 1294 if 'C' in name: 1295 short_name, addon = name.split('C',1) 1296 try: 1297 addon = 'C' + str(int(addon)) 1298 except Exception: 1299 addon = '' 1300 else: 1301 name = short_name 1302 if unknown_propa: 1303 addon += '%(propa)s' 1304 1305 # if outgoing is not None: 1306 # return '_'.join((name,) + tuple(other_names)) + addon + '_%s' % outgoing 1307 # else: 1308 # return '_'.join((name,) + tuple(other_names)) + addon 1309 1310 if outgoing is not None: 1311 return myHash('_'.join((name,) + tuple(other_names))) + addon + '_%s' % outgoing 1312 else: 1313 return myHash('_'.join((name,) + tuple(other_names))) + addon 1314
1315 -class ALOHAWriterForCPP(WriteALOHA):
1316 """Routines for writing out helicity amplitudes as C++ .h and .cc files.""" 1317 1318 extension = '.c' 1319 writer = writers.CPPWriter 1320 1321 type2def = {} 1322 type2def['int'] = 'int ' 1323 type2def['double'] = 'double ' 1324 type2def['complex'] = 'std::complex<double> ' 1325 1326 #variable overwritten by gpu 1327 realoperator = '.real()' 1328 imagoperator = '.imag()' 1329 ci_definition = 'static std::complex<double> cI = std::complex<double>(0.,1.);\n' 1330 1331
1332 - def change_number_format(self, number):
1333 """Formating the number""" 1334 1335 def isinteger(x): 1336 try: 1337 return int(x) == x 1338 except TypeError: 1339 return False
1340 1341 if isinteger(number): 1342 out = '%s.' % (str(int(number))) 1343 elif isinstance(number, complex): 1344 if number.imag: 1345 if number.real: 1346 out = '(%s + %s*cI)' % (self.change_number_format(number.real), \ 1347 self.change_number_format(number.imag)) 1348 else: 1349 if number.imag == 1: 1350 out = 'cI' 1351 elif number.imag == -1: 1352 out = '-cI' 1353 else: 1354 out = '%s * cI' % self.change_number_format(number.imag) 1355 else: 1356 out = '%s' % (self.change_number_format(number.real)) 1357 else: 1358 tmp = Fraction(str(number)) 1359 tmp = tmp.limit_denominator(100) 1360 if not abs(tmp - number) / abs(tmp + number) < 1e-8: 1361 out = '%.9f' % (number) 1362 else: 1363 out = '%s./%s.' % (tmp.numerator, tmp.denominator) 1364 return out 1365 1366
1367 - def shift_indices(self, match):
1368 """shift the indices for non impulsion object""" 1369 if match.group('var').startswith('P'): 1370 shift = 0 1371 else: 1372 shift = self.momentum_size - 1 1373 return '%s[%s]' % (match.group('var'), int(match.group('num')) + shift)
1374 1375
1376 - def change_var_format(self, name):
1377 """Format the variable name to C++ format""" 1378 1379 if '_' in name: 1380 type = name.type 1381 decla = name.split('_',1)[0] 1382 self.declaration.add(('list_%s' % type, decla)) 1383 else: 1384 self.declaration.add((name.type, name.split('_',1)[0])) 1385 name = re.sub('(?P<var>\w*)_(?P<num>\d+)$', self.shift_indices , name) 1386 return name 1387
1388 - def get_fct_format(self, fct):
1389 """Put the function in the correct format""" 1390 if not hasattr(self, 'fct_format'): 1391 one = self.change_number_format(1) 1392 self.fct_format = {'csc' : '{0}/cos(%s)'.format(one), 1393 'sec': '{0}/sin(%s)'.format(one), 1394 'acsc': 'asin({0}/(%s))'.format(one), 1395 'asec': 'acos({0}/(%s))'.format(one), 1396 're': ' real(%s)', 1397 'im': 'imag(%s)', 1398 'cmath.sqrt':'sqrt(%s)', 1399 'sqrt': 'sqrt(%s)', 1400 'complexconjugate': 'conj(dcmplx(%s))', 1401 '/' : '{0}/(%s)'.format(one), 1402 'abs': 'std::abs(%s)' 1403 } 1404 1405 if fct in self.fct_format: 1406 return self.fct_format[fct] 1407 else: 1408 self.declaration.add(('fct', fct)) 1409 return '{0}(%s)'.format(fct)
1410 1411 1412 1413
1414 - def get_header_txt(self, name=None, couplings=None,mode=''):
1415 """Define the Header of the fortran file. This include 1416 - function tag 1417 - definition of variable 1418 """ 1419 if name is None: 1420 name = self.name 1421 1422 if mode=='': 1423 mode = self.mode 1424 1425 1426 1427 out = StringIO() 1428 # define the type of function and argument 1429 if not 'no_include' in mode: 1430 out.write('#include \"%s.h\"\n\n' % self.name) 1431 args = [] 1432 for format, argname in self.define_argument_list(couplings): 1433 if format.startswith('list'): 1434 type = self.type2def[format[5:]] 1435 list_arg = '[]' 1436 else: 1437 type = self.type2def[format] 1438 list_arg = '' 1439 args.append('%s%s%s'% (type, argname, list_arg)) 1440 1441 if not self.offshell: 1442 output = 'std::complex<double> & vertex' 1443 #self.declaration.add(('complex','vertex')) 1444 else: 1445 output = 'std::complex<double> %(spin)s%(id)d[]' % { 1446 'spin': self.particles[self.outgoing -1], 1447 'id': self.outgoing} 1448 self.declaration.add(('list_complex', output)) 1449 1450 out.write('void %(name)s(%(args)s,%(output)s)' % \ 1451 {'output':output, 'name': name, 'args': ', '.join(args)}) 1452 if 'is_h' in mode: 1453 out.write(';\n') 1454 else: 1455 out.write('\n{\n') 1456 1457 return out.getvalue()
1458
1459 - def get_declaration_txt(self, add_i=True):
1460 """ Prototype for how to write the declaration of variable 1461 Include the symmetry line (entry FFV_2) 1462 """ 1463 1464 out = StringIO() 1465 argument_var = [name for type,name in self.call_arg] 1466 # define the complex number CI = 0+1j 1467 if add_i: 1468 out.write(self.ci_definition) 1469 1470 for type, name in self.declaration: 1471 if type.startswith('list'): 1472 type = type[5:] 1473 if name.startswith('P'): 1474 size = 4 1475 elif not 'tmp' in name: 1476 continue 1477 #should be define in the header 1478 elif name[0] in ['F','V']: 1479 if aloha.loop_mode: 1480 size = 8 1481 else: 1482 size = 6 1483 elif name[0] == 'S': 1484 if aloha.loop_mode: 1485 size = 5 1486 else: 1487 size = 3 1488 elif name[0] in ['R','T']: 1489 if aloha.loop_mode: 1490 size = 20 1491 else: 1492 size = 18 1493 1494 out.write(' %s %s[%s];\n' % (self.type2def[type], name, size)) 1495 elif (type, name) not in self.call_arg: 1496 out.write(' %s %s;\n' % (self.type2def[type], name)) 1497 1498 return out.getvalue()
1499
1500 - def get_foot_txt(self):
1501 """Prototype for language specific footer""" 1502 return '}\n'
1503
1504 - def get_momenta_txt(self):
1505 """Define the Header of the fortran file. This include 1506 - momentum conservation 1507 - definition of the impulsion""" 1508 1509 out = StringIO() 1510 1511 # Define all the required momenta 1512 p = [] # a list for keeping track how to write the momentum 1513 1514 signs = self.get_momentum_conservation_sign() 1515 1516 for i,type in enumerate(self.particles): 1517 if self.declaration.is_used('OM%s' % (i+1)): 1518 out.write(" OM{0} = {1};\n if (M{0} != {1})\n OM{0}={2}/(M{0}*M{0});\n".format( 1519 i+1, self.change_number_format(0), self.change_number_format(1))) 1520 1521 if i+1 == self.outgoing: 1522 out_type = type 1523 out_size = self.type_to_size[type] 1524 continue 1525 elif self.offshell: 1526 p.append('{0}{1}{2}[%(i)s]'.format(signs[i],type,i+1,type)) 1527 1528 if self.declaration.is_used('P%s' % (i+1)): 1529 self.get_one_momenta_def(i+1, out) 1530 1531 # define the resulting momenta 1532 if self.offshell: 1533 energy_pos = out_size -2 1534 type = self.particles[self.outgoing-1] 1535 if aloha.loop_mode: 1536 size_p = 4 1537 else: 1538 size_p = 2 1539 1540 for i in range(size_p): 1541 dict_energy = {'i':i} 1542 out.write(' %s%s[%s] = %s;\n' % (type,self.outgoing, i, 1543 ''.join(p) % dict_energy)) 1544 if self.declaration.is_used('P%s' % self.outgoing): 1545 self.get_one_momenta_def(self.outgoing, out) 1546 1547 1548 # Returning result 1549 return out.getvalue()
1550
1551 - def get_one_momenta_def(self, i, strfile):
1552 1553 type = self.particles[i-1] 1554 1555 if aloha.loop_mode: 1556 template ='P%(i)d[%(j)d] = %(sign)s%(type)s%(i)d[%(nb)d];\n' 1557 else: 1558 template ='P%(i)d[%(j)d] = %(sign)s%(type)s%(i)d[%(nb2)d]%(operator)s;\n' 1559 1560 nb2 = 0 1561 for j in range(4): 1562 if not aloha.loop_mode: 1563 nb = j 1564 if j == 0: 1565 assert not aloha.mp_precision 1566 operator = self.realoperator # not suppose to pass here in mp 1567 elif j == 1: 1568 nb2 += 1 1569 elif j == 2: 1570 assert not aloha.mp_precision 1571 operator = self.imagoperator # not suppose to pass here in mp 1572 elif j ==3: 1573 nb2 -= 1 1574 else: 1575 operator ='' 1576 nb = j 1577 nb2 = j 1578 strfile.write(template % {'j':j,'type': type, 'i': i, 1579 'nb': nb, 'nb2': nb2, 'operator':operator, 1580 'sign': self.get_P_sign(i)})
1581 1582
1583 - def define_expression(self):
1584 """Write the helicity amplitude in C++ format""" 1585 1586 out = StringIO() 1587 1588 if self.routine.contracted: 1589 for name,obj in self.routine.contracted.items(): 1590 out.write(' %s = %s;\n' % (name, self.write_obj(obj))) 1591 self.declaration.add(('complex', name)) 1592 1593 for name, (fct, objs) in self.routine.fct.items(): 1594 format = ' %s = %s;\n' % (name, self.get_fct_format(fct)) 1595 out.write(format % ','.join([self.write_obj(obj) for obj in objs])) 1596 1597 1598 1599 numerator = self.routine.expr 1600 if not 'Coup(1)' in self.routine.infostr: 1601 coup_name = 'COUP' 1602 else: 1603 coup_name = '%s' % self.change_number_format(1) 1604 if not self.offshell: 1605 if coup_name == 'COUP': 1606 out.write(' vertex = COUP*%s;\n' % self.write_obj(numerator.get_rep([0]))) 1607 else: 1608 out.write(' vertex = %s;\n' % self.write_obj(numerator.get_rep([0]))) 1609 else: 1610 OffShellParticle = '%s%d' % (self.particles[self.offshell-1],\ 1611 self.offshell) 1612 if 'L' not in self.tag: 1613 coeff = 'denom' 1614 if not aloha.complex_mass: 1615 if self.routine.denominator: 1616 out.write(' denom = %(COUP)s/(%(denom)s)\n' % {'COUP': coup_name,\ 1617 'denom':self.write_obj(self.routine.denominator)}) 1618 else: 1619 out.write(' denom = %(coup)s/((P%(i)s[0]*P%(i)s[0])-(P%(i)s[1]*P%(i)s[1])-(P%(i)s[2]*P%(i)s[2])-(P%(i)s[3]*P%(i)s[3]) - M%(i)s * (M%(i)s -cI* W%(i)s));\n' % \ 1620 {'i': self.outgoing, 'coup': coup_name}) 1621 else: 1622 if self.routine.denominator: 1623 raise Exception, 'modify denominator are not compatible with complex mass scheme' 1624 1625 out.write(' denom = %(coup)s/((P%(i)s[0]*P%(i)s[0])-(P%(i)s[1]*P%(i)s[1])-(P%(i)s[2]*P%(i)s[2])-(P%(i)s[3]*P%(i)s[3]) - (M%(i)s*M%(i)s));\n' % \ 1626 {'i': self.outgoing, 'coup': coup_name}) 1627 self.declaration.add(('complex','denom')) 1628 if aloha.loop_mode: 1629 ptype = 'list_complex' 1630 else: 1631 ptype = 'list_double' 1632 self.declaration.add((ptype,'P%s' % self.outgoing)) 1633 else: 1634 coeff = 'COUP' 1635 1636 for ind in numerator.listindices(): 1637 out.write(' %s[%d]= %s*%s;\n' % (self.outname, 1638 self.pass_to_HELAS(ind), coeff, 1639 self.write_obj(numerator.get_rep(ind)))) 1640 return out.getvalue()
1641 1642 remove_double = re.compile('std::complex<double> (?P<name>[\w]+)\[\]')
1643 - def define_symmetry(self, new_nb, couplings=None):
1644 """Write the call for symmetric routines""" 1645 number = self.offshell 1646 arguments = [name for format, name in self.define_argument_list()] 1647 new_name = self.name.rsplit('_')[0] + '_%s' % new_nb 1648 output = '%(spin)s%(id)d' % { 1649 'spin': self.particles[self.offshell -1], 1650 'id': self.outgoing} 1651 return '%s\n %s(%s,%s);\n}' % \ 1652 (self.get_header_txt(new_name, couplings, mode='no_include'), 1653 self.name, ','.join(arguments), output)
1654
1655 - def get_h_text(self,couplings=None):
1656 """Return the full contents of the .h file""" 1657 1658 h_string = StringIO() 1659 if not self.mode == 'no_include': 1660 h_string.write('#ifndef '+ self.name + '_guard\n') 1661 h_string.write('#define ' + self.name + '_guard\n') 1662 h_string.write('#include <complex>\n\n') 1663 1664 h_header = self.get_header_txt(mode='no_include__is_h', couplings=couplings) 1665 h_string.write(h_header) 1666 1667 for elem in self.routine.symmetries: 1668 symmetryhead = h_header.replace( \ 1669 self.name,self.name[0:-1]+'%s' %(elem)) 1670 h_string.write(symmetryhead) 1671 1672 if not self.mode == 'no_include': 1673 h_string.write('#endif\n\n') 1674 1675 return h_string.getvalue()
1676 1677
1678 - def write_combined_cc(self, lor_names, offshell=None, sym=True, mode=''):
1679 "Return the content of the .cc file linked to multiple lorentz call." 1680 1681 # Set some usefull command 1682 if offshell is None: 1683 offshell = self.offshell 1684 1685 name = combine_name(self.routine.name, lor_names, offshell, self.tag) 1686 self.name = name 1687 # write head - momenta - body - foot 1688 text = StringIO() 1689 routine = StringIO() 1690 data = {} # for the formating of the line 1691 1692 # write header 1693 new_couplings = ['COUP%s' % (i+1) for i in range(len(lor_names)+1)] 1694 text.write(self.get_header_txt(name=name, couplings=new_couplings, mode=mode)) 1695 1696 # Define which part of the routine should be called 1697 data['addon'] = ''.join(self.tag) + '_%s' % self.offshell 1698 1699 # how to call the routine 1700 argument = [name for format, name in self.define_argument_list(new_couplings)] 1701 index= argument.index('COUP1') 1702 data['before_coup'] = ','.join(argument[:index]) 1703 data['after_coup'] = ','.join(argument[index+len(lor_names)+1:]) 1704 if data['after_coup']: 1705 data['after_coup'] = ',' + data['after_coup'] 1706 1707 lor_list = (self.routine.name,) + lor_names 1708 line = " %(name)s%(addon)s(%(before_coup)s,%(coup)s%(after_coup)s,%(out)s);\n" 1709 main = '%(spin)s%(id)d' % {'spin': self.particles[self.offshell -1], 1710 'id': self.outgoing} 1711 for i, name in enumerate(lor_list): 1712 data['name'] = name 1713 data['coup'] = 'COUP%d' % (i+1) 1714 if i == 0: 1715 if not offshell: 1716 data['out'] = 'vertex' 1717 else: 1718 data['out'] = main 1719 elif i==1: 1720 if self.offshell: 1721 type = self.particles[self.offshell-1] 1722 self.declaration.add(('list_complex','%stmp' % type)) 1723 else: 1724 type = '' 1725 self.declaration.add(('complex','%stmp' % type)) 1726 data['out'] = '%stmp' % type 1727 routine.write(line % data) 1728 if i: 1729 if not offshell: 1730 routine.write( ' vertex = vertex + tmp;\n') 1731 else: 1732 size = self.type_to_size[self.particles[offshell -1]] -2 1733 routine.write(""" i= %s;\nwhile (i < %s)\n{\n""" % (self.momentum_size, self.momentum_size+size)) 1734 routine.write(" %(main)s[i] = %(main)s[i] + %(tmp)s[i];\n i++;\n" %\ 1735 {'main': main, 'tmp': data['out']}) 1736 routine.write('}\n') 1737 self.declaration.add(('int','i')) 1738 self.declaration.discard(('complex','COUP')) 1739 self.declaration.discard(('complex', 'denom')) 1740 if self.outgoing: 1741 self.declaration.discard(('list_double', 'P%s' % self.outgoing)) 1742 self.declaration.discard(('double', 'OM%s' % self.outgoing)) 1743 for name in aloha_lib.KERNEL.reduced_expr2: 1744 self.declaration.discard(('complex', name)) 1745 1746 #clean pointless declaration 1747 #self.declaration.discard 1748 text.write(self.get_declaration_txt(add_i=False)) 1749 text.write(routine.getvalue()) 1750 text.write(self.get_foot_txt()) 1751 1752 text = text.getvalue() 1753 return text
1754 1755
1756 - def write(self, **opt):
1757 """Write the .h and .cc files""" 1758 1759 cc_text = WriteALOHA.write(self, **opt) 1760 h_text = self.get_h_text() 1761 1762 # write in two file 1763 if self.out_path: 1764 writer_h = writers.CPPWriter(self.out_path[:-len(self.extension)] + ".h") 1765 commentstring = 'This File is Automatically generated by ALOHA \n' 1766 commentstring += 'The process calculated in this file is: \n' 1767 commentstring += self.routine.infostr + '\n' 1768 writer_h.write_comments(commentstring) 1769 writer_h.writelines(h_text) 1770 1771 return h_text, cc_text
1772 1773 1774
1775 - def write_combined(self, lor_names, mode='', offshell=None, **opt):
1776 """Write the .h and .cc files associated to the combined file""" 1777 1778 # Set some usefull command 1779 if offshell is None: 1780 sym = 1 1781 offshell = self.offshell 1782 else: 1783 sym = None 1784 1785 if mode == 'self': 1786 # added to another file 1787 self.mode = 'no_include' 1788 1789 1790 1791 #h_text = self.write_combined_h(lor_names, offshell, **opt) 1792 cc_text, h_text = StringIO() , StringIO() 1793 cc_text.write(self.write_combined_cc(lor_names, offshell, mode=mode,**opt)) 1794 couplings = ['COUP%d' % (i+1) for i in range(len(lor_names)+1)] 1795 1796 if mode == 'self': 1797 self.mode = 'self' 1798 h_text.write(self.get_h_text(couplings=couplings)) 1799 1800 #ADD SYMETRY 1801 if sym: 1802 for elem in self.routine.symmetries: 1803 self.mode = 'no_include' 1804 cc_text.write( self.write_combined_cc(lor_names, elem)) 1805 1806 1807 if self.out_path: 1808 # Prepare a specific file 1809 path = os.path.join(os.path.dirname(self.out_path), self.name) 1810 commentstring = 'This File is Automatically generated by ALOHA \n' 1811 1812 writer_h = writers.CPPWriter(path + ".h") 1813 writer_h.write_comments(commentstring) 1814 writer_h.writelines(h_text.getvalue()) 1815 1816 writer_cc = writers.CPPWriter(path + ".cc") 1817 writer_cc.write_comments(commentstring) 1818 writer_cc.writelines(cc_text.getvalue()) 1819 1820 return h_text.getvalue(), cc_text.getvalue()
1821
1822 1823 -class ALOHAWriterForGPU(ALOHAWriterForCPP):
1824 1825 extension = '.cu' 1826 realoperator = '.re' 1827 imagoperator = '.im' 1828 ci_definition = 'complex<double> cI = mkcmplx(0., 1.);\n' 1829
1830 - def get_header_txt(self, name=None, couplings=None, mode=''):
1831 """Define the Header of the fortran file. This include 1832 - function tag 1833 - definition of variable 1834 """ 1835 text = StringIO() 1836 if not 'is_h' in mode: 1837 text.write('__device__=__forceinclude__\n') 1838 text.write(ALOHAWriterForCPP.get_header_txt(self, name, couplings, mode)) 1839 return text.getvalue()
1840
1841 - def get_h_text(self,couplings=None):
1842 """Return the full contents of the .h file""" 1843 1844 h_string = StringIO() 1845 if not self.mode == 'no_include': 1846 h_string.write('#ifndef '+ self.name + '_guard\n') 1847 h_string.write('#define ' + self.name + '_guard\n') 1848 h_string.write('#include "cmplx.h"\n') 1849 h_string.write('using namespace std;\n\n') 1850 1851 h_header = self.get_header_txt(mode='no_include__is_h', couplings=couplings) 1852 h_string.write(h_header) 1853 1854 for elem in self.routine.symmetries: 1855 symmetryhead = h_header.replace( \ 1856 self.name,self.name[0:-1]+'%s' %(elem)) 1857 h_string.write(symmetryhead) 1858 1859 if not self.mode == 'no_include': 1860 h_string.write('#endif\n\n') 1861 1862 return h_string.getvalue()
1863
1864 1865 -class ALOHAWriterForPython(WriteALOHA):
1866 """ A class for returning a file/a string for python evaluation """ 1867 1868 extension = '.py' 1869 writer = writers.PythonWriter 1870 1871 @staticmethod
1872 - def change_number_format(obj, pure_complex=''):
1873 change_number_format = ALOHAWriterForPython.change_number_format 1874 if obj.real == 0 and obj.imag: 1875 if int(obj.imag) == obj.imag: 1876 return '%ij' % obj.imag 1877 else: 1878 return change_number_format(obj.imag, pure_complex='j') 1879 elif obj.imag != 0: 1880 return '(%s+%s)' % (change_number_format(obj.real), 1881 change_number_format(obj.imag, pure_complex='j')) 1882 elif obj.imag == 0: 1883 if int(obj.real) == obj: 1884 return '%i%s' % (obj.real,pure_complex) 1885 obj = obj.real 1886 tmp = Fraction(str(obj)) 1887 tmp = tmp.limit_denominator(100) 1888 if not abs(tmp - obj) / abs(tmp + obj) < 1e-8: 1889 out = str(obj) 1890 elif tmp.denominator != 1: 1891 out = '%i%s/%i' % (tmp.numerator, pure_complex, tmp.denominator) 1892 else: 1893 out = '%i%s' % (tmp.numerator, pure_complex) 1894 return out
1895 1896
1897 - def shift_indices(self, match):
1898 """shift the indices for non impulsion object""" 1899 if match.group('var').startswith('P'): 1900 shift = 0 1901 else: 1902 shift = -1 + self.momentum_size 1903 1904 return '%s[%s]' % (match.group('var'), int(match.group('num')) + shift)
1905
1906 - def change_var_format(self, name):
1907 """Formatting the variable name to Python format 1908 start to count at zero. 1909 No neeed to define the variable in python -> no need to keep track of 1910 the various variable 1911 """ 1912 1913 if '_' not in name: 1914 self.declaration.add((name.type, name)) 1915 else: 1916 self.declaration.add(('', name.split('_',1)[0])) 1917 name = re.sub('(?P<var>\w*)_(?P<num>\d+)$', self.shift_indices , name) 1918 1919 return name
1920
1921 - def get_fct_format(self, fct):
1922 """Put the function in the correct format""" 1923 if not hasattr(self, 'fct_format'): 1924 one = self.change_number_format(1) 1925 self.fct_format = {'csc' : '{0}/cmath.cos(%s)'.format(one), 1926 'sec': '{0}/cmath.sin(%s)'.format(one), 1927 'acsc': 'cmath.asin({0}/(%s))'.format(one), 1928 'asec': 'cmath.acos({0}/(%s))'.format(one), 1929 're': ' complex(%s).real', 1930 'im': 'complex(%s).imag', 1931 'cmath.sqrt': 'cmath.sqrt(%s)', 1932 'sqrt': 'cmath.sqrt(%s)', 1933 'pow': 'pow(%s, %s)', 1934 'complexconjugate': 'complex(%s).conjugate()', 1935 '/' : '{0}/%s'.format(one), 1936 'abs': 'cmath.fabs(%s)' 1937 } 1938 1939 if fct in self.fct_format: 1940 return self.fct_format[fct] 1941 elif hasattr(cmath, fct): 1942 self.declaration.add(('fct', fct)) 1943 return 'cmath.{0}(%s)'.format(fct) 1944 else: 1945 raise Exception, "Unable to handle function name %s (no special rule defined and not in cmath)" % fct
1946
1947 - def define_expression(self):
1948 """Define the functions in a 100% way """ 1949 1950 out = StringIO() 1951 1952 if self.routine.contracted: 1953 for name,obj in self.routine.contracted.items(): 1954 out.write(' %s = %s\n' % (name, self.write_obj(obj))) 1955 1956 def sort_fct(a, b): 1957 if len(a) < len(b): 1958 return -1 1959 elif len(a) > len(b): 1960 return 1 1961 elif a < b: 1962 return -1 1963 else: 1964 return +1
1965 1966 keys = self.routine.fct.keys() 1967 keys.sort(sort_fct) 1968 for name in keys: 1969 fct, objs = self.routine.fct[name] 1970 format = ' %s = %s\n' % (name, self.get_fct_format(fct)) 1971 try: 1972 text = format % ','.join([self.write_obj(obj) for obj in objs]) 1973 except TypeError: 1974 text = format % tuple([self.write_obj(obj) for obj in objs]) 1975 finally: 1976 out.write(text) 1977 1978 1979 1980 numerator = self.routine.expr 1981 if not 'Coup(1)' in self.routine.infostr: 1982 coup_name = 'COUP' 1983 else: 1984 coup_name = '%s' % self.change_number_format(1) 1985 1986 if not self.offshell: 1987 if coup_name == 'COUP': 1988 out.write(' vertex = COUP*%s\n' % self.write_obj(numerator.get_rep([0]))) 1989 else: 1990 out.write(' vertex = %s\n' % self.write_obj(numerator.get_rep([0]))) 1991 else: 1992 OffShellParticle = '%s%d' % (self.particles[self.offshell-1],\ 1993 self.offshell) 1994 1995 if not 'L' in self.tag: 1996 coeff = 'denom' 1997 if not aloha.complex_mass: 1998 if self.routine.denominator: 1999 out.write(' denom = %(COUP)s/(%(denom)s)\n' % {'COUP': coup_name,\ 2000 'denom':self.write_obj(self.routine.denominator)}) 2001 else: 2002 out.write(' denom = %(coup)s/(P%(i)s[0]**2-P%(i)s[1]**2-P%(i)s[2]**2-P%(i)s[3]**2 - M%(i)s * (M%(i)s -1j* W%(i)s))\n' % 2003 {'i': self.outgoing,'coup':coup_name}) 2004 else: 2005 if self.routine.denominator: 2006 raise Exception, 'modify denominator are not compatible with complex mass scheme' 2007 2008 out.write(' denom = %(coup)s/(P%(i)s[0]**2-P%(i)s[1]**2-P%(i)s[2]**2-P%(i)s[3]**2 - M%(i)s**2)\n' % 2009 {'i': self.outgoing,'coup':coup_name}) 2010 else: 2011 coeff = 'COUP' 2012 2013 for ind in numerator.listindices(): 2014 out.write(' %s[%d]= %s*%s\n' % (self.outname, 2015 self.pass_to_HELAS(ind), coeff, 2016 self.write_obj(numerator.get_rep(ind)))) 2017 return out.getvalue() 2018
2019 - def get_foot_txt(self):
2020 if not self.offshell: 2021 return ' return vertex\n\n' 2022 else: 2023 return ' return %s\n\n' % (self.outname)
2024 2025
2026 - def get_header_txt(self, name=None, couplings=None, mode=''):
2027 """Define the Header of the fortran file. This include 2028 - function tag 2029 - definition of variable 2030 """ 2031 if name is None: 2032 name = self.name 2033 2034 out = StringIO() 2035 out.write("import cmath\n") 2036 if self.mode == 'mg5': 2037 out.write('import aloha.template_files.wavefunctions as wavefunctions\n') 2038 else: 2039 out.write('import wavefunctions\n') 2040 2041 2042 # define the type of function and argument 2043 2044 arguments = [arg for format, arg in self.define_argument_list(couplings)] 2045 out.write('def %(name)s(%(args)s):\n' % \ 2046 {'name': name, 'args': ','.join(arguments)}) 2047 2048 return out.getvalue()
2049
2050 - def get_momenta_txt(self):
2051 """Define the Header of the fortran file. This include 2052 - momentum conservation 2053 - definition of the impulsion""" 2054 2055 out = StringIO() 2056 2057 # Define all the required momenta 2058 p = [] # a list for keeping track how to write the momentum 2059 2060 signs = self.get_momentum_conservation_sign() 2061 2062 for i,type in enumerate(self.particles): 2063 if self.declaration.is_used('OM%s' % (i+1)): 2064 out.write(" OM{0} = 0.0\n if (M{0}): OM{0}=1.0/M{0}**2\n".format( (i+1) )) 2065 if i+1 == self.outgoing: 2066 out_type = type 2067 out_size = self.type_to_size[type] 2068 continue 2069 elif self.offshell: 2070 p.append('{0}{1}{2}[%(i)s]'.format(signs[i],type,i+1)) 2071 2072 if self.declaration.is_used('P%s' % (i+1)): 2073 self.get_one_momenta_def(i+1, out) 2074 2075 # define the resulting momenta 2076 if self.offshell: 2077 type = self.particles[self.outgoing-1] 2078 out.write(' %s%s = wavefunctions.WaveFunction(size=%s)\n' % (type, self.outgoing, out_size)) 2079 if aloha.loop_mode: 2080 size_p = 4 2081 else: 2082 size_p = 2 2083 for i in range(size_p): 2084 dict_energy = {'i':i} 2085 2086 out.write(' %s%s[%s] = %s\n' % (type,self.outgoing, i, 2087 ''.join(p) % dict_energy)) 2088 2089 self.get_one_momenta_def(self.outgoing, out) 2090 2091 2092 # Returning result 2093 return out.getvalue()
2094
2095 - def get_one_momenta_def(self, i, strfile):
2096 """return the string defining the momentum""" 2097 2098 type = self.particles[i-1] 2099 2100 main = ' P%d = [' % i 2101 if aloha.loop_mode: 2102 template ='%(sign)s%(type)s%(i)d[%(nb)d]' 2103 else: 2104 template ='%(sign)scomplex(%(type)s%(i)d[%(nb2)d])%(operator)s' 2105 2106 nb2 = 0 2107 strfile.write(main) 2108 data = [] 2109 for j in range(4): 2110 if not aloha.loop_mode: 2111 nb = j 2112 if j == 0: 2113 assert not aloha.mp_precision 2114 operator = '.real' # not suppose to pass here in mp 2115 elif j == 1: 2116 nb2 += 1 2117 elif j == 2: 2118 assert not aloha.mp_precision 2119 operator = '.imag' # not suppose to pass here in mp 2120 elif j ==3: 2121 nb2 -= 1 2122 else: 2123 operator ='' 2124 nb = j 2125 nb2 = j 2126 data.append(template % {'j':j,'type': type, 'i': i, 2127 'nb': nb, 'nb2': nb2, 'operator':operator, 2128 'sign': self.get_P_sign(i)}) 2129 2130 strfile.write(', '.join(data)) 2131 strfile.write(']\n')
2132 2133
2134 - def define_symmetry(self, new_nb, couplings=None):
2135 number = self.offshell 2136 arguments = [name for format, name in self.define_argument_list()] 2137 new_name = self.name.rsplit('_')[0] + '_%s' % new_nb 2138 return '%s\n return %s(%s)' % \ 2139 (self.get_header_txt(new_name, couplings), self.name, ','.join(arguments))
2140
2141 - def write_combined(self, lor_names, mode='self', offshell=None):
2142 """Write routine for combine ALOHA call (more than one coupling)""" 2143 2144 # Set some usefull command 2145 if offshell is None: 2146 sym = 1 2147 offshell = self.offshell 2148 else: 2149 sym = None 2150 name = combine_name(self.routine.name, lor_names, offshell, self.tag) 2151 # write head - momenta - body - foot 2152 text = StringIO() 2153 data = {} # for the formating of the line 2154 2155 # write header 2156 new_couplings = ['COUP%s' % (i+1) for i in range(len(lor_names)+1)] 2157 text.write(self.get_header_txt(name=name, couplings=new_couplings)) 2158 2159 # Define which part of the routine should be called 2160 data['addon'] = ''.join(self.tag) + '_%s' % self.offshell 2161 2162 # how to call the routine 2163 argument = [name for format, name in self.define_argument_list(new_couplings)] 2164 index= argument.index('COUP1') 2165 data['before_coup'] = ','.join(argument[:index]) 2166 data['after_coup'] = ','.join(argument[index+len(lor_names)+1:]) 2167 if data['after_coup']: 2168 data['after_coup'] = ',' + data['after_coup'] 2169 2170 lor_list = (self.routine.name,) + lor_names 2171 line = " %(out)s = %(name)s%(addon)s(%(before_coup)s,%(coup)s%(after_coup)s)\n" 2172 main = '%(spin)s%(id)d' % {'spin': self.particles[self.offshell -1], 2173 'id': self.outgoing} 2174 for i, name in enumerate(lor_list): 2175 data['name'] = name 2176 data['coup'] = 'COUP%d' % (i+1) 2177 if i == 0: 2178 if not offshell: 2179 data['out'] = 'vertex' 2180 else: 2181 data['out'] = main 2182 elif i==1: 2183 data['out'] = 'tmp' 2184 text.write(line % data) 2185 if i: 2186 if not offshell: 2187 text.write( ' vertex += tmp\n') 2188 else: 2189 size = self.type_to_size[self.particles[offshell -1]] -2 2190 text.write(" for i in range(%s,%s):\n" % (self.momentum_size, self.momentum_size+size)) 2191 text.write(" %(main)s[i] += tmp[i]\n" %{'main': main}) 2192 2193 text.write(self.get_foot_txt()) 2194 2195 #ADD SYMETRY 2196 if sym: 2197 for elem in self.routine.symmetries: 2198 text.write(self.write_combined(lor_names, mode, elem)) 2199 2200 text = text.getvalue() 2201 if self.out_path: 2202 writer = self.writer(self.out_path) 2203 commentstring = 'This File is Automatically generated by ALOHA \n' 2204 commentstring += 'The process calculated in this file is: \n' 2205 commentstring += self.routine.infostr + '\n' 2206 writer.write_comments(commentstring) 2207 writer.writelines(text) 2208 2209 2210 return text
2211
2212 2213 -class Declaration_list(set):
2214
2215 - def is_used(self, var):
2216 if hasattr(self, 'var_name'): 2217 return var in self.var_name 2218 self.var_name = [name for type,name in self] 2219 return var in self.var_name
2220
2221 - def add(self,obj):
2222 if __debug__: 2223 type, name = obj 2224 samename = [t for t,n in self if n ==name] 2225 for type2 in samename: 2226 assert type2 == type, '%s is defined with two different type "%s" and "%s"' % \ 2227 (name, type2, type) 2228 2229 set.add(self,obj)
2230
2231 2232 -class WriterFactory(object):
2233
2234 - def __new__(cls, data, language, outputdir, tags):
2235 language = language.lower() 2236 if isinstance(data.expr, aloha_lib.SplitCoefficient): 2237 assert language == 'fortran' 2238 if 'MP' in tags: 2239 return ALOHAWriterForFortranLoopQP(data, outputdir) 2240 else: 2241 return ALOHAWriterForFortranLoop(data, outputdir) 2242 if language == 'fortran': 2243 if 'MP' in tags: 2244 return ALOHAWriterForFortranQP(data, outputdir) 2245 else: 2246 return ALOHAWriterForFortran(data, outputdir) 2247 elif language == 'python': 2248 return ALOHAWriterForPython(data, outputdir) 2249 elif language == 'cpp': 2250 return ALOHAWriterForCPP(data, outputdir) 2251 elif language == 'gpu': 2252 return ALOHAWriterForGPU(data, outputdir) 2253 else: 2254 raise Exception, 'Unknown output format'
2255 2256 2257 2258 #unknow_fct_template = """ 2259 #cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc 2260 # double complex %(fct_name)s(%(args)s) 2261 # implicit none 2262 #c Include Model parameter / coupling 2263 # include \"../MODEL/input.inc\" 2264 # include \"../MODEL/coupl.inc\" 2265 #c Defintion of the arguments 2266 #%(definitions)s 2267 # 2268 #c enter HERE the code corresponding to your function. 2269 #c The output value should be put to the %(fct_name)s variable. 2270 # 2271 # 2272 # return 2273 # end 2274 #cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc 2275 # 2276 #""" 2277 # 2278 #def write_template_fct(fct_name, nb_args, output_dir): 2279 # """create a template for function not recognized by ALOHA""" 2280 # 2281 # dico = {'fct_name' : fct_name, 2282 # 'args': ','.join(['S%i' %(i+1) for i in range(nb_args)]), 2283 # 'definitions': '\n'.join([' double complex S%i' %(i+1) for i in range(nb_args)])} 2284 # 2285 # ff = open(pjoin(output_dir, 'additional_aloha_function.f'), 'a') 2286 # ff.write(unknow_fct_template % dico) 2287 # ff.close() 2288