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 658 self.has_model_parameter = True 659 return '%s%s' % (aloha.aloha_prefix, name) 660 661 if '_' in name: 662 vtype = name.type 663 decla = name.split('_',1)[0] 664 self.declaration.add(('list_%s' % vtype, decla)) 665 else: 666 self.declaration.add((name.type, name)) 667 name = re.sub('(?P<var>\w*)_(?P<num>\d+)$', self.shift_indices , name) 668 return name 669
670 - def change_number_format(self, number):
671 """Formating the number""" 672 673 def isinteger(x): 674 try: 675 return int(x) == x 676 except TypeError: 677 return False
678 679 if isinteger(number): 680 out = '%s%s' % (str(int(number)),self.format) 681 elif isinstance(number, complex): 682 if number.imag: 683 if number.real: 684 out = '(%s + %s*CI)' % (self.change_number_format(number.real), \ 685 self.change_number_format(number.imag)) 686 else: 687 if number.imag == 1: 688 out = 'CI' 689 elif number.imag == -1: 690 out = '-CI' 691 else: 692 out = '%s * CI' % self.change_number_format(number.imag) 693 else: 694 out = '%s' % (self.change_number_format(number.real)) 695 else: 696 tmp = Fraction(str(number)) 697 tmp = tmp.limit_denominator(100) 698 if not abs(tmp - number) / abs(tmp + number) < 1e-8: 699 out = '%s%s' % (number, self.format) 700 else: 701 out = '%s%s/%s%s' % (tmp.numerator, self.format, tmp.denominator, self.format) 702 return out 703
704 - def define_expression(self):
705 """Define the functions in a 100% way """ 706 707 out = StringIO() 708 709 if self.routine.contracted: 710 for name,obj in self.routine.contracted.items(): 711 out.write(' %s = %s\n' % (name, self.write_obj(obj))) 712 self.declaration.add(('complex', name)) 713 714 715 def sort_fct(a, b): 716 if len(a) < len(b): 717 return -1 718 elif len(a) > len(b): 719 return 1 720 elif a < b: 721 return -1 722 else: 723 return +1
724 725 keys = self.routine.fct.keys() 726 keys.sort(sort_fct) 727 for name in keys: 728 fct, objs = self.routine.fct[name] 729 730 format = ' %s = %s\n' % (name, self.get_fct_format(fct)) 731 try: 732 text = format % ','.join([self.write_obj(obj) for obj in objs]) 733 except TypeError: 734 text = format % tuple([self.write_obj(obj) for obj in objs]) 735 finally: 736 out.write(text) 737 738 739 numerator = self.routine.expr 740 if not 'Coup(1)' in self.routine.infostr: 741 coup_name = 'COUP' 742 else: 743 coup_name = '%s' % self.change_number_format(1) 744 745 746 if not self.offshell: 747 if coup_name == 'COUP': 748 out.write(' vertex = COUP*%s\n' % self.write_obj(numerator.get_rep([0]))) 749 else: 750 out.write(' vertex = %s\n' % self.write_obj(numerator.get_rep([0]))) 751 else: 752 OffShellParticle = '%s%d' % (self.particles[self.offshell-1],\ 753 self.offshell) 754 if 'L' not in self.tag: 755 coeff = 'denom*' 756 if not aloha.complex_mass: 757 if self.routine.denominator: 758 out.write(' denom = %(COUP)s/(%(denom)s)\n' % {'COUP': coup_name,\ 759 'denom':self.write_obj(self.routine.denominator)}) 760 else: 761 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' % \ 762 {'i': self.outgoing, 'COUP': coup_name}) 763 else: 764 if self.routine.denominator: 765 raise Exception, 'modify denominator are not compatible with complex mass scheme' 766 767 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' % \ 768 {'i': self.outgoing, 'COUP': coup_name}) 769 self.declaration.add(('complex','denom')) 770 if aloha.loop_mode: 771 ptype = 'list_complex' 772 else: 773 ptype = 'list_double' 774 self.declaration.add((ptype,'P%s' % self.outgoing)) 775 else: 776 if coup_name == 'COUP': 777 coeff = 'COUP*' 778 else: 779 coeff = '' 780 to_order = {} 781 for ind in numerator.listindices(): 782 to_order[self.pass_to_HELAS(ind)] = \ 783 ' %s(%d)= %s%s\n' % (self.outname, self.pass_to_HELAS(ind)+1, 784 coeff, self.write_obj(numerator.get_rep(ind))) 785 key = to_order.keys() 786 key.sort() 787 for i in key: 788 out.write(to_order[i]) 789 return out.getvalue() 790
791 - def define_symmetry(self, new_nb, couplings=None):
792 return ''
793 #number = self.offshell 794 #arguments = [name for format, name in self.define_argument_list()] 795 #new_name = self.name.rsplit('_')[0] + '_%s' % new_nb 796 #return '%s\n call %s(%s)' % \ 797 # (self.get_header_txt(new_name, couplings), self.name, ','.join(arguments)) 798
799 - def get_foot_txt(self):
800 return 'end\n\n'
801
802 - def write_combined(self, lor_names, mode='self', offshell=None):
803 """Write routine for combine ALOHA call (more than one coupling)""" 804 805 # Set some usefull command 806 if offshell is None: 807 sym = 1 808 offshell = self.offshell 809 else: 810 sym = None 811 name = combine_name(self.routine.name, lor_names, offshell, self.tag) 812 self.name = name 813 # write head - momenta - body - foot 814 text = StringIO() 815 routine = StringIO() 816 data = {} # for the formating of the line 817 818 # write header 819 new_couplings = ['COUP%s' % (i+1) for i in range(len(lor_names)+1)] 820 text.write(self.get_header_txt(name=name, couplings=new_couplings)) 821 822 # Define which part of the routine should be called 823 data['addon'] = ''.join(self.tag) + '_%s' % self.offshell 824 825 # how to call the routine 826 argument = [name for format, name in self.define_argument_list(new_couplings)] 827 index= argument.index('COUP1') 828 data['before_coup'] = ','.join(argument[:index]) 829 data['after_coup'] = ','.join(argument[index+len(lor_names)+1:]) 830 if data['after_coup']: 831 data['after_coup'] = ',' + data['after_coup'] 832 833 lor_list = (self.routine.name,) + lor_names 834 line = " call %(name)s%(addon)s(%(before_coup)s,%(coup)s%(after_coup)s,%(out)s)\n" 835 main = '%(spin)s%(id)d' % {'spin': self.particles[self.offshell -1], 836 'id': self.outgoing} 837 for i, name in enumerate(lor_list): 838 data['name'] = name 839 data['coup'] = 'COUP%d' % (i+1) 840 if i == 0: 841 if not offshell: 842 data['out'] = 'vertex' 843 else: 844 data['out'] = main 845 elif i==1: 846 if self.offshell: 847 type = self.particles[self.offshell-1] 848 self.declaration.add(('list_complex','%stmp' % type)) 849 else: 850 type = '' 851 self.declaration.add(('complex','%stmp' % type)) 852 data['out'] = '%stmp' % type 853 routine.write(line % data) 854 if i: 855 if not offshell: 856 routine.write( ' vertex = vertex + tmp\n') 857 else: 858 size = self.type_to_size[self.particles[offshell -1]] -2 859 routine.write(" do i = %s, %s\n" % (self.momentum_size+1, self.momentum_size+size)) 860 routine.write(" %(main)s(i) = %(main)s(i) + %(tmp)s(i)\n" %\ 861 {'main': main, 'tmp': data['out']}) 862 routine.write(' enddo\n') 863 self.declaration.add(('int','i')) 864 865 self.declaration.discard(('complex','COUP')) 866 for name in aloha_lib.KERNEL.reduced_expr2: 867 self.declaration.discard(('complex', name)) 868 869 #clean pointless declaration 870 #self.declaration.discard 871 872 873 text.write(self.get_declaration_txt()) 874 text.write(routine.getvalue()) 875 text.write(self.get_foot_txt()) 876 877 878 text = text.getvalue() 879 if self.out_path: 880 writer = self.writer(self.out_path,'a') 881 commentstring = 'This File is Automatically generated by ALOHA \n' 882 commentstring += 'The process calculated in this file is: \n' 883 commentstring += self.routine.infostr + '\n' 884 writer.write_comments(commentstring) 885 writer.writelines(text) 886 return text
887
888 -class QP(object):
889 """routines for writing out Fortran""" 890 891 type2def = {} 892 type2def['int'] = 'integer*4' 893 type2def['double'] = 'real*16' 894 type2def['complex'] = 'complex*32' 895 format = 'q0' 896
897 -class ALOHAWriterForFortranQP(QP, ALOHAWriterForFortran):
898
899 - def __init__(self, *arg):
900 return ALOHAWriterForFortran.__init__(self, *arg)
901
902 -class ALOHAWriterForFortranLoop(ALOHAWriterForFortran):
903 """routines for writing out Fortran""" 904
905 - def __init__(self, abstract_routine, dirpath):
906 907 ALOHAWriterForFortran.__init__(self, abstract_routine, dirpath) 908 # position of the outgoing in particle list 909 self.l_id = [int(c[1:]) for c in abstract_routine.tag if c[0] == 'L'][0] 910 self.l_helas_id = self.l_id # expected position for the argument list 911 if 'C%s' %((self.l_id + 1) // 2) in abstract_routine.tag: 912 #flip the outgoing tag if in conjugate 913 self.l_helas_id += self.l_id % 2 - (self.l_id +1) % 2
914 915
916 - def define_expression(self):
917 """Define the functions in a 100% way """ 918 919 out = StringIO() 920 921 if self.routine.contracted: 922 for name,obj in self.routine.contracted.items(): 923 out.write(' %s = %s\n' % (name, self.write_obj(obj))) 924 self.declaration.add(('complex', name)) 925 926 if not 'Coup(1)' in self.routine.infostr: 927 coup = True 928 else: 929 coup = False 930 931 rank = self.routine.expr.get_max_rank() 932 poly_object = q_polynomial.Polynomial(rank) 933 nb_coeff = q_polynomial.get_number_of_coefs_for_rank(rank) 934 size = self.type_to_size[self.particles[self.l_id-1]] - 2 935 for K in range(size): 936 for J in range(nb_coeff): 937 data = poly_object.get_coef_at_position(J) 938 arg = [data.count(i) for i in range(4)] # momentum 939 arg += [0] * (K) + [1] + [0] * (size-1-K) 940 try: 941 expr = self.routine.expr[tuple(arg)] 942 except KeyError: 943 expr = None 944 for ind in self.routine.expr.values()[0].listindices(): 945 if expr: 946 data = expr.get_rep(ind) 947 else: 948 data = 0 949 if data and coup: 950 out.write(' COEFF(%s,%s,%s)= coup*%s\n' % ( 951 self.pass_to_HELAS(ind)+1-self.momentum_size, 952 J, K+1, self.write_obj(data))) 953 else: 954 out.write(' COEFF(%s,%s,%s)= %s\n' % ( 955 self.pass_to_HELAS(ind)+1-self.momentum_size, 956 J, K+1, self.write_obj(data))) 957 958 959 return out.getvalue()
960
961 - def get_declaration_txt(self):
962 """ Prototype for how to write the declaration of variable""" 963 964 out = StringIO() 965 out.write('implicit none\n') 966 # define the complex number CI = 0+1j 967 if 'MP' in self.tag: 968 out.write(' complex*32 CI\n') 969 else: 970 out.write(' complex*16 CI\n') 971 out.write(' parameter (CI=(%s,%s))\n' % 972 (self.change_number_format(0),self.change_number_format(1))) 973 argument_var = [name for type,name in self.call_arg] 974 for type, name in self.declaration: 975 if type.startswith('list'): 976 type = type[5:] 977 #determine the size of the list 978 if name.startswith('P'): 979 size='0:3' 980 elif name in argument_var: 981 size ='*' 982 elif name[0] in ['F','V']: 983 if aloha.loop_mode: 984 size = 8 985 else: 986 size = 6 987 elif name[0] == 'S': 988 if aloha.loop_mode: 989 size = 5 990 else: 991 size = 3 992 elif name[0] in ['R','T']: 993 if aloha.loop_mode: 994 size = 20 995 else: 996 size = 18 997 elif name == 'coeff': 998 out.write("include 'coef_specs.inc'\n") 999 size = 'MAXLWFSIZE,0:VERTEXMAXCOEFS-1,MAXLWFSIZE' 1000 1001 out.write(' %s %s(%s)\n' % (self.type2def[type], name, size)) 1002 elif type == 'fct': 1003 if name.upper() in ['EXP','LOG','SIN','COS','ASIN','ACOS']: 1004 continue 1005 out.write(' %s %s\n' % (self.type2def['complex'], name)) 1006 out.write(' external %s\n' % (name)) 1007 else: 1008 out.write(' %s %s\n' % (self.type2def[type], name)) 1009 1010 return out.getvalue()
1011 1012
1013 - def define_argument_list(self, couplings=None):
1014 """define a list with the string of object required as incoming argument""" 1015 1016 conjugate = [2*(int(c[1:])-1) for c in self.tag if c[0] == 'C'] 1017 call_arg = [] 1018 #incoming argument of the routine 1019 call_arg.append( ('list_complex', 'P%s'% self.l_helas_id) ) 1020 1021 self.declaration.add(call_arg[0]) 1022 1023 for index,spin in enumerate(self.particles): 1024 if self.outgoing == index + 1: 1025 continue 1026 if self.l_helas_id == index + 1: 1027 continue 1028 call_arg.append(('complex','%s%d' % (spin, index +1))) 1029 self.declaration.add(('list_complex', call_arg[-1][-1])) 1030 1031 # couplings 1032 if couplings is None: 1033 detected_couplings = [name for type, name in self.declaration if name.startswith('COUP')] 1034 coup_sort = lambda x,y: int(x[4:])-int(y[4:]) 1035 detected_couplings.sort(coup_sort) 1036 if detected_couplings: 1037 couplings = detected_couplings 1038 else: 1039 couplings = ['COUP'] 1040 1041 for coup in couplings: 1042 call_arg.append(('complex', coup)) 1043 self.declaration.add(('complex',coup)) 1044 1045 if self.offshell: 1046 if aloha.complex_mass: 1047 call_arg.append(('complex','M%s' % self.outgoing)) 1048 self.declaration.add(('complex','M%s' % self.outgoing)) 1049 else: 1050 call_arg.append(('double','M%s' % self.outgoing)) 1051 self.declaration.add(('double','M%s' % self.outgoing)) 1052 call_arg.append(('double','W%s' % self.outgoing)) 1053 self.declaration.add(('double','W%s' % self.outgoing)) 1054 1055 self.call_arg = call_arg 1056 1057 return call_arg
1058
1059 - def get_momenta_txt(self):
1060 """Define the Header of the ortran file. This include 1061 - momentum conservation 1062 - definition of the impulsion""" 1063 1064 out = StringIO() 1065 1066 # Define all the required momenta 1067 p = [] # a list for keeping track how to write the momentum 1068 size = [] 1069 1070 signs = self.get_momentum_conservation_sign() 1071 1072 for i,type in enumerate(self.particles): 1073 if self.declaration.is_used('OM%s' % (i+1)): 1074 out.write(" OM{0} = {1}\n if (M{0}.ne.{1}) OM{0}={2}/M{0}**2\n".format( 1075 i+1, self.change_number_format(0), self.change_number_format(1))) 1076 1077 if i+1 == self.outgoing: 1078 out_type = 'P' 1079 continue 1080 elif i+1 == self.l_helas_id: 1081 p.append('%sP%s({%s})' % (signs[i],i+1,len(size))) 1082 size.append(0) 1083 continue 1084 elif self.offshell: 1085 p.append('%s%s%s({%s})' % (signs[i],type,i+1,len(size))) 1086 size.append(1) 1087 1088 if self.declaration.is_used('P%s' % (i+1)): 1089 self.get_one_momenta_def(i+1, out) 1090 1091 # define the resulting momenta 1092 if self.offshell: 1093 if aloha.loop_mode: 1094 size_p = 4 1095 else: 1096 size_p = 2 1097 for i in range(size_p): 1098 out.write(' P%s(%s) = %s\n' % (self.outgoing, i, 1099 ''.join(p).format(*[s+i for s in size]))) 1100 1101 1102 # Returning result 1103 return out.getvalue()
1104 1105
1106 - def get_loop_argument(self, key):
1107 """return the position for the argument in the HELAS convention""" 1108 1109 loop_momentum = key[:4] 1110 basis = key[4:] 1111 1112 loop_pos = sum([loop_momentum[i] * (i+1) for i in range(4)]) 1113 basis_pos = sum([basis[i] * (i+1) for i in range(len(basis))]) 1114 return (str(loop_pos), str(basis_pos))
1115 1116 1117 1118 1119 1120
1121 - def get_header_txt(self, name=None, couplings=None, **opt):
1122 """Define the Header of the fortran file. This include 1123 - function tag 1124 - definition of variable 1125 """ 1126 if name is None: 1127 name = self.name 1128 1129 out = StringIO() 1130 # define the type of function and argument 1131 1132 arguments = [arg for format, arg in self.define_argument_list(couplings)] 1133 self.declaration.add(('list_complex', 'P%s'% self.outgoing)) 1134 self.declaration.add(('list_complex', 'P%s'% self.l_helas_id)) 1135 self.declaration.add(('list_complex', 'coeff')) 1136 out.write('subroutine %(name)s(%(args)s, P%(out)s, COEFF)\n' % \ 1137 {'name': name, 'args': ', '.join(arguments), 1138 'out':self.outgoing}) 1139 1140 return out.getvalue()
1141
1142 -class ALOHAWriterForFortranLoopQP(QP, ALOHAWriterForFortranLoop):
1143 """routines for writing out Fortran""" 1144
1145 - def __init__(self, *arg):
1146 return ALOHAWriterForFortranLoop.__init__(self, *arg)
1147
1148 -def get_routine_name(name=None, outgoing=None, tag=None, abstract=None):
1149 """ build the name of the aloha function """ 1150 1151 assert (name and outgoing is not None) or abstract 1152 1153 if tag is None: 1154 tag = list(abstract.tag) 1155 else: 1156 tag=list(tag) 1157 1158 if name is None: 1159 prefix='' 1160 if 'MP' in tag: 1161 prefix = 'MP_' 1162 tag.remove('MP') 1163 if any(t.startswith('P') for t in tag): 1164 #put the propagator tag at the end 1165 propa = [t for t in tag if t.startswith('P')][0] 1166 tag.remove(propa) 1167 tag.append(propa) 1168 name = prefix + abstract.name + ''.join(tag) 1169 1170 if outgoing is None: 1171 outgoing = abstract.outgoing 1172 1173 return '%s_%s' % (name, outgoing)
1174
1175 -def combine_name(name, other_names, outgoing, tag=None, unknown_propa=False):
1176 """ build the name for combined aloha function """ 1177 1178 # Two possible scheme FFV1C1_2_X or FFV1__FFV2C1_X 1179 # If they are all in FFVX scheme then use the first 1180 p=re.compile('^(?P<type>[RFSVT]{2,})(?P<id>\d+)$') 1181 routine = '' 1182 if p.search(name): 1183 base, id = p.search(name).groups() 1184 routine = name 1185 for s in other_names: 1186 try: 1187 base2,id2 = p.search(s).groups() 1188 except Exception: 1189 routine = '' 1190 break # one matching not good -> other scheme 1191 if base != base2: 1192 routine = '' 1193 break # one matching not good -> other scheme 1194 else: 1195 routine += '_%s' % id2 1196 1197 if routine: 1198 if tag is not None: 1199 routine += ''.join(tag) 1200 if unknown_propa and outgoing: 1201 routine += '%(propa)s' 1202 if outgoing is not None: 1203 return routine +'_%s' % outgoing 1204 else: 1205 return routine 1206 1207 if tag is not None: 1208 addon = ''.join(tag) 1209 else: 1210 addon = '' 1211 if 'C' in name: 1212 short_name, addon = name.split('C',1) 1213 try: 1214 addon = 'C' + str(int(addon)) 1215 except Exception: 1216 addon = '' 1217 else: 1218 name = short_name 1219 if unknown_propa: 1220 addon += '%(propa)s' 1221 1222 if outgoing is not None: 1223 return '_'.join((name,) + tuple(other_names)) + addon + '_%s' % outgoing 1224 else: 1225 return '_'.join((name,) + tuple(other_names)) + addon
1226
1227 -class ALOHAWriterForCPP(WriteALOHA):
1228 """Routines for writing out helicity amplitudes as C++ .h and .cc files.""" 1229 1230 extension = '.c' 1231 writer = writers.CPPWriter 1232 1233 type2def = {} 1234 type2def['int'] = 'int' 1235 type2def['double'] = 'double ' 1236 type2def['complex'] = 'complex<double> ' 1237 1238 #variable overwritten by gpu 1239 realoperator = '.real()' 1240 imagoperator = '.imag()' 1241 ci_definition = ' complex<double> cI = complex<double>(0.,1.);\n' 1242 1243
1244 - def change_number_format(self, number):
1245 """Formating the number""" 1246 1247 def isinteger(x): 1248 try: 1249 return int(x) == x 1250 except TypeError: 1251 return False
1252 1253 if isinteger(number): 1254 out = '%s.' % (str(int(number))) 1255 elif isinstance(number, complex): 1256 if number.imag: 1257 if number.real: 1258 out = '(%s + %s*cI)' % (self.change_number_format(number.real), \ 1259 self.change_number_format(number.imag)) 1260 else: 1261 if number.imag == 1: 1262 out = 'cI' 1263 elif number.imag == -1: 1264 out = '-cI' 1265 else: 1266 out = '%s * cI' % self.change_number_format(number.imag) 1267 else: 1268 out = '%s' % (self.change_number_format(number.real)) 1269 else: 1270 tmp = Fraction(str(number)) 1271 tmp = tmp.limit_denominator(100) 1272 if not abs(tmp - number) / abs(tmp + number) < 1e-8: 1273 out = '%.9f' % (number) 1274 else: 1275 out = '%s./%s.' % (tmp.numerator, tmp.denominator) 1276 return out 1277 1278
1279 - def shift_indices(self, match):
1280 """shift the indices for non impulsion object""" 1281 if match.group('var').startswith('P'): 1282 shift = 0 1283 else: 1284 shift = self.momentum_size - 1 1285 return '%s[%s]' % (match.group('var'), int(match.group('num')) + shift)
1286 1287
1288 - def change_var_format(self, name):
1289 """Format the variable name to C++ format""" 1290 1291 if '_' in name: 1292 type = name.type 1293 decla = name.split('_',1)[0] 1294 self.declaration.add(('list_%s' % type, decla)) 1295 else: 1296 self.declaration.add((name.type, name.split('_',1)[0])) 1297 name = re.sub('(?P<var>\w*)_(?P<num>\d+)$', self.shift_indices , name) 1298 return name 1299
1300 - def get_fct_format(self, fct):
1301 """Put the function in the correct format""" 1302 if not hasattr(self, 'fct_format'): 1303 one = self.change_number_format(1) 1304 self.fct_format = {'csc' : '{0}/cos(%s)'.format(one), 1305 'sec': '{0}/sin(%s)'.format(one), 1306 'acsc': 'asin({0}/(%s))'.format(one), 1307 'asec': 'acos({0}/(%s))'.format(one), 1308 're': ' real(%s)', 1309 'im': 'imag(%s)', 1310 'cmath.sqrt':'sqrt(%s)', 1311 'sqrt': 'sqrt(%s)', 1312 'complexconjugate': 'conj(dcmplx(%s))', 1313 '/' : '{0}/%s'.format(one), 1314 'abs': 'abs(%s)' 1315 } 1316 1317 if fct in self.fct_format: 1318 return self.fct_format[fct] 1319 else: 1320 self.declaration.add(('fct', fct)) 1321 return '{0}(%s)'.format(fct)
1322 1323 1324 1325
1326 - def get_header_txt(self, name=None, couplings=None,mode=''):
1327 """Define the Header of the fortran file. This include 1328 - function tag 1329 - definition of variable 1330 """ 1331 if name is None: 1332 name = self.name 1333 1334 if mode=='': 1335 mode = self.mode 1336 1337 1338 1339 out = StringIO() 1340 # define the type of function and argument 1341 if not 'no_include' in mode: 1342 out.write('#include \"%s.h\"\n\n' % self.name) 1343 args = [] 1344 for format, argname in self.define_argument_list(couplings): 1345 if format.startswith('list'): 1346 type = self.type2def[format[5:]] 1347 list_arg = '[]' 1348 else: 1349 type = self.type2def[format] 1350 list_arg = '' 1351 args.append('%s%s%s'% (type, argname, list_arg)) 1352 1353 if not self.offshell: 1354 output = 'complex<double> & vertex' 1355 #self.declaration.add(('complex','vertex')) 1356 else: 1357 output = 'complex<double> %(spin)s%(id)d[]' % { 1358 'spin': self.particles[self.outgoing -1], 1359 'id': self.outgoing} 1360 self.declaration.add(('list_complex', output)) 1361 1362 out.write('void %(name)s(%(args)s,%(output)s)' % \ 1363 {'output':output, 'name': name, 'args': ', '.join(args)}) 1364 if 'is_h' in mode: 1365 out.write(';\n') 1366 else: 1367 out.write('\n{\n') 1368 1369 return out.getvalue()
1370
1371 - def get_declaration_txt(self):
1372 """ Prototype for how to write the declaration of variable 1373 Include the symmetry line (entry FFV_2) 1374 """ 1375 1376 out = StringIO() 1377 argument_var = [name for type,name in self.call_arg] 1378 # define the complex number CI = 0+1j 1379 out.write(self.ci_definition) 1380 1381 for type, name in self.declaration: 1382 if type.startswith('list'): 1383 type = type[5:] 1384 if name.startswith('P'): 1385 size = 4 1386 elif not 'tmp' in name: 1387 continue 1388 #should be define in the header 1389 elif name[0] in ['F','V']: 1390 if aloha.loop_mode: 1391 size = 8 1392 else: 1393 size = 6 1394 elif name[0] == 'S': 1395 if aloha.loop_mode: 1396 size = 5 1397 else: 1398 size = 3 1399 elif name[0] in ['R','T']: 1400 if aloha.loop_mode: 1401 size = 20 1402 else: 1403 size = 18 1404 1405 out.write(' %s %s[%s];\n' % (self.type2def[type], name, size)) 1406 elif (type, name) not in self.call_arg: 1407 out.write(' %s %s;\n' % (self.type2def[type], name)) 1408 1409 return out.getvalue()
1410
1411 - def get_foot_txt(self):
1412 """Prototype for language specific footer""" 1413 return '}\n'
1414
1415 - def get_momenta_txt(self):
1416 """Define the Header of the fortran file. This include 1417 - momentum conservation 1418 - definition of the impulsion""" 1419 1420 out = StringIO() 1421 1422 # Define all the required momenta 1423 p = [] # a list for keeping track how to write the momentum 1424 1425 signs = self.get_momentum_conservation_sign() 1426 1427 for i,type in enumerate(self.particles): 1428 if self.declaration.is_used('OM%s' % (i+1)): 1429 out.write(" OM{0} = {1};\n if (M{0} != {1})\n OM{0}={2}/pow(M{0},2);\n".format( 1430 i+1, self.change_number_format(0), self.change_number_format(1))) 1431 1432 if i+1 == self.outgoing: 1433 out_type = type 1434 out_size = self.type_to_size[type] 1435 continue 1436 elif self.offshell: 1437 p.append('{0}{1}{2}[%(i)s]'.format(signs[i],type,i+1,type)) 1438 1439 if self.declaration.is_used('P%s' % (i+1)): 1440 self.get_one_momenta_def(i+1, out) 1441 1442 # define the resulting momenta 1443 if self.offshell: 1444 energy_pos = out_size -2 1445 type = self.particles[self.outgoing-1] 1446 if aloha.loop_mode: 1447 size_p = 4 1448 else: 1449 size_p = 2 1450 1451 for i in range(size_p): 1452 dict_energy = {'i':i} 1453 out.write(' %s%s[%s] = %s;\n' % (type,self.outgoing, i, 1454 ''.join(p) % dict_energy)) 1455 if self.declaration.is_used('P%s' % self.outgoing): 1456 self.get_one_momenta_def(self.outgoing, out) 1457 1458 1459 # Returning result 1460 return out.getvalue()
1461
1462 - def get_one_momenta_def(self, i, strfile):
1463 1464 type = self.particles[i-1] 1465 1466 if aloha.loop_mode: 1467 template ='P%(i)d[%(j)d] = %(sign)s%(type)s%(i)d[%(nb)d];\n' 1468 else: 1469 template ='P%(i)d[%(j)d] = %(sign)s%(type)s%(i)d[%(nb2)d]%(operator)s;\n' 1470 1471 nb2 = 0 1472 for j in range(4): 1473 if not aloha.loop_mode: 1474 nb = j 1475 if j == 0: 1476 assert not aloha.mp_precision 1477 operator = self.realoperator # not suppose to pass here in mp 1478 elif j == 1: 1479 nb2 += 1 1480 elif j == 2: 1481 assert not aloha.mp_precision 1482 operator = self.imagoperator # not suppose to pass here in mp 1483 elif j ==3: 1484 nb2 -= 1 1485 else: 1486 operator ='' 1487 nb = j 1488 nb2 = j 1489 strfile.write(template % {'j':j,'type': type, 'i': i, 1490 'nb': nb, 'nb2': nb2, 'operator':operator, 1491 'sign': self.get_P_sign(i)})
1492 1493
1494 - def define_expression(self):
1495 """Write the helicity amplitude in C++ format""" 1496 1497 out = StringIO() 1498 1499 if self.routine.contracted: 1500 for name,obj in self.routine.contracted.items(): 1501 out.write(' %s = %s;\n' % (name, self.write_obj(obj))) 1502 self.declaration.add(('complex', name)) 1503 1504 for name, (fct, objs) in self.routine.fct.items(): 1505 format = ' %s = %s;\n' % (name, self.get_fct_format(fct)) 1506 out.write(format % ','.join([self.write_obj(obj) for obj in objs])) 1507 1508 1509 1510 numerator = self.routine.expr 1511 if not 'Coup(1)' in self.routine.infostr: 1512 coup_name = 'COUP' 1513 else: 1514 coup_name = '%s' % self.change_number_format(1) 1515 if not self.offshell: 1516 if coup_name == 'COUP': 1517 out.write(' vertex = COUP*%s;\n' % self.write_obj(numerator.get_rep([0]))) 1518 else: 1519 out.write(' vertex = %s;\n' % self.write_obj(numerator.get_rep([0]))) 1520 else: 1521 OffShellParticle = '%s%d' % (self.particles[self.offshell-1],\ 1522 self.offshell) 1523 if 'L' not in self.tag: 1524 coeff = 'denom' 1525 if not aloha.complex_mass: 1526 if self.routine.denominator: 1527 out.write(' denom = %(COUP)s/(%(denom)s)\n' % {'COUP': coup_name,\ 1528 'denom':self.write_obj(self.routine.denominator)}) 1529 else: 1530 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' % \ 1531 {'i': self.outgoing, 'coup': coup_name}) 1532 else: 1533 if self.routine.denominator: 1534 raise Exception, 'modify denominator are not compatible with complex mass scheme' 1535 1536 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' % \ 1537 {'i': self.outgoing, 'coup': coup_name}) 1538 self.declaration.add(('complex','denom')) 1539 if aloha.loop_mode: 1540 ptype = 'list_complex' 1541 else: 1542 ptype = 'list_double' 1543 self.declaration.add((ptype,'P%s' % self.outgoing)) 1544 else: 1545 coeff = 'COUP' 1546 1547 for ind in numerator.listindices(): 1548 out.write(' %s[%d]= %s*%s;\n' % (self.outname, 1549 self.pass_to_HELAS(ind), coeff, 1550 self.write_obj(numerator.get_rep(ind)))) 1551 return out.getvalue()
1552 1553 remove_double = re.compile('complex<double> (?P<name>[\w]+)\[\]')
1554 - def define_symmetry(self, new_nb, couplings=None):
1555 """Write the call for symmetric routines""" 1556 number = self.offshell 1557 arguments = [name for format, name in self.define_argument_list()] 1558 new_name = self.name.rsplit('_')[0] + '_%s' % new_nb 1559 output = '%(spin)s%(id)d' % { 1560 'spin': self.particles[self.offshell -1], 1561 'id': self.outgoing} 1562 return '%s\n %s(%s,%s);\n}' % \ 1563 (self.get_header_txt(new_name, couplings, mode='no_include'), 1564 self.name, ','.join(arguments), output)
1565
1566 - def get_h_text(self,couplings=None):
1567 """Return the full contents of the .h file""" 1568 1569 h_string = StringIO() 1570 if not self.mode == 'no_include': 1571 h_string.write('#ifndef '+ self.name + '_guard\n') 1572 h_string.write('#define ' + self.name + '_guard\n') 1573 h_string.write('#include <complex>\n') 1574 h_string.write('using namespace std;\n\n') 1575 1576 h_header = self.get_header_txt(mode='no_include__is_h', couplings=couplings) 1577 h_string.write(h_header) 1578 1579 for elem in self.routine.symmetries: 1580 symmetryhead = h_header.replace( \ 1581 self.name,self.name[0:-1]+'%s' %(elem)) 1582 h_string.write(symmetryhead) 1583 1584 if not self.mode == 'no_include': 1585 h_string.write('#endif\n\n') 1586 1587 return h_string.getvalue()
1588 1589
1590 - def write_combined_cc(self, lor_names, offshell=None, sym=True, mode=''):
1591 "Return the content of the .cc file linked to multiple lorentz call." 1592 1593 # Set some usefull command 1594 if offshell is None: 1595 offshell = self.offshell 1596 1597 name = combine_name(self.routine.name, lor_names, offshell, self.tag) 1598 self.name = name 1599 # write head - momenta - body - foot 1600 text = StringIO() 1601 routine = StringIO() 1602 data = {} # for the formating of the line 1603 1604 # write header 1605 new_couplings = ['COUP%s' % (i+1) for i in range(len(lor_names)+1)] 1606 text.write(self.get_header_txt(name=name, couplings=new_couplings, mode=mode)) 1607 1608 # Define which part of the routine should be called 1609 data['addon'] = ''.join(self.tag) + '_%s' % self.offshell 1610 1611 # how to call the routine 1612 argument = [name for format, name in self.define_argument_list(new_couplings)] 1613 index= argument.index('COUP1') 1614 data['before_coup'] = ','.join(argument[:index]) 1615 data['after_coup'] = ','.join(argument[index+len(lor_names)+1:]) 1616 if data['after_coup']: 1617 data['after_coup'] = ',' + data['after_coup'] 1618 1619 lor_list = (self.routine.name,) + lor_names 1620 line = " %(name)s%(addon)s(%(before_coup)s,%(coup)s%(after_coup)s,%(out)s);\n" 1621 main = '%(spin)s%(id)d' % {'spin': self.particles[self.offshell -1], 1622 'id': self.outgoing} 1623 for i, name in enumerate(lor_list): 1624 data['name'] = name 1625 data['coup'] = 'COUP%d' % (i+1) 1626 if i == 0: 1627 if not offshell: 1628 data['out'] = 'vertex' 1629 else: 1630 data['out'] = main 1631 elif i==1: 1632 if self.offshell: 1633 type = self.particles[self.offshell-1] 1634 self.declaration.add(('list_complex','%stmp' % type)) 1635 else: 1636 type = '' 1637 self.declaration.add(('complex','%stmp' % type)) 1638 data['out'] = '%stmp' % type 1639 routine.write(line % data) 1640 if i: 1641 if not offshell: 1642 routine.write( ' vertex = vertex + tmp;\n') 1643 else: 1644 size = self.type_to_size[self.particles[offshell -1]] -2 1645 routine.write(""" i= %s;\nwhile (i < %s)\n{\n""" % (self.momentum_size, self.momentum_size+size)) 1646 routine.write(" %(main)s[i] = %(main)s[i] + %(tmp)s[i];\n i++;\n" %\ 1647 {'main': main, 'tmp': data['out']}) 1648 routine.write('}\n') 1649 self.declaration.add(('int','i')) 1650 self.declaration.discard(('complex','COUP')) 1651 for name in aloha_lib.KERNEL.reduced_expr2: 1652 self.declaration.discard(('complex', name)) 1653 1654 #clean pointless declaration 1655 #self.declaration.discard 1656 1657 text.write(self.get_declaration_txt()) 1658 text.write(routine.getvalue()) 1659 text.write(self.get_foot_txt()) 1660 1661 text = text.getvalue() 1662 1663 return text
1664 1665
1666 - def write(self, **opt):
1667 """Write the .h and .cc files""" 1668 1669 cc_text = WriteALOHA.write(self, **opt) 1670 h_text = self.get_h_text() 1671 1672 # write in two file 1673 if self.out_path: 1674 writer_h = writers.CPPWriter(self.out_path[:-len(self.extension)] + ".h") 1675 commentstring = 'This File is Automatically generated by ALOHA \n' 1676 commentstring += 'The process calculated in this file is: \n' 1677 commentstring += self.routine.infostr + '\n' 1678 writer_h.write_comments(commentstring) 1679 writer_h.writelines(h_text) 1680 1681 return h_text, cc_text
1682 1683 1684
1685 - def write_combined(self, lor_names, mode='', offshell=None, **opt):
1686 """Write the .h and .cc files associated to the combined file""" 1687 1688 # Set some usefull command 1689 if offshell is None: 1690 sym = 1 1691 offshell = self.offshell 1692 else: 1693 sym = None 1694 1695 if mode == 'self': 1696 # added to another file 1697 self.mode = 'no_include' 1698 1699 #name = combine_name(self.name, lor_names, offshell, self.tag) 1700 1701 #h_text = self.write_combined_h(lor_names, offshell, **opt) 1702 cc_text, h_text = StringIO() , StringIO() 1703 cc_text.write(self.write_combined_cc(lor_names, offshell, mode=mode,**opt)) 1704 couplings = ['COUP%d' % (i+1) for i in range(len(lor_names)+1)] 1705 1706 if mode == 'self': 1707 self.mode = 'self' 1708 h_text.write(self.get_h_text(couplings=couplings)) 1709 1710 #ADD SYMETRY 1711 if sym: 1712 for elem in self.routine.symmetries: 1713 self.mode = 'no_include' 1714 cc_text.write( self.write_combined_cc(lor_names, elem)) 1715 1716 1717 if self.out_path: 1718 # Prepare a specific file 1719 path = os.path.join(os.path.dirname(self.out_path), self.name) 1720 commentstring = 'This File is Automatically generated by ALOHA \n' 1721 1722 writer_h = writers.CPPWriter(path + ".h") 1723 writer_h.write_comments(commentstring) 1724 writer_h.writelines(h_text) 1725 1726 writer_cc = writers.CPPWriter(path + ".cc") 1727 writer_cc.write_comments(commentstring) 1728 writer_cc.writelines(cc_text) 1729 1730 return h_text.getvalue(), cc_text.getvalue()
1731
1732 1733 -class ALOHAWriterForGPU(ALOHAWriterForCPP):
1734 1735 extension = '.cu' 1736 realoperator = '.re' 1737 imagoperator = '.im' 1738 ci_definition = 'complex<double> cI = mkcmplx(0., 1.);\n' 1739
1740 - def get_header_txt(self, name=None, couplings=None, mode=''):
1741 """Define the Header of the fortran file. This include 1742 - function tag 1743 - definition of variable 1744 """ 1745 text = StringIO() 1746 if not 'is_h' in mode: 1747 text.write('__device__=__forceinclude__\n') 1748 text.write(ALOHAWriterForCPP.get_header_txt(self, name, couplings, mode)) 1749 return text.getvalue()
1750
1751 - def get_h_text(self,couplings=None):
1752 """Return the full contents of the .h file""" 1753 1754 h_string = StringIO() 1755 if not self.mode == 'no_include': 1756 h_string.write('#ifndef '+ self.name + '_guard\n') 1757 h_string.write('#define ' + self.name + '_guard\n') 1758 h_string.write('#include "cmplx.h"\n') 1759 h_string.write('using namespace std;\n\n') 1760 1761 h_header = self.get_header_txt(mode='no_include__is_h', couplings=couplings) 1762 h_string.write(h_header) 1763 1764 for elem in self.routine.symmetries: 1765 symmetryhead = h_header.replace( \ 1766 self.name,self.name[0:-1]+'%s' %(elem)) 1767 h_string.write(symmetryhead) 1768 1769 if not self.mode == 'no_include': 1770 h_string.write('#endif\n\n') 1771 1772 return h_string.getvalue()
1773
1774 1775 -class ALOHAWriterForPython(WriteALOHA):
1776 """ A class for returning a file/a string for python evaluation """ 1777 1778 extension = '.py' 1779 writer = writers.PythonWriter 1780 1781 @staticmethod
1782 - def change_number_format(obj, pure_complex=''):
1783 change_number_format = ALOHAWriterForPython.change_number_format 1784 if obj.real == 0 and obj.imag: 1785 if int(obj.imag) == obj.imag: 1786 return '%ij' % obj.imag 1787 else: 1788 return change_number_format(obj.imag, pure_complex='j') 1789 elif obj.imag != 0: 1790 return '(%s+%s)' % (change_number_format(obj.real), 1791 change_number_format(obj.imag, pure_complex='j')) 1792 elif obj.imag == 0: 1793 if int(obj.real) == obj: 1794 return '%i%s' % (obj.real,pure_complex) 1795 obj = obj.real 1796 tmp = Fraction(str(obj)) 1797 tmp = tmp.limit_denominator(100) 1798 if not abs(tmp - obj) / abs(tmp + obj) < 1e-8: 1799 out = str(obj) 1800 elif tmp.denominator != 1: 1801 out = '%i%s/%i' % (tmp.numerator, pure_complex, tmp.denominator) 1802 else: 1803 out = '%i%s' % (tmp.numerator, pure_complex) 1804 return out
1805 1806
1807 - def shift_indices(self, match):
1808 """shift the indices for non impulsion object""" 1809 if match.group('var').startswith('P'): 1810 shift = 0 1811 else: 1812 shift = -1 + self.momentum_size 1813 1814 return '%s[%s]' % (match.group('var'), int(match.group('num')) + shift)
1815
1816 - def change_var_format(self, name):
1817 """Formatting the variable name to Python format 1818 start to count at zero. 1819 No neeed to define the variable in python -> no need to keep track of 1820 the various variable 1821 """ 1822 1823 if '_' not in name: 1824 self.declaration.add((name.type, name)) 1825 else: 1826 self.declaration.add(('', name.split('_',1)[0])) 1827 name = re.sub('(?P<var>\w*)_(?P<num>\d+)$', self.shift_indices , name) 1828 1829 return name
1830
1831 - def define_expression(self):
1832 """Define the functions in a 100% way """ 1833 1834 out = StringIO() 1835 1836 if self.routine.contracted: 1837 for name,obj in self.routine.contracted.items(): 1838 out.write(' %s = %s\n' % (name, self.write_obj(obj))) 1839 1840 numerator = self.routine.expr 1841 if not 'Coup(1)' in self.routine.infostr: 1842 coup_name = 'COUP' 1843 else: 1844 coup_name = '%s' % self.change_number_format(1) 1845 1846 if not self.offshell: 1847 if coup_name == 'COUP': 1848 out.write(' vertex = COUP*%s\n' % self.write_obj(numerator.get_rep([0]))) 1849 else: 1850 out.write(' vertex = %s\n' % self.write_obj(numerator.get_rep([0]))) 1851 else: 1852 OffShellParticle = '%s%d' % (self.particles[self.offshell-1],\ 1853 self.offshell) 1854 1855 if not 'L' in self.tag: 1856 coeff = 'denom' 1857 if not aloha.complex_mass: 1858 if self.routine.denominator: 1859 out.write(' denom = %(COUP)s/(%(denom)s)\n' % {'COUP': coup_name,\ 1860 'denom':self.write_obj(self.routine.denominator)}) 1861 else: 1862 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' % 1863 {'i': self.outgoing,'coup':coup_name}) 1864 else: 1865 if self.routine.denominator: 1866 raise Exception, 'modify denominator are not compatible with complex mass scheme' 1867 1868 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' % 1869 {'i': self.outgoing,'coup':coup_name}) 1870 else: 1871 coeff = 'COUP' 1872 1873 for ind in numerator.listindices(): 1874 out.write(' %s[%d]= %s*%s\n' % (self.outname, 1875 self.pass_to_HELAS(ind), coeff, 1876 self.write_obj(numerator.get_rep(ind)))) 1877 return out.getvalue()
1878
1879 - def get_foot_txt(self):
1880 if not self.offshell: 1881 return ' return vertex\n\n' 1882 else: 1883 return ' return %s\n\n' % (self.outname)
1884 1885
1886 - def get_header_txt(self, name=None, couplings=None, mode=''):
1887 """Define the Header of the fortran file. This include 1888 - function tag 1889 - definition of variable 1890 """ 1891 if name is None: 1892 name = self.name 1893 1894 out = StringIO() 1895 1896 if self.mode == 'mg5': 1897 out.write('import aloha.template_files.wavefunctions as wavefunctions\n') 1898 else: 1899 out.write('import wavefunctions\n') 1900 1901 1902 # define the type of function and argument 1903 1904 arguments = [arg for format, arg in self.define_argument_list(couplings)] 1905 out.write('def %(name)s(%(args)s):\n' % \ 1906 {'name': name, 'args': ','.join(arguments)}) 1907 1908 return out.getvalue()
1909
1910 - def get_momenta_txt(self):
1911 """Define the Header of the fortran file. This include 1912 - momentum conservation 1913 - definition of the impulsion""" 1914 1915 out = StringIO() 1916 1917 # Define all the required momenta 1918 p = [] # a list for keeping track how to write the momentum 1919 1920 signs = self.get_momentum_conservation_sign() 1921 1922 for i,type in enumerate(self.particles): 1923 if self.declaration.is_used('OM%s' % (i+1)): 1924 out.write(" OM{0} = 0.0\n if (M{0}): OM{0}=1.0/M{0}**2\n".format( (i+1) )) 1925 if i+1 == self.outgoing: 1926 out_type = type 1927 out_size = self.type_to_size[type] 1928 continue 1929 elif self.offshell: 1930 p.append('{0}{1}{2}[%(i)s]'.format(signs[i],type,i+1)) 1931 1932 if self.declaration.is_used('P%s' % (i+1)): 1933 self.get_one_momenta_def(i+1, out) 1934 1935 # define the resulting momenta 1936 if self.offshell: 1937 type = self.particles[self.outgoing-1] 1938 out.write(' %s%s = wavefunctions.WaveFunction(size=%s)\n' % (type, self.outgoing, out_size)) 1939 if aloha.loop_mode: 1940 size_p = 4 1941 else: 1942 size_p = 2 1943 for i in range(size_p): 1944 dict_energy = {'i':i} 1945 1946 out.write(' %s%s[%s] = %s\n' % (type,self.outgoing, i, 1947 ''.join(p) % dict_energy)) 1948 1949 self.get_one_momenta_def(self.outgoing, out) 1950 1951 1952 # Returning result 1953 return out.getvalue()
1954
1955 - def get_one_momenta_def(self, i, strfile):
1956 """return the string defining the momentum""" 1957 1958 type = self.particles[i-1] 1959 1960 main = ' P%d = [' % i 1961 if aloha.loop_mode: 1962 template ='%(sign)s%(type)s%(i)d[%(nb)d]' 1963 else: 1964 template ='%(sign)scomplex(%(type)s%(i)d[%(nb2)d])%(operator)s' 1965 1966 nb2 = 0 1967 strfile.write(main) 1968 data = [] 1969 for j in range(4): 1970 if not aloha.loop_mode: 1971 nb = j 1972 if j == 0: 1973 assert not aloha.mp_precision 1974 operator = '.real' # not suppose to pass here in mp 1975 elif j == 1: 1976 nb2 += 1 1977 elif j == 2: 1978 assert not aloha.mp_precision 1979 operator = '.imag' # not suppose to pass here in mp 1980 elif j ==3: 1981 nb2 -= 1 1982 else: 1983 operator ='' 1984 nb = j 1985 nb2 = j 1986 data.append(template % {'j':j,'type': type, 'i': i, 1987 'nb': nb, 'nb2': nb2, 'operator':operator, 1988 'sign': self.get_P_sign(i)}) 1989 1990 strfile.write(', '.join(data)) 1991 strfile.write(']\n')
1992 1993
1994 - def define_symmetry(self, new_nb, couplings=None):
1995 number = self.offshell 1996 arguments = [name for format, name in self.define_argument_list()] 1997 new_name = self.name.rsplit('_')[0] + '_%s' % new_nb 1998 return '%s\n return %s(%s)' % \ 1999 (self.get_header_txt(new_name, couplings), self.name, ','.join(arguments))
2000
2001 - def write_combined(self, lor_names, mode='self', offshell=None):
2002 """Write routine for combine ALOHA call (more than one coupling)""" 2003 2004 # Set some usefull command 2005 if offshell is None: 2006 sym = 1 2007 offshell = self.offshell 2008 else: 2009 sym = None 2010 name = combine_name(self.routine.name, lor_names, offshell, self.tag) 2011 # write head - momenta - body - foot 2012 text = StringIO() 2013 data = {} # for the formating of the line 2014 2015 # write header 2016 new_couplings = ['COUP%s' % (i+1) for i in range(len(lor_names)+1)] 2017 text.write(self.get_header_txt(name=name, couplings=new_couplings)) 2018 2019 # Define which part of the routine should be called 2020 data['addon'] = ''.join(self.tag) + '_%s' % self.offshell 2021 2022 # how to call the routine 2023 argument = [name for format, name in self.define_argument_list(new_couplings)] 2024 index= argument.index('COUP1') 2025 data['before_coup'] = ','.join(argument[:index]) 2026 data['after_coup'] = ','.join(argument[index+len(lor_names)+1:]) 2027 if data['after_coup']: 2028 data['after_coup'] = ',' + data['after_coup'] 2029 2030 lor_list = (self.routine.name,) + lor_names 2031 line = " %(out)s = %(name)s%(addon)s(%(before_coup)s,%(coup)s%(after_coup)s)\n" 2032 main = '%(spin)s%(id)d' % {'spin': self.particles[self.offshell -1], 2033 'id': self.outgoing} 2034 for i, name in enumerate(lor_list): 2035 data['name'] = name 2036 data['coup'] = 'COUP%d' % (i+1) 2037 if i == 0: 2038 if not offshell: 2039 data['out'] = 'vertex' 2040 else: 2041 data['out'] = main 2042 elif i==1: 2043 data['out'] = 'tmp' 2044 text.write(line % data) 2045 if i: 2046 if not offshell: 2047 text.write( ' vertex += tmp\n') 2048 else: 2049 size = self.type_to_size[self.particles[offshell -1]] -2 2050 text.write(" for i in range(%s,%s):\n" % (self.momentum_size, self.momentum_size+size)) 2051 text.write(" %(main)s[i] += tmp[i]\n" %{'main': main}) 2052 2053 text.write(self.get_foot_txt()) 2054 2055 #ADD SYMETRY 2056 if sym: 2057 for elem in self.routine.symmetries: 2058 text.write(self.write_combined(lor_names, mode, elem)) 2059 2060 text = text.getvalue() 2061 if self.out_path: 2062 writer = self.writer(self.out_path) 2063 commentstring = 'This File is Automatically generated by ALOHA \n' 2064 commentstring += 'The process calculated in this file is: \n' 2065 commentstring += self.routine.infostr + '\n' 2066 writer.write_comments(commentstring) 2067 writer.writelines(text) 2068 2069 2070 return text
2071
2072 2073 -class Declaration_list(set):
2074
2075 - def is_used(self, var):
2076 if hasattr(self, 'var_name'): 2077 return var in self.var_name 2078 self.var_name = [name for type,name in self] 2079 return var in self.var_name
2080
2081 - def add(self,obj):
2082 if __debug__: 2083 type, name = obj 2084 samename = [t for t,n in self if n ==name] 2085 for type2 in samename: 2086 assert type2 == type, '%s is defined with two different type "%s" and "%s"' % \ 2087 (name, type2, type) 2088 2089 set.add(self,obj)
2090
2091 2092 -class WriterFactory(object):
2093
2094 - def __new__(cls, data, language, outputdir, tags):
2095 language = language.lower() 2096 if isinstance(data.expr, aloha_lib.SplitCoefficient): 2097 assert language == 'fortran' 2098 if 'MP' in tags: 2099 return ALOHAWriterForFortranLoopQP(data, outputdir) 2100 else: 2101 return ALOHAWriterForFortranLoop(data, outputdir) 2102 if language == 'fortran': 2103 if 'MP' in tags: 2104 return ALOHAWriterForFortranQP(data, outputdir) 2105 else: 2106 return ALOHAWriterForFortran(data, outputdir) 2107 elif language == 'python': 2108 return ALOHAWriterForPython(data, outputdir) 2109 elif language == 'cpp': 2110 return ALOHAWriterForCPP(data, outputdir) 2111 elif language == 'gpu': 2112 return ALOHAWriterForGPU(data, outputdir) 2113 else: 2114 raise Exception, 'Unknown output format'
2115 2116 2117 2118 #unknow_fct_template = """ 2119 #cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc 2120 # double complex %(fct_name)s(%(args)s) 2121 # implicit none 2122 #c Include Model parameter / coupling 2123 # include \"../MODEL/input.inc\" 2124 # include \"../MODEL/coupl.inc\" 2125 #c Defintion of the arguments 2126 #%(definitions)s 2127 # 2128 #c enter HERE the code corresponding to your function. 2129 #c The output value should be put to the %(fct_name)s variable. 2130 # 2131 # 2132 # return 2133 # end 2134 #cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc 2135 # 2136 #""" 2137 # 2138 #def write_template_fct(fct_name, nb_args, output_dir): 2139 # """create a template for function not recognized by ALOHA""" 2140 # 2141 # dico = {'fct_name' : fct_name, 2142 # 'args': ','.join(['S%i' %(i+1) for i in range(nb_args)]), 2143 # 'definitions': '\n'.join([' double complex S%i' %(i+1) for i in range(nb_args)])} 2144 # 2145 # ff = open(pjoin(output_dir, 'additional_aloha_function.f'), 'a') 2146 # ff.write(unknow_fct_template % dico) 2147 # ff.close() 2148