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