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