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