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