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

Source Code for Module aloha.aloha_writers

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