1
2
3
4
5
6
7
8
9
10
11
12
13
14
15 from __future__ import division
16 import cmath
17 import copy
18 import cPickle
19 import glob
20 import logging
21 import numbers
22 import os
23 import re
24 import shutil
25 import sys
26 import time
27
28 root_path = os.path.split(os.path.dirname(os.path.realpath( __file__ )))[0]
29 sys.path.append(root_path)
30 from aloha.aloha_object import *
31 import aloha.aloha_writers as aloha_writers
32 import aloha.aloha_lib as aloha_lib
33 try:
34 import madgraph.iolibs.files as files
35 except:
36 import aloha.files as files
37
38 aloha_path = os.path.dirname(os.path.realpath(__file__))
39 logger = logging.getLogger('ALOHA')
40
41 _conjugate_gap = 50
42 _spin2_mult = 1000
45
47 """ store the result of the computation of Helicity Routine
48 this is use for storing and passing to writer """
49
50 - def __init__(self, expr, outgoing, spins, name, infostr):
51 """ store the information """
52
53 self.spins = spins
54 self.expr = expr
55 self.name = name
56 self.outgoing = outgoing
57 self.infostr = infostr
58 self.symmetries = []
59 self.combined = []
60
62 """ add an outgoing """
63
64 if not outgoing in self.symmetries:
65 self.symmetries.append(outgoing)
66
68 """add a combine rule """
69
70 if lor_list not in self.combined:
71 self.combined.append(lor_list)
72
73 - def write(self, output_dir, language='Fortran', mode='self', **opt):
74 """ write the content of the object """
75
76 writer = getattr(aloha_writers, 'ALOHAWriterFor%s' % language)(self, output_dir)
77 text = writer.write(mode=mode, **opt)
78 for grouped in self.combined:
79 if isinstance(text, tuple):
80 text = tuple([old.__add__(new) for old, new in zip(text,
81 writer.write_combined(grouped, mode=mode, **opt))])
82 else:
83 text += writer.write_combined(grouped, mode=mode, **opt)
84 return text
85
87 """ Launch the creation of the Helicity Routine"""
88
89 aloha_lib = None
90 counter = 0
91
93 """ An error class for ALOHA"""
94
96 """ initialize the run
97 lorentz: the lorentz information analyzed (UFO format)
98 language: define in which language we write the output
99 modes: 0 for all incoming particles
100 >0 defines the outgoing part (start to count at 1)
101 """
102
103 self.spins = lorentz.spins
104 self.name = lorentz.name
105 self.conjg = []
106 self.outgoing = None
107 self.lorentz_expr = lorentz.structure
108 self.routine_kernel = None
109 self.spin2_massless = False
110
111
117
119 """ return the full set of AbstractRoutineBuilder linked to fermion
120 clash"""
121
122 solution = []
123
124 for i, pair in enumerate(pair_list):
125 new_builder = self.define_conjugate_builder(pair)
126 solution.append(new_builder)
127 solution += new_builder.define_all_conjugate_builder(pair_list[i+1:])
128 return solution
129
131 """ return a AbstractRoutineBuilder for the conjugate operation.
132 If they are more than one pair of fermion. Then use pair to claim which
133 one is conjugated"""
134
135 new_builder = copy.copy(self)
136 new_builder.conjg = self.conjg[:]
137 try:
138 for index in pairs:
139 new_builder.apply_conjugation(index)
140 except TypeError:
141 new_builder.apply_conjugation(pairs)
142 return new_builder
143
145 """ apply conjugation on self object"""
146
147
148 old_id = 2 * pair - 1
149 new_id = _conjugate_gap + old_id
150
151 if self.routine_kernel is None:
152 self.kernel_tag = set()
153 self.routine_kernel = eval(self.lorentz_expr)
154
155 self.routine_kernel = \
156 C(new_id, old_id + 1) * self.routine_kernel * C(new_id + 1, old_id)
157 self.name += 'C'
158
159 if pair:
160 self.name += str(pair)
161 self.conjg.append(pair)
162
163
164
166 """ define a simple output for this AbstractRoutine """
167
168 infostr = str(self.lorentz_expr)
169 return AbstractRoutine(self.expr, self.outgoing, self.spins, self.name, \
170 infostr)
171
173 """change the sign of P for outcoming fermion in order to
174 correct the mismatch convention between HELAS and FR"""
175
176 flip_sign = []
177 for i in range(1,len(self.spins),2):
178 if self.spins[i] == 2:
179 flip_sign.append(str(i))
180
181 if not flip_sign:
182 return self.lorentz_expr
183 momentum_pattern = re.compile(r'\bP\(([\+\-\d]+),(%s)\)' % '|'.join(flip_sign))
184 lorentz_expr = momentum_pattern.sub(r'P(\1,\2, -1)', self.lorentz_expr)
185 return lorentz_expr
186
187
189 """compute the abstract routine associate to this mode """
190
191
192 aloha_lib.USE_TAG=set()
193
194 nb_spinor = 0
195 if not self.routine_kernel:
196 AbstractRoutineBuilder.counter += 1
197 logger.info('aloha creates %s routines' % self.name)
198 try:
199 lorentz = self.change_sign_for_outcoming_fermion()
200 lorentz = eval(lorentz)
201 except NameError:
202 logger.error('unknow type in Lorentz Evaluation')
203 raise ALOHAERROR, 'unknow type in Lorentz Evaluation'
204 else:
205 self.routine_kernel = lorentz
206 self.kernel_tag = set(aloha_lib.USE_TAG)
207 else:
208 lorentz = self.routine_kernel
209 aloha_lib.USE_TAG = set(self.kernel_tag)
210 for (i, spin ) in enumerate(self.spins):
211 id = i + 1
212
213 if id == self.outgoing:
214 if spin == 1:
215 lorentz *= complex(0,1)
216 elif spin == 2:
217
218 if (id+1) // 2 in self.conjg:
219 id += _conjugate_gap
220 nb_spinor += 1
221 if nb_spinor %2:
222 lorentz *= SpinorPropagator(id, 'I2', self.outgoing)
223 else:
224 lorentz *= SpinorPropagator('I2', id, self.outgoing)
225 elif spin == 3 :
226 lorentz *= VectorPropagator(id, 'I2', id)
227 elif spin == 5 :
228 lorentz *= 1
229
230
231
232
233
234
235 else:
236 raise self.AbstractALOHAError(
237 'The spin value %s is not supported yet' % spin)
238 else:
239
240 if spin == 1:
241 lorentz *= Scalar(id)
242 elif spin == 2:
243
244 if (id+1) // 2 in self.conjg:
245 spin_id = id + _conjugate_gap
246 else:
247 spin_id = id
248 nb_spinor += 1
249 lorentz *= Spinor(spin_id, id)
250 elif spin == 3:
251 lorentz *= Vector(id, id)
252 elif spin == 5:
253 lorentz *= Spin2(1 * _spin2_mult + id, 2 * _spin2_mult + id, id)
254 else:
255 raise self.AbstractALOHAError(
256 'The spin value %s is not supported yet' % spin)
257
258
259 if self.outgoing:
260 lorentz /= DenominatorPropagator(self.outgoing)
261
262
263 else:
264 lorentz *= complex(0,-1)
265
266
267 lorentz = lorentz.simplify()
268
269 lorentz = lorentz.expand()
270 if self.outgoing and self.spins[self.outgoing-1] == 5:
271 if not self.aloha_lib:
272 AbstractRoutineBuilder.load_library()
273 if self.spin2_massless:
274 lorentz *= self.aloha_lib[('Spin2PropMassless', id)]
275 else:
276 lorentz *= self.aloha_lib[('Spin2Prop', id)]
277 aloha_lib.USE_TAG.add('OM%d' % id)
278 aloha_lib.USE_TAG.add('P%d' % id)
279
280 lorentz = lorentz.simplify()
281 if factorize:
282 lorentz = lorentz.factorize()
283 lorentz.tag = set(aloha_lib.USE_TAG)
284 return lorentz
285
287 """Define the expression"""
288
289 self.expr = lorentz_expr
290
292 """Define the kernel at low level"""
293
294 if not lorentz:
295 logger.info('compute kernel %s' % self.counter)
296 AbstractRoutineBuilder.counter += 1
297 lorentz = eval(self.lorentz_expr)
298
299 if isinstance(lorentz, numbers.Number):
300 self.routine_kernel = lorentz
301 return lorentz
302 lorentz = lorentz.simplify()
303 lorentz = lorentz.expand()
304 lorentz = lorentz.simplify()
305
306 self.routine_kernel = lorentz
307 return lorentz
308
309
310 @staticmethod
312 """return the name of the """
313
314 name = '%s_%s' % (name, outgoing)
315 return name
316
317 @classmethod
327
330 """ A class to build and store the full set of Abstract ALOHA Routine"""
331
332 - def __init__(self, model_name, write_dir=None, format='Fortran'):
333 """ load the UFO model and init the dictionary """
334
335
336 model_name_pattern = re.compile("^(?P<name>.+)-(?P<rest>[\w\d_]+)$")
337 model_name_re = model_name_pattern.match(model_name)
338 if model_name_re:
339 name = model_name_re.group('name')
340 rest = model_name_re.group("rest")
341 if rest == 'full' or \
342 os.path.isfile(os.path.join(root_path, "models", name,
343 "restrict_%s.dat" % rest)):
344 model_name = model_name_re.group("name")
345
346
347 try:
348 python_pos = model_name
349 __import__(python_pos)
350 except:
351 python_pos = 'models.%s' % model_name
352 __import__(python_pos)
353 self.model = sys.modules[python_pos]
354
355 self.model_pos = os.path.dirname(self.model.__file__)
356
357
358 self.external_routines = []
359
360
361 dict.__init__(self)
362 self.symmetries = {}
363 self.multiple_lor = {}
364
365
366 self.massless_spin2 = self.has_massless_spin2()
367
368 if write_dir:
369 self.main(write_dir,format=format)
370
371 - def main(self, output_dir, format='Fortran'):
372 """ Compute if not already compute.
373 Write file in models/MY_MODEL/MY_FORMAT.
374 copy the file to output_dir
375 """
376 ext = {'Fortran':'f','Python':'py','CPP':'h'}
377
378
379
380 if not self.load():
381 self.compute_all()
382 logger.info(' %s aloha routine' % len(self))
383
384
385 if not output_dir:
386 output_dir = os.path.join(self.model_pos, format.lower())
387 logger.debug('aloha output dir is %s' % output_dir)
388 if not os.path.exists(output_dir):
389 os.mkdir(output_dir)
390
391
392 for (name, outgoing), abstract in self.items():
393 routine_name = AbstractRoutineBuilder.get_routine_name(name, outgoing)
394 if not glob.glob(os.path.join(output_dir, routine_name) + '.' + ext[format]):
395 abstract.write(output_dir, format)
396 else:
397 logger.info('File for %s already present, skip the writing of this file' % routine_name)
398
399
400 - def save(self, filepos=None):
401 """ save the current model in a pkl file """
402
403 logger.info('save the aloha abstract routine in a pickle file')
404 if not filepos:
405 filepos = os.path.join(self.model_pos,'aloha.pkl')
406
407 fsock = open(filepos, 'w')
408 cPickle.dump(dict(self), fsock)
409
410 - def load(self, filepos=None):
411 """ reload the pickle file """
412 return False
413 if not filepos:
414 filepos = os.path.join(self.model_pos,'aloha.pkl')
415 if os.path.exists(filepos):
416 fsock = open(filepos, 'r')
417 self.update(cPickle.load(fsock))
418 return True
419 else:
420 return False
421
422 - def get(self, lorentzname, outgoing):
423 """ return the AbstractRoutine with a given lorentz name, and for a given
424 outgoing particle """
425
426 try:
427 return self[(lorentzname, outgoing)]
428 except:
429 logger.warning('(%s, %s) is not a valid key' %
430 (lorentzname, outgoing) )
431 return None
432
433 - def set(self, lorentzname, outgoing, abstract_routine):
434 """ add in the dictionary """
435
436 self[(lorentzname, outgoing)] = abstract_routine
437
438 - def compute_all(self, save=True, wanted_lorentz = []):
439 """ define all the AbstractRoutine linked to a model """
440
441
442
443
444 self.look_for_symmetries()
445 conjugate_list = self.look_for_conjugate()
446 self.look_for_multiple_lorentz_interactions()
447
448 if not wanted_lorentz:
449 wanted_lorentz = [l.name for l in self.model.all_lorentz]
450 for lorentz in self.model.all_lorentz:
451 if not lorentz.name in wanted_lorentz:
452
453 continue
454
455 if -1 in lorentz.spins:
456
457 continue
458
459 if lorentz.structure == 'external':
460 self.external_routines.append(lorentz.name)
461 continue
462
463 builder = AbstractRoutineBuilder(lorentz)
464
465 if 5 in lorentz.spins and self.massless_spin2 is not None:
466 builder.spin2_massless = self.massless_spin2
467 self.compute_aloha(builder)
468
469 if lorentz.name in self.multiple_lor:
470 for m in self.multiple_lor[lorentz.name]:
471 for outgoing in range(len(lorentz.spins)+1):
472 try:
473 self[(lorentz.name, outgoing)].add_combine(m)
474 except:
475 pass
476
477
478 if lorentz.name in conjugate_list:
479 conjg_builder_list= builder.define_all_conjugate_builder(\
480 conjugate_list[lorentz.name])
481 for conjg_builder in conjg_builder_list:
482
483 assert conjg_builder_list.count(conjg_builder) == 1
484 self.compute_aloha(conjg_builder, lorentz.name)
485 if lorentz.name in self.multiple_lor:
486 for m in self.multiple_lor[lorentz.name]:
487 for outgoing in range(len(lorentz.spins)+1):
488 self[(conjg_builder.name, outgoing)].add_combine(m)
489
490
491
492 if save:
493 self.save()
494
495
497 """ create the requested ALOHA routine.
498 data should be a list of tuple (lorentz, conjugate, outgoing)
499 conjugate should be a tuple with the pair number to conjugate.
500 outgoing a tuple of the requested routines."""
501
502
503
504 self.look_for_symmetries()
505
506
507
508 request = {}
509 for list_l_name, conjugate, outgoing in data:
510 for l_name in list_l_name:
511 try:
512 request[l_name][conjugate].append(outgoing)
513 except:
514 try:
515 request[l_name][conjugate] = [outgoing]
516 except:
517 request[l_name] = {conjugate: [outgoing]}
518
519
520 for l_name in request:
521 lorentz = eval('self.model.lorentz.%s' % l_name)
522 if lorentz.structure == 'external':
523 if lorentz.name not in self.external_routines:
524 self.external_routines.append(lorentz.name)
525 continue
526
527 builder = AbstractRoutineBuilder(lorentz)
528
529 if 5 in lorentz.spins and self.massless_spin2 is not None:
530 builder.spin2_massless = self.massless_spin2
531
532 for conjg in request[l_name]:
533
534 routines = sorted(request[l_name][conjg])
535 if not conjg:
536
537 self.compute_aloha(builder, routines=routines)
538 else:
539
540 conjg_builder = builder.define_conjugate_builder(conjg)
541
542 self.compute_aloha(conjg_builder, symmetry=lorentz.name,
543 routines=routines)
544
545
546 for list_l_name, conjugate, outgoing in data:
547 if len(list_l_name) >1:
548 lorentzname = list_l_name[0]
549 for c in conjugate:
550 lorentzname += 'C%s' % c
551 self[(lorentzname, outgoing)].add_combine(list_l_name[1:])
552
554 """ define all the AbstractRoutine linked to a given lorentz structure
555 symmetry authorizes to use the symmetry of anoter lorentz structure.
556 routines to define only a subset of the routines."""
557
558 name = builder.name
559 if not symmetry:
560 symmetry = name
561 if not routines:
562 routines = range(len(builder.spins) + 1)
563
564
565 for outgoing in routines:
566 symmetric = self.has_symmetries(symmetry, outgoing, valid_output=routines)
567 if symmetric:
568 self.get(name, symmetric).add_symmetry(outgoing)
569 else:
570 wavefunction = builder.compute_routine(outgoing)
571
572 self.set(name, outgoing, wavefunction)
573
575 """define all the AbstractRoutine linked to a given lorentz structure
576 symmetry authorizes to use the symmetry of anoter lorentz structure.
577 routines to define only a subset of the routines.
578 Compare to compute_aloha, each routines are computed independently.
579 """
580
581 name = builder.name
582 if not routines:
583 routines = range(len(builder.spins) + 1 )
584
585 for outgoing in routines:
586 builder.routine_kernel = None
587 wavefunction = builder.compute_routine(outgoing)
588 self.set(name, outgoing, wavefunction)
589
590
591 - def write(self, output_dir, language):
592 """ write the full set of Helicity Routine in output_dir"""
593
594 for abstract_routine in self.values():
595 abstract_routine.write(output_dir, language)
596
597 for routine in self.external_routines:
598 self.locate_external(routine, language, output_dir)
599
600
601
603 """search a valid external file and copy it to output_dir directory"""
604
605 language_to_ext = {'Python': 'py',
606 'Fortran' : 'f',
607 'CPP': 'C'}
608 ext = language_to_ext[language]
609
610 if os.path.exists(os.path.join(self.model_pos, '%s.%s' % (name, ext))):
611 filepos = '%s/%s.%s' % (self.model_pos, name, ext)
612
613 elif os.path.exists(os.path.join(root_path, 'aloha', 'template_files',
614 '%s.%s' %(name, ext))):
615 filepos = '%s/aloha/template_files/%s.%s' % (root_path, name, ext)
616 else:
617 path1 = self.model_pos
618 path2 = os.path.join(root_path, 'aloha', 'template_files', )
619 raise ALOHAERROR, 'No external routine \"%s.%s\" in directories\n %s\n %s' % \
620 (name, ext, path1, path2)
621
622 if output_dir:
623 files.cp(filepos, output_dir)
624 return filepos
625
626
627
629 """Search some symmetries in the vertices.
630 We search if some identical particles are in a vertices in order
631 to avoid to compute symmetrical contributions"""
632
633 for vertex in self.model.all_vertices:
634 for i, part1 in enumerate(vertex.particles):
635 for j in range(i-1,-1,-1):
636 part2 = vertex.particles[j]
637 if part1.pdg_code == part2.pdg_code and part1.color == 1:
638 if part1.spin == 2 and (i % 2 != j % 2 ):
639 continue
640 for lorentz in vertex.lorentz:
641 if self.symmetries.has_key(lorentz.name):
642 if self.symmetries[lorentz.name].has_key(i+1):
643 self.symmetries[lorentz.name][i+1] = max(self.symmetries[lorentz.name][i+1], j+1)
644 else:
645 self.symmetries[lorentz.name][i+1] = j+1
646 else:
647 self.symmetries[lorentz.name] = {i+1:j+1}
648 break
649
651 """Search the interaction associate with more than one lorentz structure.
652 If those lorentz structure have the same order and the same color then
653 associate a multiple lorentz routines to ALOHA """
654
655 orders = {}
656 for coup in self.model.all_couplings:
657 orders[coup.name] = str(coup.order)
658
659 for vertex in self.model.all_vertices:
660 if len(vertex.lorentz) == 1:
661 continue
662
663 if -1 in vertex.lorentz[0].spins:
664 continue
665
666
667 combine = {}
668 for (id_col, id_lor), coup in vertex.couplings.items():
669 order = orders[coup.name]
670 key = (id_col, order)
671 if key in combine:
672 combine[key].append(id_lor)
673 else:
674 combine[key] = [id_lor]
675
676
677 for list_lor in combine.values():
678 if len(list_lor) == 1:
679 continue
680 list_lor.sort()
681 main = vertex.lorentz[list_lor[0]].name
682 if main not in self.multiple_lor:
683 self.multiple_lor[main] = []
684
685 info = tuple([vertex.lorentz[id].name for id in list_lor[1:]])
686 if info not in self.multiple_lor[main]:
687 self.multiple_lor[main].append(info)
688
690 """Search if the spin2 particles are massless or not"""
691
692 massless = None
693 for particle in self.model.all_particles:
694 if particle.spin == 5:
695 if massless is None:
696 massless = (particle.mass == 'Zero')
697 elif massless != (particle.mass == 'Zero'):
698 raise ALOHAERROR, 'All spin 2 should be massive or massless'
699 return massless
700
701 - def has_symmetries(self, l_name, outgoing, out=None, valid_output=None):
702 """ This returns out if no symmetries are available, otherwise it finds
703 the lowest equivalent outgoing by recursivally calling this function.
704 auth is a list of authorize output, if define"""
705
706 try:
707 equiv = self.symmetries[l_name][outgoing]
708 except:
709 return out
710 else:
711 if not valid_output or equiv in valid_output:
712 return self.has_symmetries(l_name, equiv, out=equiv,
713 valid_output=valid_output)
714 else:
715 return self.has_symmetries(l_name, equiv, out=out,
716 valid_output=valid_output)
717
719 """ create a list for the routine needing to be conjugate """
720
721
722 need = False
723 for particle in self.model.all_particles:
724 if particle.spin == 2 and particle.selfconjugate:
725 need = True
726 break
727
728 if not need:
729 for interaction in self.model.all_vertices:
730 fermions = [p for p in interaction.particles if p.spin == 2]
731 for i in range(0, len(fermions), 2):
732 if fermions[i].pdg_code * fermions[i+1].pdg_code > 0:
733
734 need = True
735 break
736
737
738 if not need:
739 return {}
740
741 conjugate_request = {}
742
743 for vertex in self.model.all_vertices:
744 for i in range(0, len(vertex.particles), 2):
745 part1 = vertex.particles[i]
746 if part1.spin !=2:
747
748 break
749
750 if part1.selfconjugate:
751 continue
752 part2 = vertex.particles[i + 1]
753 if part2.selfconjugate:
754 continue
755
756
757 for lorentz in vertex.lorentz:
758 try:
759 conjugate_request[lorentz.name].add(i//2+1)
760 except:
761 conjugate_request[lorentz.name] = set([i//2+1])
762
763 for elem in conjugate_request:
764 conjugate_request[elem] = list(conjugate_request[elem])
765
766 return conjugate_request
767
771 """find the list of Helicity routine in the directory and create a list
772 of those files (but with compile extension)"""
773
774 aloha_files = []
775
776
777 alohafile_pattern = re.compile(r'''_\d%s''' % file_ext)
778 for filename in os.listdir(aloha_dir):
779 if os.path.isfile(os.path.join(aloha_dir, filename)):
780 if alohafile_pattern.search(filename):
781 aloha_files.append(filename.replace(file_ext, comp_ext))
782
783 text="ALOHARoutine = "
784 text += ' '.join(aloha_files)
785 text +='\n'
786 file(os.path.join(aloha_dir, 'aloha_file.inc'), 'w').write(text)
787
791
792 def create(obj):
793 """ """
794 obj= obj.simplify()
795 obj = obj.expand()
796 obj = obj.simplify()
797 return obj
798
799
800 old_tag = set(aloha_lib.USE_TAG)
801
802 lib = {}
803 for i in range(1, 10):
804 logger.info('step %s/9' % i)
805
806
807
808
809
810
811
812
813
814 lib[('Spin2Prop',i)] = create( Spin2Propagator(_spin2_mult + i, \
815 2 * _spin2_mult + i,'I2','I3', i) )
816 lib[('Spin2PropMassless',i)] = create( Spin2masslessPropagator(
817 _spin2_mult + i, 2 * _spin2_mult + i,'I2','I3'))
818 logger.info('writing Spin2 lib')
819 fsock = open(os.path.join(aloha_path, 'ALOHALib.pkl'),'wb')
820 cPickle.dump(lib, fsock, -1)
821 aloha_lib.USE_TAG = old_tag
822 return lib
823
824 if '__main__' == __name__:
825 logging.basicConfig(level=0)
826
827 import profile
828
829
830 start = time.time()
835 - def write(alohagenerator):
837 alohagenerator = main()
838 logger.info('done in %s s' % (time.time()-start))
839 write(alohagenerator)
840
841
842 stop = time.time()
843 logger.info('done in %s s' % (stop-start))
844