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