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