1 from __future__ import division
2 from __future__ import absolute_import
3 from __future__ import print_function
4 import collections
5 import random
6 import re
7 import operator
8 import numbers
9 import math
10 import time
11 import os
12 import shutil
13 import sys
14 import six
15 from six.moves import filter
16 from six.moves import range
17 from six.moves import zip
18 from functools import reduce
19
20 pjoin = os.path.join
21
22 if '__main__' == __name__:
23 import sys
24 import os
25 sys.path.append('../../')
26 root = os.path.dirname(__file__)
27 if os.path.basename(root) == 'internal':
28 __package__ = "internal"
29 sys.path.append(os.path.dirname(root))
30 import internal
31 else:
32 __package__ = "madgraph.various"
33
34 try:
35 import madgraph
36 except ImportError:
37 from . import misc
38 from . import banner as banner_mod
39 else:
40 import madgraph.various.misc as misc
41 import madgraph.various.banner as banner_mod
42 import logging
43 import gzip
44
45
46 try:
47 import madgraph.various.hepmc_parser as hepmc_parser
48 except Exception as error:
49 hepmc_parser = False
50 misc.sprint("No hepmc reader since", error)
51 pass
52
53
54 logger = logging.getLogger("madgraph.lhe_parser")
55
56 if six.PY3:
57 six.text_type = str
58
59 -class Particle(object):
60 """ """
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81 - def __init__(self, line=None, event=None):
82 """ """
83
84 if isinstance(line, Particle):
85 for key in line.__dict__:
86 setattr(self, key, getattr(line, key))
87 if event:
88 self.event = event
89 return
90 elif hepmc_parser and isinstance(line, hepmc_parser.HEPMC_Particle):
91 self.event = event
92 self.event_id = len(event)
93 for key in ['pid', 'status', 'E','px','py','pz','mass']:
94 setattr(self, key, getattr(line, key))
95 self.mother1 = 1
96 self.mother2 = 1
97 self.color1 = 0
98 self.color2 = 0
99 self.vtim = 0
100 self.comment = ''
101 self.helicity = 9
102 self.rwgt = 0
103 return
104
105
106 self.event = event
107 self.event_id = len(event)
108
109 self.pid = 0
110 self.status = 0
111 self.mother1 = None
112 self.mother2 = None
113 self.color1 = 0
114 self.color2 = None
115 self.px = 0
116 self.py = 0
117 self.pz = 0
118 self.E = 0
119 self.mass = 0
120 self.vtim = 0
121 self.helicity = 9
122 self.rwgt = 0
123 self.comment = ''
124
125 if line:
126 self.parse(line)
127
128 @property
130 "convenient alias"
131 return self.pid
132
133 - def parse(self, line):
134 """parse the line"""
135
136 args = line.split()
137 keys = ['pid', 'status','mother1','mother2','color1', 'color2', 'px','py','pz','E',
138 'mass','vtim','helicity']
139
140 for key,value in zip(keys,args):
141 setattr(self, key, float(value))
142 self.pid = int(self.pid)
143
144 self.comment = ' '.join(args[len(keys):])
145 if self.comment.startswith(('|','#')):
146 self.comment = self.comment[1:]
147
148
149
150
152 """string representing the particles"""
153 return " %8d %2d %4d %4d %4d %4d %+13.10e %+13.10e %+13.10e %14.10e %14.10e %10.4e %10.4e" \
154 % (self.pid,
155 self.status,
156 (self.mother1 if isinstance(self.mother1, numbers.Number) else self.mother1.event_id+1) if self.mother1 else 0,
157 (self.mother2 if isinstance(self.mother2, numbers.Number) else self.mother2.event_id+1) if self.mother2 else 0,
158 self.color1,
159 self.color2,
160 self.px,
161 self.py,
162 self.pz,
163 self.E,
164 self.mass,
165 self.vtim,
166 self.helicity)
167
168 - def __eq__(self, other):
169
170 if not isinstance(other, Particle):
171 return False
172 if self.pid == other.pid and \
173 self.status == other.status and \
174 self.mother1 == other.mother1 and \
175 self.mother2 == other.mother2 and \
176 self.color1 == other.color1 and \
177 self.color2 == other.color2 and \
178 self.px == other.px and \
179 self.py == other.py and \
180 self.pz == other.pz and \
181 self.E == other.E and \
182 self.mass == other.mass and \
183 self.vtim == other.vtim and \
184 self.helicity == other.helicity:
185 return True
186 return False
187
188 - def set_momentum(self, momentum):
189
190 self.E = momentum.E
191 self.px = momentum.px
192 self.py = momentum.py
193 self.pz = momentum.pz
194
195 - def add_decay(self, decay_event):
196 """associate to this particle the decay in the associate event"""
197
198 return self.event.add_decay_to_particle(self.event_id, decay_event)
199
200
201 - def __repr__(self):
202 return 'Particle("%s", event=%s)' % (str(self), self.event)
203
206 """A class to allow to read both gzip and not gzip file"""
207
208 allow_empty_event = False
209 encoding = 'UTF-8'
210
211 - def __init__(self, path, mode='r', *args, **opt):
212 """open file and read the banner [if in read mode]"""
213
214 if mode in ['r','rb']:
215 mode ='r'
216 self.mode = mode
217
218 self.to_zip = False
219 self.zip_mode = False
220
221 if not path.endswith(".gz"):
222 self.file = open(path, mode, *args, **opt)
223 elif mode == 'r' and not os.path.exists(path) and os.path.exists(path[:-3]):
224 self.file = open(path[:-3], mode, *args, **opt)
225 path = path[:-3]
226 else:
227 try:
228 self.file = gzip.GzipFile(path, mode, *args, **opt)
229 self.zip_mode =True
230 except IOError as error:
231 raise
232 except Exception as error:
233 misc.sprint(error)
234 if mode == 'r':
235 misc.gunzip(path)
236 else:
237 self.to_zip = True
238 self.file = open(path[:-3], mode, *args, **opt)
239 path = path[:-3]
240
241
242 self.path = path
243
244
245 self.parsing = True
246 self.eventgroup = False
247
248 self.banner = ''
249 if mode in ['r', 'rb']:
250 line = ''
251 while '</init>' not in line.lower():
252 line = self.file.readline()
253 if not line:
254 self.seek(0)
255 self.banner = ''
256 break
257 if 'b' in mode or self.zip_mode:
258 line = str(line.decode(self.encoding))
259 if '<event' in line.lower():
260 self.seek(0)
261 self.banner = ''
262 break
263
264 self.banner += line
265
275
276 @property
279
280 @property
283
284 @property
286 """return the cross-section of the file #from the banner"""
287 try:
288 return self._cross
289 except Exception:
290 pass
291
292 onebanner = self.get_banner()
293 self._cross = onebanner.get_cross()
294 return self._cross
295
309
312
319
320 __next__ = next
321
323
324 text = ''
325 line = ''
326 mode = 0
327
328 while True:
329
330 line = self.file.readline()
331 if not line:
332 raise StopIteration
333 if 'b' in self.mode or self.zip_mode:
334 line = line.decode(self.encoding)
335
336 if '<event' in line:
337 mode = 1
338 text = []
339 if mode:
340 text.append(line)
341
342 if '</event>' in line:
343 if self.parsing:
344 out = Event(text)
345 if len(out) == 0 and not self.allow_empty_event:
346 raise Exception
347 return out
348 else:
349 return text
350
351
353 events = []
354 text = ''
355 line = ''
356 mode = 0
357 while '</eventgroup>' not in line:
358
359
360 line = self.file.readline()
361 if not line:
362 raise StopIteration
363 if 'b' in self.mode:
364 line = line.decode(self.encoding)
365
366 if '<eventgroup' in line:
367 events=[]
368 text = ''
369 elif '<event' in line:
370 text = []
371 mode=1
372 elif '</event>' in line:
373 if self.parsing:
374 events.append(Event(text))
375 else:
376 events.append('\n'.join(text))
377 text = ''
378 mode = 0
379 if mode:
380 text += line
381 if len(events) == 0:
382 return self.next_eventgroup()
383
384 return events
385
386
388 """ scan once the file to return
389 - the list of the hightest weight (of size trunc_error*NB_EVENT
390 - the cross-section by type of process
391 - the total number of events in the file
392 """
393
394
395
396 self.seek(0)
397 all_wgt = []
398 cross = collections.defaultdict(int)
399 nb_event = 0
400 for event in self:
401 nb_event +=1
402 wgt = get_wgt(event)
403 cross['all'] += wgt
404 cross['abs'] += abs(wgt)
405 cross[event.ievent] += wgt
406 all_wgt.append(abs(wgt))
407
408 if nb_event % 20000 == 0:
409 all_wgt.sort()
410
411 nb_keep = max(20, int(nb_event*trunc_error*15))
412 all_wgt = all_wgt[-nb_keep:]
413
414
415 all_wgt.sort()
416
417 nb_keep = max(20, int(nb_event*trunc_error*10))
418 all_wgt = all_wgt[-nb_keep:]
419 self.seek(0)
420 return all_wgt, cross, nb_event
421
423 """ write a single events or a list of event
424 if self.eventgroup is ON, then add <eventgroup> around the lists of events
425 """
426 if isinstance(event, Event):
427 if self.eventgroup:
428 tmp = '<eventgroup>\n%s\n</eventgroup>\n' % event
429 self.write(tmp)
430 else:
431 self.write(str(event))
432 elif isinstance(event, list):
433 if self.eventgroup:
434 self.write('<eventgroup>\n')
435 for evt in event:
436 self.write(str(evt))
437 if self.eventgroup:
438 self.write('</eventgroup>\n')
439
440 - def unweight(self, outputpath, get_wgt=None, max_wgt=0, trunc_error=0,
441 event_target=0, log_level=logging.INFO, normalization='average'):
442 """unweight the current file according to wgt information wgt.
443 which can either be a fct of the event or a tag in the rwgt list.
444 max_wgt allow to do partial unweighting.
445 trunc_error allow for dynamical partial unweighting
446 event_target reweight for that many event with maximal trunc_error.
447 (stop to write event when target is reached)
448 """
449 if not get_wgt:
450 def weight(event):
451 return event.wgt
452 get_wgt = weight
453 unwgt_name = "central weight"
454 elif isinstance(get_wgt, (str,six.text_type)):
455 unwgt_name =get_wgt
456 def get_wgt(event):
457 event.parse_reweight()
458 return event.reweight_data[unwgt_name]
459 else:
460 unwgt_name = get_wgt.__name__
461
462
463 if hasattr(self, "written_weight"):
464 written_weight = lambda x: math.copysign(self.written_weight,float(x))
465 else:
466 written_weight = lambda x: x
467
468 all_wgt, cross, nb_event = self.initialize_unweighting(get_wgt, trunc_error)
469
470
471 def max_wgt_for_trunc(trunc):
472 """find the weight with the maximal truncation."""
473
474 xsum = 0
475 i=1
476 while (xsum - all_wgt[-i] * (i-1) <= cross['abs'] * trunc):
477 max_wgt = all_wgt[-i]
478 xsum += all_wgt[-i]
479 i +=1
480 if i == len(all_wgt):
481 break
482
483 return max_wgt
484
485
486
487 if not max_wgt:
488 if trunc_error == 0 or len(all_wgt)<2 or event_target:
489 max_wgt = all_wgt[-1]
490 else:
491 max_wgt = max_wgt_for_trunc(trunc_error)
492
493
494 if self.banner:
495 try:
496 import internal
497 except:
498 import madgraph.various.banner as banner_module
499 else:
500 import internal.banner as banner_module
501 if not isinstance(self.banner, banner_module.Banner):
502 banner = self.get_banner()
503
504 banner.modify_init_cross(cross)
505
506 banner["unweight"] = "unweighted by %s" % unwgt_name
507 else:
508 banner = self.banner
509
510 curr_strategy = banner.get_lha_strategy()
511 if normalization in ['unit', 'sum']:
512 strategy = 3
513 else:
514 strategy = 4
515 if curr_strategy >0:
516 banner.set_lha_strategy(abs(strategy))
517 else:
518 banner.set_lha_strategy(-1*abs(strategy))
519
520
521 nb_try = 20
522 nb_keep = 0
523 for i in range(nb_try):
524 self.seek(0)
525 if event_target:
526 if i==0:
527 max_wgt = max_wgt_for_trunc(0)
528 else:
529
530 efficiency = nb_keep/nb_event
531 needed_efficiency = event_target/nb_event
532 last_max_wgt = max_wgt
533 needed_max_wgt = last_max_wgt * efficiency / needed_efficiency
534
535 min_max_wgt = max_wgt_for_trunc(trunc_error)
536 max_wgt = max(min_max_wgt, needed_max_wgt)
537 max_wgt = min(max_wgt, all_wgt[-1])
538 if max_wgt == last_max_wgt:
539 if nb_keep < event_target and log_level>=10:
540 logger.log(log_level+10,"fail to reach target %s", event_target)
541 break
542 else:
543 break
544
545
546 if outputpath:
547 outfile = EventFile(outputpath, "w")
548
549
550 if self.banner and outputpath:
551 banner.write(outfile, close_tag=False)
552
553
554 nb_keep = 0
555 trunc_cross = 0
556 for event in self:
557 r = random.random()
558 wgt = get_wgt(event)
559 if abs(wgt) < r * max_wgt:
560 continue
561 elif wgt > 0:
562 nb_keep += 1
563 event.wgt = written_weight(max(wgt, max_wgt))
564 if abs(wgt) > max_wgt:
565 trunc_cross += abs(wgt) - max_wgt
566 if event_target ==0 or nb_keep <= event_target:
567 if outputpath:
568 outfile.write(str(event))
569
570 elif wgt < 0:
571 nb_keep += 1
572 event.wgt = -1* written_weight(max(abs(wgt), max_wgt))
573 if abs(wgt) > max_wgt:
574 trunc_cross += abs(wgt) - max_wgt
575 if outputpath and (event_target ==0 or nb_keep <= event_target):
576 outfile.write(str(event))
577
578 if event_target and nb_keep > event_target:
579 if not outputpath:
580
581 continue
582 elif event_target and i != nb_try-1 and nb_keep >= event_target *1.05:
583 outfile.write("</LesHouchesEvents>\n")
584 outfile.close()
585
586 continue
587 else:
588 outfile.write("</LesHouchesEvents>\n")
589 outfile.close()
590 break
591 elif event_target == 0:
592 if outputpath:
593 outfile.write("</LesHouchesEvents>\n")
594 outfile.close()
595 break
596 elif outputpath:
597 outfile.write("</LesHouchesEvents>\n")
598 outfile.close()
599
600
601 else:
602
603 logger.log(log_level+10,"fail to reach target event %s (iteration=%s)", event_target,i)
604
605
606
607
608 if event_target:
609 nb_events_unweighted = nb_keep
610 nb_keep = min( event_target, nb_keep)
611 else:
612 nb_events_unweighted = nb_keep
613
614 logger.log(log_level, "write %i event (efficiency %.2g %%, truncation %.2g %%) after %i iteration(s)",
615 nb_keep, nb_events_unweighted/nb_event*100, trunc_cross/cross['abs']*100, i)
616
617
618 if nb_keep != event_target and hasattr(self, "written_weight") and strategy !=4:
619 written_weight = lambda x: math.copysign(self.written_weight*event_target/nb_keep, float(x))
620 startfile = EventFile(outputpath)
621 tmpname = pjoin(os.path.dirname(outputpath), "wgtcorrected_"+ os.path.basename(outputpath))
622 outfile = EventFile(tmpname, "w")
623 outfile.write(startfile.banner)
624 for event in startfile:
625 event.wgt = written_weight(event.wgt)
626 outfile.write(str(event))
627 outfile.write("</LesHouchesEvents>\n")
628 startfile.close()
629 outfile.close()
630 shutil.move(tmpname, outputpath)
631
632
633
634
635 self.max_wgt = max_wgt
636 return nb_keep
637
639 """ apply one or more fct on all event. """
640
641 opt= {"print_step": 5000, "maxevent":float("inf"),'no_output':False}
642 opt.update(opts)
643 start = time.time()
644 nb_fct = len(fcts)
645 out = []
646 for i in range(nb_fct):
647 out.append([])
648 self.seek(0)
649 nb_event = 0
650 for event in self:
651 nb_event += 1
652 if opt["print_step"] and (nb_event % opt["print_step"]) == 0:
653 if hasattr(self,"len"):
654 print(("currently at %s/%s event [%is]" % (nb_event, self.len, time.time()-start)))
655 else:
656 print(("currently at %s event [%is]" % (nb_event, time.time()-start)))
657 for i in range(nb_fct):
658 value = fcts[i](event)
659 if not opt['no_output']:
660 out[i].append(value)
661 if nb_event > opt['maxevent']:
662 break
663 if nb_fct == 1:
664 return out[0]
665 else:
666 return out
667
668 - def split(self, nb_event=0, partition=None, cwd=os.path.curdir, zip=False):
669 """split the file in multiple file. Do not change the weight!"""
670
671 nb_file = -1
672 for i, event in enumerate(self):
673 if (not (partition is None) and i==sum(partition[:nb_file+1])) or \
674 (partition is None and i % nb_event == 0):
675 if i:
676
677 current.write('</LesHouchesEvent>\n')
678 current.close()
679
680 nb_file +=1
681
682 if not partition is None and (nb_file+1>len(partition)):
683 return nb_file
684 if zip:
685 current = EventFile(pjoin(cwd,'%s_%s.lhe.gz' % (self.path, nb_file)),'w')
686 else:
687 current = open(pjoin(cwd,'%s_%s.lhe' % (self.path, nb_file)),'w')
688 current.write(self.banner)
689 current.write(str(event))
690 if i!=0:
691 current.write('</LesHouchesEvent>\n')
692 current.close()
693
694 return nb_file +1
695
696 - def update_HwU(self, hwu, fct, name='lhe', keep_wgt=False, maxevents=sys.maxsize):
697 """take a HwU and add this event file for the function fct"""
698
699 if not isinstance(hwu, list):
700 hwu = [hwu]
701
702 class HwUUpdater(object):
703
704 def __init__(self, fct, keep_wgt):
705
706 self.fct = fct
707 self.first = True
708 self.keep_wgt = keep_wgt
709
710 def add(self, event):
711
712 value = self.fct(event)
713
714 if self.first:
715 for h in hwu:
716
717 if isinstance(value, dict):
718 h.add_line(list(value.keys()))
719 else:
720
721 h.add_line(name)
722 if self.keep_wgt is True:
723 event.parse_reweight()
724 h.add_line(['%s_%s' % (name, key)
725 for key in event.reweight_data])
726 elif self.keep_wgt:
727 h.add_line(list(self.keep_wgt.values()))
728 self.first = False
729
730 for h in hwu:
731 if isinstance(value, tuple):
732 h.addEvent(value[0], value[1])
733 else:
734 h.addEvent(value,{name:event.wgt})
735 if self.keep_wgt:
736 event.parse_reweight()
737 if self.keep_wgt is True:
738 data = dict(('%s_%s' % (name, key),event.reweight_data[key])
739 for key in event.reweight_data)
740 h.addEvent(value, data)
741 else:
742 data = dict(( value,event.reweight_data[key])
743 for key,value in self.keep_wgt.items())
744 h.addEvent(value, data)
745
746
747
748 self.apply_fct_on_event(HwUUpdater(fct,keep_wgt).add, no_output=True,maxevent=maxevents)
749 return hwu
750
752 """take the lhe file and add the matchscale from the pythia_input file"""
753
754 if pythia_input:
755 def next_data():
756 for line in open(pythia_input):
757 if line.startswith('#'):
758 continue
759 data = line.split()
760 print((int(data[0]), data[-3], data[-2], data[-1]))
761 yield (int(data[0]), data[-3], data[-2], data[-1])
762 else:
763 def next_data():
764 i=0
765 while 1:
766 yield [i,0,0,0]
767 i+=1
768 sys_iterator = next_data()
769
770 self.seek(0)
771 out = open(out_path,'w')
772
773 pdf_pattern = re.compile(r'''<init>(.*)</init>''', re.M+re.S)
774 init = pdf_pattern.findall(self.banner)[0].split('\n',2)[1]
775 id1, id2, _, _, _, _, pdf1,pdf2,_,_ = init.split()
776 id = [int(id1), int(id2)]
777 type = []
778 for i in range(2):
779 if abs(id[i]) == 2212:
780 if i > 0:
781 type.append(1)
782 else:
783 type.append(-1)
784 else:
785 type.append(0)
786 pdf = max(int(pdf1),int(pdf2))
787
788 out.write("<header>\n" + \
789 "<orgpdf>%i</orgpdf>\n" % pdf + \
790 "<beams> %s %s</beams>\n" % tuple(type) + \
791 "</header>\n")
792
793
794 nevt, smin, smax, scomp = next(sys_iterator)
795 for i, orig_event in enumerate(self):
796 if i < nevt:
797 continue
798 new_event = Event()
799 sys = orig_event.parse_syscalc_info()
800 new_event.syscalc_data = sys
801 if smin:
802 new_event.syscalc_data['matchscale'] = "%s %s %s" % (smin, scomp, smax)
803 out.write(str(new_event), nevt)
804 try:
805 nevt, smin, smax, scomp = next(sys_iterator)
806 except StopIteration:
807 break
808
809 - def get_alphas(self, scale, lhapdf_config='lhapdf-config'):
810 """return the alphas value associated to a given scale"""
811
812 if hasattr(self, 'alpsrunner'):
813 return self.alpsrunner(scale)
814
815
816 banner = banner_mod.Banner(self.banner)
817 run_card = banner.charge_card('run_card')
818 use_runner = False
819 if abs(run_card['lpp1']) != 1 and abs(run_card['lpp2']) != 1:
820
821 use_runner = True
822 else:
823
824 lhapdf = misc.import_python_lhapdf(lhapdf_config)
825 if not lhapdf:
826 logger.warning('fail to link to lhapdf for the alphas-running. Use Two loop computation')
827 use_runner = True
828 try:
829 self.pdf = lhapdf.mkPDF(int(self.banner.run_card.get_lhapdf_id()))
830 except Exception:
831 logger.warning('fail to link to lhapdf for the alphas-running. Use Two loop computation')
832 use_runner = True
833
834 if not use_runner:
835 self.alpsrunner = lambda scale: self.pdf.alphasQ(scale)
836 else:
837 try:
838 from models.model_reader import Alphas_Runner
839 except ImportError:
840 root = os.path.dirname(__file__)
841 root_path = pjoin(root, os.pardir, os.pardir)
842 try:
843 import internal.madevent_interface as me_int
844 cmd = me_int.MadEventCmd(root_path,force_run=True)
845 except ImportError:
846 import internal.amcnlo_run_interface as me_int
847 cmd = me_int.Cmd(root_path,force_run=True)
848 if 'mg5_path' in cmd.options and cmd.options['mg5_path']:
849 sys.path.append(cmd.options['mg5_path'])
850 from models.model_reader import Alphas_Runner
851
852 if not hasattr(banner, 'param_card'):
853 param_card = banner.charge_card('param_card')
854 else:
855 param_card = banner.param_card
856
857 asmz = param_card.get_value('sminputs', 3, 0.13)
858 nloop =2
859 zmass = param_card.get_value('mass', 23, 91.188)
860 cmass = param_card.get_value('mass', 4, 1.4)
861 if cmass == 0:
862 cmass = 1.4
863 bmass = param_card.get_value('mass', 5, 4.7)
864 if bmass == 0:
865 bmass = 4.7
866 self.alpsrunner = Alphas_Runner(asmz, nloop, zmass, cmass, bmass)
867
868
869
870 return self.alpsrunner(scale)
871
872 - def seek(self, *args, **opts):
874
876 if self.zipmode:
877 currpos = self.file.tell()
878 if not currpos:
879 currpos = self.size
880 return currpos
881 else:
882 self.file.tell()
883
885
886 if self.zip_mode or 'b' in self.mode:
887 self.file.write(text.encode())
888 else:
889 self.file.write(text)
890
891 - def close(self,*args, **opts):
896
898 try:
899 self.file.close()
900 except Exception:
901 pass
902
905 """a class to read simultaneously multiple file and read them in mixing them.
906 Unweighting can be done at the same time.
907 The number of events in each file need to be provide in advance
908 (if not provide the file is first read to find that number"""
909
910 - def __new__(cls, start_list=[],parse=True):
912
913 - def __init__(self, start_list=[], parse=True):
914 """if trunc_error is define here then this allow
915 to only read all the files twice and not three times."""
916 self.eventgroup = False
917 self.files = []
918
919 self.parsefile = parse
920 self.banner = ''
921 self.initial_nb_events = []
922 self.total_event_in_files = 0
923 self.curr_nb_events = []
924 self.allcross = []
925 self.error = []
926 self.across = []
927 self.scales = []
928 if start_list:
929 if parse:
930 for p in start_list:
931 self.add(p)
932 else:
933 self.files = start_list
934
935 self._configure = False
936
937 - def close(self,*args,**opts):
940
941 - def add(self, path, cross, error, across, nb_event=0, scale=1):
942 """ add a file to the pool, across allow to reweight the sum of weight
943 in the file to the given cross-section
944 """
945
946 if across == 0:
947
948 return
949 obj = EventFile(path)
950 obj.eventgroup = self.eventgroup
951 if len(self.files) == 0 and not self.banner:
952 self.banner = obj.banner
953 self.curr_nb_events.append(0)
954 self.initial_nb_events.append(0)
955 self.allcross.append(cross)
956 self.across.append(across)
957 self.error.append(error)
958 self.scales.append(scale)
959 self.files.append(obj)
960
961 if nb_event:
962 obj.len = nb_event
963 self._configure = False
964 return obj
965
967
968 if not self._configure:
969 self.configure()
970 return self
971
973 if not self._configure:
974 self.configure()
975 remaining_event = self.total_event_in_files - sum(self.curr_nb_events)
976 if remaining_event == 0:
977 raise StopIteration
978
979 nb_event = random.randint(1, remaining_event)
980 sum_nb=0
981 for i, obj in enumerate(self.files):
982 sum_nb += self.initial_nb_events[i] - self.curr_nb_events[i]
983 if nb_event <= sum_nb:
984 self.curr_nb_events[i] += 1
985 event = next(obj)
986 if not self.eventgroup:
987 event.sample_scale = self.scales[i]
988 else:
989 for evt in event:
990 evt.sample_scale = self.scales[i]
991 return event
992 else:
993 raise StopIteration
994
995
996 __next__ = next
997
999 """define the part of the init_banner"""
1000
1001 if not self.banner:
1002 return
1003
1004
1005 grouped_cross = {}
1006 grouped_error = {}
1007 for i,ff in enumerate(self.files):
1008 filename = ff.path
1009 from_init = False
1010 Pdir = [P for P in filename.split(os.path.sep) if P.startswith('P')]
1011 if Pdir:
1012 Pdir = Pdir[-1]
1013 group = Pdir.split("_")[0][1:]
1014 if not group.isdigit():
1015 from_init = True
1016 else:
1017 from_init = True
1018
1019 if not from_init:
1020 if group in grouped_cross:
1021 grouped_cross[group] += self.allcross[i]
1022 grouped_error[group] += self.error[i]**2
1023 else:
1024 grouped_cross[group] = self.allcross[i]
1025 grouped_error[group] = self.error[i]**2
1026 else:
1027 ban = banner_mod.Banner(ff.banner)
1028 for line in ban['init'].split('\n'):
1029 splitline = line.split()
1030 if len(splitline)==4:
1031 cross, error, _, group = splitline
1032 if int(group) in grouped_cross:
1033 grouped_cross[group] += float(cross)
1034 grouped_error[group] += float(error)**2
1035 else:
1036 grouped_cross[group] = float(cross)
1037 grouped_error[group] = float(error)**2
1038 nb_group = len(grouped_cross)
1039
1040
1041 try:
1042 run_card = self.banner.run_card
1043 except:
1044 run_card = self.banner.charge_card("run_card")
1045
1046 init_information = run_card.get_banner_init_information()
1047
1048 if proc_charac and proc_charac['ninitial'] == 1:
1049
1050 init_information = run_card.get_banner_init_information()
1051 event = next(self)
1052 init_information["idbmup1"] = event[0].pdg
1053 init_information["ebmup1"] = event[0].mass
1054 init_information["idbmup2"] = 0
1055 init_information["ebmup2"] = 0
1056 self.seek(0)
1057 else:
1058
1059 if init_information["idbmup1"] in [0,9]:
1060 event = next(self)
1061 init_information["idbmup1"]= event[0].pdg
1062 if init_information["idbmup2"] == 0:
1063 init_information["idbmup2"]= event[1].pdg
1064 self.seek(0)
1065 if init_information["idbmup2"] in [0,9]:
1066 event = next(self)
1067 init_information["idbmup2"] = event[1].pdg
1068 self.seek(0)
1069
1070 init_information["nprup"] = nb_group
1071
1072 if run_card["lhe_version"] < 3:
1073 init_information["generator_info"] = ""
1074 else:
1075 init_information["generator_info"] = "<generator name='MadGraph5_aMC@NLO' version='%s'>please cite 1405.0301 </generator>\n" \
1076 % misc.get_pkg_info()['version']
1077
1078
1079 cross_info = "%(cross)e %(error)e %(wgt)e %(id)i"
1080 init_information["cross_info"] = []
1081 for id in grouped_cross:
1082 conv = {"id": int(id), "cross": grouped_cross[id], "error": math.sqrt(grouped_error[id]),
1083 "wgt": wgt}
1084 init_information["cross_info"].append( cross_info % conv)
1085 init_information["cross_info"] = '\n'.join(init_information["cross_info"])
1086 init_information['lha_stra'] = -1 * abs(lha_strategy)
1087
1088 template_init =\
1089 """ %(idbmup1)i %(idbmup2)i %(ebmup1)e %(ebmup2)e %(pdfgup1)i %(pdfgup2)i %(pdfsup1)i %(pdfsup2)i %(lha_stra)i %(nprup)i
1090 %(cross_info)s
1091 %(generator_info)s
1092 """
1093
1094 self.banner["init"] = template_init % init_information
1095
1096
1097
1099 """ scan once the file to return
1100 - the list of the hightest weight (of size trunc_error*NB_EVENT
1101 - the cross-section by type of process
1102 - the total number of events in the files
1103 In top of that it initialise the information for the next routine
1104 to determine how to choose which file to read
1105 """
1106 self.seek(0)
1107 all_wgt = []
1108 total_event = 0
1109 sum_cross = collections.defaultdict(int)
1110 for i,f in enumerate(self.files):
1111 nb_event = 0
1112
1113
1114 cross = collections.defaultdict(int)
1115 new_wgt =[]
1116 for event in f:
1117 nb_event += 1
1118 total_event += 1
1119 event.sample_scale = 1
1120 wgt = getwgt(event)
1121 cross['all'] += wgt
1122 cross['abs'] += abs(wgt)
1123 cross[event.ievent] += wgt
1124 new_wgt.append(abs(wgt))
1125
1126 if nb_event % 20000 == 0:
1127 new_wgt.sort()
1128
1129 nb_keep = max(20, int(nb_event*trunc_error*15))
1130 new_wgt = new_wgt[-nb_keep:]
1131 if nb_event == 0:
1132 raise Exception
1133
1134 self.initial_nb_events[i] = nb_event
1135 self.scales[i] = self.across[i]/cross['abs'] if self.across[i] else 1
1136
1137
1138 for key in cross:
1139 sum_cross[key] += cross[key]* self.scales[i]
1140 all_wgt +=[self.scales[i] * w for w in new_wgt]
1141 all_wgt.sort()
1142 nb_keep = max(20, int(total_event*trunc_error*10))
1143 all_wgt = all_wgt[-nb_keep:]
1144
1145 self.total_event_in_files = total_event
1146
1147 all_wgt.sort()
1148
1149 nb_keep = max(20, int(total_event*trunc_error*10))
1150 all_wgt = all_wgt[-nb_keep:]
1151 self.seek(0)
1152 self._configure = True
1153 return all_wgt, sum_cross, total_event
1154
1161
1163
1164 return len(self.files)
1165
1166 - def seek(self, pos):
1167 """ """
1168
1169 if pos !=0:
1170 raise Exception
1171 for i in range(len(self)):
1172 self.curr_nb_events[i] = 0
1173 for f in self.files:
1174 f.seek(pos)
1175
1176 - def unweight(self, outputpath, get_wgt, **opts):
1177 """unweight the current file according to wgt information wgt.
1178 which can either be a fct of the event or a tag in the rwgt list.
1179 max_wgt allow to do partial unweighting.
1180 trunc_error allow for dynamical partial unweighting
1181 event_target reweight for that many event with maximal trunc_error.
1182 (stop to write event when target is reached)
1183 """
1184
1185
1186 if isinstance(get_wgt, (str,six.text_type)):
1187 unwgt_name =get_wgt
1188 def get_wgt_multi(event):
1189 event.parse_reweight()
1190 return event.reweight_data[unwgt_name] * event.sample_scale
1191 else:
1192 unwgt_name = get_wgt.__name__
1193 get_wgt_multi = lambda event: get_wgt(event) * event.sample_scale
1194
1195
1196 if 'proc_charac' in opts:
1197 if opts['proc_charac']:
1198 proc_charac = opts['proc_charac']
1199 else:
1200 proc_charac=None
1201 del opts['proc_charac']
1202 else:
1203 proc_charac = None
1204
1205 if 'event_target' in opts and opts['event_target']:
1206 if 'normalization' in opts:
1207 if opts['normalization'] == 'sum':
1208 new_wgt = sum(self.across)/opts['event_target']
1209 strategy = 3
1210 elif opts['normalization'] == 'average':
1211 strategy = 4
1212 new_wgt = sum(self.across)
1213 elif opts['normalization'] == 'unit':
1214 strategy =3
1215 new_wgt = 1.
1216 else:
1217 strategy = 4
1218 new_wgt = sum(self.across)
1219 self.define_init_banner(new_wgt, strategy, proc_charac=proc_charac)
1220 self.written_weight = new_wgt
1221 elif 'write_init' in opts and opts['write_init']:
1222 self.define_init_banner(0,0, proc_charac=proc_charac)
1223 del opts['write_init']
1224 return super(MultiEventFile, self).unweight(outputpath, get_wgt_multi, **opts)
1225
1226 - def write(self, path, random=False, banner=None, get_info=False):
1280
1289
1290
1291
1292 -class Event(list):
1293 """Class storing a single event information (list of particles + global information)"""
1294
1295 warning_order = True
1296
1298 """The initialization of an empty Event (or one associate to a text file)"""
1299 list.__init__(self)
1300
1301
1302 self.nexternal = 0
1303 self.ievent = 0
1304 self.wgt = 0
1305 self.aqcd = 0
1306 self.scale = 0
1307 self.aqed = 0
1308 self.aqcd = 0
1309
1310 self.tag = ''
1311 self.eventflag = {}
1312 self.comment = ''
1313 self.reweight_data = {}
1314 self.matched_scale_data = None
1315 self.syscalc_data = {}
1316 if text:
1317 self.parse(text)
1318
1319
1320
1322 """Take the input file and create the structured information"""
1323
1324 status = 'first'
1325
1326 if not isinstance(text, list):
1327 text = text.split('\n')
1328
1329 for line in text:
1330 line = line.strip()
1331 if not line:
1332 continue
1333 elif line[0] == '#':
1334 self.comment += '%s\n' % line
1335 continue
1336 elif line.startswith('<event'):
1337 if '=' in line:
1338 found = re.findall(r"""(\w*)=(?:(?:['"])([^'"]*)(?=['"])|(\S*))""",line)
1339
1340
1341 self.eventflag = dict((n, a1) if a1 else (n,a2) for n,a1,a2 in found)
1342
1343 continue
1344
1345 elif 'first' == status:
1346 if '<rwgt>' in line:
1347 status = 'tag'
1348 else:
1349 self.assign_scale_line(line)
1350 status = 'part'
1351 continue
1352 if '<' in line:
1353 status = 'tag'
1354
1355 if 'part' == status:
1356 part = Particle(line, event=self)
1357 if part.E != 0 or part.status==-1:
1358 self.append(part)
1359 elif self.nexternal:
1360 self.nexternal-=1
1361 else:
1362 if '</event>' in line:
1363 line = line.replace('</event>','',1)
1364 self.tag += '%s\n' % line
1365
1366 self.assign_mother()
1367
1368
1370 """convert the number in actual particle"""
1371
1372 if all(p.status != -1 for p in self):
1373 if not self.nexternal:
1374 return
1375 if self.warning_order:
1376 Event.warning_order = False
1377 logger.warning("Weird format for lhe format: no incoming particle... adding a fake one")
1378 raise Exception
1379 mother = Particle(event=self)
1380 mother.status = -1
1381 mother.pid = 0
1382 self.insert(0,mother)
1383 mother.color2 = 0
1384 mother.event_id = 0
1385 self.nexternal += 1
1386 for p in self[1:]:
1387 if isinstance(p.mother1, int) and p.mother1 > 1:
1388 p.mother1 += 1
1389 if isinstance(p.mother2, int) and p.mother2 > 1:
1390 p.mother2 += 1
1391 p.event_id += 1
1392
1393
1394
1395 for i,particle in enumerate(self):
1396 if i < particle.mother1 or i < particle.mother2:
1397 if self.warning_order:
1398 logger.warning("Order of particle in the event did not agree with parent/child order. This might be problematic for some code.")
1399 Event.warning_order = False
1400 self.reorder_mother_child()
1401 return self.assign_mother()
1402
1403 if particle.mother1:
1404 try:
1405 particle.mother1 = self[int(particle.mother1) -1]
1406 except Exception:
1407 logger.warning("WRONG MOTHER INFO %s", self)
1408 particle.mother1 = 0
1409 if particle.mother2:
1410 try:
1411 particle.mother2 = self[int(particle.mother2) -1]
1412 except Exception:
1413 logger.warning("WRONG MOTHER INFO %s", self)
1414 particle.mother2 = 0
1415
1417 """change all the weights by a given ratio"""
1418
1419 self.wgt *= ratio
1420 self.parse_reweight()
1421 for key in self.reweight_data:
1422 self.reweight_data[key] *= ratio
1423 return self.wgt
1424
1426 """check and correct the mother/child position.
1427 only correct one order by call (but this is a recursive call)"""
1428
1429 tomove, position = None, None
1430 for i,particle in enumerate(self):
1431 if i < particle.mother1:
1432
1433 tomove, position = i, particle.mother1-1
1434 break
1435 if i < particle.mother2:
1436 tomove, position = i, particle.mother2-1
1437
1438
1439 if not tomove:
1440 return
1441
1442
1443 particle = self.pop(tomove)
1444 self.insert(int(position), particle)
1445
1446
1447 for i, particle in enumerate(self):
1448 particle.event_id = i
1449
1450 m1, m2 = particle.mother1, particle.mother2
1451 if m1 == tomove +1:
1452 particle.mother1 = position+1
1453 elif tomove < m1 <= position +1:
1454 particle.mother1 -= 1
1455 if m2 == tomove +1:
1456 particle.mother2 = position+1
1457 elif tomove < m2 <= position +1:
1458 particle.mother2 -= 1
1459
1460 return self.reorder_mother_child()
1461
1462
1463
1464
1465
1466
1468 """Parse the re-weight information in order to return a dictionary
1469 {key: value}. If no group is define group should be '' """
1470 if self.reweight_data:
1471 return self.reweight_data
1472 self.reweight_data = {}
1473 self.reweight_order = []
1474 start, stop = self.tag.find('<rwgt>'), self.tag.find('</rwgt>')
1475 if start != -1 != stop :
1476 pattern = re.compile(r'''<\s*wgt id=(?:\'|\")(?P<id>[^\'\"]+)(?:\'|\")\s*>\s*(?P<val>[\ded+-.]*)\s*</wgt>''',re.I)
1477 data = pattern.findall(self.tag[start:stop])
1478 try:
1479 self.reweight_data = dict([(pid, float(value)) for (pid, value) in data
1480 if not self.reweight_order.append(pid)])
1481
1482 except ValueError as error:
1483 raise Exception('Event File has unvalid weight. %s' % error)
1484 self.tag = self.tag[:start] + self.tag[stop+7:]
1485 return self.reweight_data
1486
1488 """ """
1489 if hasattr(self, 'nloweight'):
1490 return self.nloweight
1491
1492 start, stop = self.tag.find('<mgrwgt>'), self.tag.find('</mgrwgt>')
1493 if start != -1 != stop :
1494
1495 text = self.tag[start+8:stop]
1496 self.nloweight = NLO_PARTIALWEIGHT(text, self, real_type=real_type,
1497 threshold=threshold)
1498 return self.nloweight
1499
1501 """get the string associate to the weight"""
1502
1503 text="""<mgrwgt>
1504 %(total_wgt).10e %(nb_wgt)i %(nb_event)i 0
1505 %(event)s
1506 %(wgt)s
1507 </mgrwgt>"""
1508
1509
1510 if not wgt:
1511 if not hasattr(self, 'nloweight'):
1512 return
1513 wgt = self.nloweight
1514
1515 data = {'total_wgt': wgt.total_wgt,
1516 'nb_wgt': wgt.nb_wgt,
1517 'nb_event': wgt.nb_event,
1518 'event': '\n'.join(p.__str__(mode='fortran') for p in wgt.momenta),
1519 'wgt':'\n'.join(w.__str__(mode='formatted')
1520 for e in wgt.cevents for w in e.wgts)}
1521
1522 data['total_wgt'] = sum([w.ref_wgt for e in wgt.cevents for w in e.wgts])
1523 start, stop = self.tag.find('<mgrwgt>'), self.tag.find('</mgrwgt>')
1524
1525 self.tag = self.tag[:start] + text % data + self.tag[stop+9:]
1526
1527
1529 """ """
1530
1531
1532 if hasattr(self, 'loweight'):
1533 return self.loweight
1534
1535 if not hasattr(Event, 'loweight_pattern'):
1536 Event.loweight_pattern = re.compile('''<rscale>\s*(?P<nqcd>\d+)\s+(?P<ren_scale>[\d.e+-]+)\s*</rscale>\s*\n\s*
1537 <asrwt>\s*(?P<asrwt>[\s\d.+-e]+)\s*</asrwt>\s*\n\s*
1538 <pdfrwt\s+beam=["']?(?P<idb1>1|2)["']?\>\s*(?P<beam1>[\s\d.e+-]*)\s*</pdfrwt>\s*\n\s*
1539 <pdfrwt\s+beam=["']?(?P<idb2>1|2)["']?\>\s*(?P<beam2>[\s\d.e+-]*)\s*</pdfrwt>\s*\n\s*
1540 <totfact>\s*(?P<totfact>[\d.e+-]*)\s*</totfact>
1541 ''',re.X+re.I+re.M)
1542
1543 start, stop = self.tag.find('<mgrwt>'), self.tag.find('</mgrwt>')
1544
1545 if start != -1 != stop :
1546 text = self.tag[start+8:stop]
1547
1548 info = Event.loweight_pattern.search(text)
1549 if not info:
1550 raise Exception('%s not parsed'% text)
1551 self.loweight={}
1552 self.loweight['n_qcd'] = int(info.group('nqcd'))
1553 self.loweight['ren_scale'] = float(info.group('ren_scale'))
1554 self.loweight['asrwt'] =[float(x) for x in info.group('asrwt').split()[1:]]
1555 self.loweight['tot_fact'] = float(info.group('totfact'))
1556
1557 if info.group('idb1') == info.group('idb2'):
1558 raise Exception('%s not parsed'% text)
1559
1560 if info.group('idb1') =="1":
1561 args = info.group('beam1').split()
1562 else:
1563 args = info.group('beam2').split()
1564 npdf = int(args[0])
1565 self.loweight['n_pdfrw1'] = npdf
1566 self.loweight['pdf_pdg_code1'] = [int(i) for i in args[1:1+npdf]]
1567 self.loweight['pdf_x1'] = [float(i) for i in args[1+npdf:1+2*npdf]]
1568 self.loweight['pdf_q1'] = [float(i) for i in args[1+2*npdf:1+3*npdf]]
1569 if info.group('idb2') =="2":
1570 args = info.group('beam2').split()
1571 else:
1572 args = info.group('beam1').split()
1573 npdf = int(args[0])
1574 self.loweight['n_pdfrw2'] = npdf
1575 self.loweight['pdf_pdg_code2'] = [int(i) for i in args[1:1+npdf]]
1576 self.loweight['pdf_x2'] = [float(i) for i in args[1+npdf:1+2*npdf]]
1577 self.loweight['pdf_q2'] = [float(i) for i in args[1+2*npdf:1+3*npdf]]
1578
1579 else:
1580 return None
1581 return self.loweight
1582
1583
1585 """Parse the line containing the starting scale for the shower"""
1586
1587 if self.matched_scale_data is not None:
1588 return self.matched_scale_data
1589
1590 self.matched_scale_data = []
1591
1592
1593 pattern = re.compile("<scales\s|</scales>")
1594 data = re.split(pattern,self.tag)
1595 if len(data) == 1:
1596 return []
1597 else:
1598 tmp = {}
1599 start,content, end = data
1600 self.tag = "%s%s" % (start, end)
1601 pattern = re.compile("pt_clust_(\d*)=\"([\de+-.]*)\"")
1602 for id,value in pattern.findall(content):
1603 tmp[int(id)] = float(value)
1604 for i in range(1, len(self)+1):
1605 if i in tmp:
1606 self.matched_scale_data.append(tmp[i])
1607 else:
1608 self.matched_scale_data.append(-1)
1609 return self.matched_scale_data
1610
1612 """ parse the flag for syscalc between <mgrwt></mgrwt>
1613 <mgrwt>
1614 <rscale> 3 0.26552898E+03</rscale>
1615 <asrwt>0</asrwt>
1616 <pdfrwt beam="1"> 1 21 0.14527945E+00 0.26552898E+03</pdfrwt>
1617 <pdfrwt beam="2"> 1 21 0.15249110E-01 0.26552898E+03</pdfrwt>
1618 <totfact> 0.10344054E+04</totfact>
1619 </mgrwt>
1620 """
1621 if self.syscalc_data:
1622 return self.syscalc_data
1623
1624 pattern = re.compile("<mgrwt>|</mgrwt>")
1625 pattern2 = re.compile("<(?P<tag>[\w]*)(?:\s*(\w*)=[\"'](.*)[\"']\s*|\s*)>(.*)</(?P=tag)>")
1626 data = re.split(pattern,self.tag)
1627 if len(data) == 1:
1628 return []
1629 else:
1630 tmp = {}
1631 start,content, end = data
1632 self.tag = "%s%s" % (start, end)
1633 for tag, key, keyval, tagval in pattern2.findall(content):
1634 if key:
1635 self.syscalc_data[(tag, key, keyval)] = tagval
1636 else:
1637 self.syscalc_data[tag] = tagval
1638 return self.syscalc_data
1639
1640
1641 - def add_decay_to_particle(self, position, decay_event):
1642 """define the decay of the particle id by the event pass in argument"""
1643
1644 this_particle = self[position]
1645
1646 this_particle.status = 2
1647 this_particle.helicity = 0
1648
1649
1650 decay_particle = decay_event[0]
1651 this_4mom = FourMomentum(this_particle)
1652 nb_part = len(self)
1653
1654 thres = decay_particle.E*1e-10
1655 assert max(decay_particle.px, decay_particle.py, decay_particle.pz) < thres,\
1656 "not on rest particle %s %s %s %s" % (decay_particle.E, decay_particle.px,decay_particle.py,decay_particle.pz)
1657
1658 self.nexternal += decay_event.nexternal -1
1659 old_scales = list(self.parse_matching_scale())
1660 if old_scales:
1661 jet_position = sum(1 for i in range(position) if self[i].status==1)
1662 initial_pos = sum(1 for i in range(position) if self[i].status==-1)
1663 self.matched_scale_data.pop(initial_pos+jet_position)
1664
1665
1666 for particle in decay_event[1:]:
1667
1668 new_particle = Particle(particle, self)
1669 new_particle.event_id = len(self)
1670 self.append(new_particle)
1671 if old_scales:
1672 self.matched_scale_data.append(old_scales[initial_pos+jet_position])
1673
1674 new_momentum = FourMomentum(new_particle).boost(this_4mom)
1675 new_particle.set_momentum(new_momentum)
1676
1677 for tag in ['mother1', 'mother2']:
1678 mother = getattr(particle, tag)
1679 if isinstance(mother, Particle):
1680 mother_id = getattr(particle, tag).event_id
1681 if mother_id == 0:
1682 setattr(new_particle, tag, this_particle)
1683 else:
1684 try:
1685 setattr(new_particle, tag, self[nb_part + mother_id -1])
1686 except Exception as error:
1687 print(error)
1688 misc.sprint( self)
1689 misc.sprint(nb_part + mother_id -1)
1690 misc.sprint(tag)
1691 misc.sprint(position, decay_event)
1692 misc.sprint(particle)
1693 misc.sprint(len(self), nb_part + mother_id -1)
1694 raise
1695 elif tag == "mother2" and isinstance(particle.mother1, Particle):
1696 new_particle.mother2 = this_particle
1697 else:
1698 raise Exception("Something weird happens. Please report it for investigation")
1699
1700
1701 max_color=501
1702 for particle in self[:nb_part]:
1703 max_color=max(max_color, particle.color1, particle.color2)
1704
1705
1706 color_mapping = {}
1707 color_mapping[decay_particle.color1] = this_particle.color1
1708 color_mapping[decay_particle.color2] = this_particle.color2
1709 for particle in self[nb_part:]:
1710 if particle.color1:
1711 if particle.color1 not in color_mapping:
1712 max_color +=1
1713 color_mapping[particle.color1] = max_color
1714 particle.color1 = max_color
1715 else:
1716 particle.color1 = color_mapping[particle.color1]
1717 if particle.color2:
1718 if particle.color2 not in color_mapping:
1719 max_color +=1
1720 color_mapping[particle.color2] = max_color
1721 particle.color2 = max_color
1722 else:
1723 particle.color2 = color_mapping[particle.color2]
1724
1726 """use auto-recursion"""
1727
1728 pdg_to_decay = dict(pdg_to_decay)
1729
1730 for i,particle in enumerate(self):
1731 if particle.status != 1:
1732 continue
1733 if particle.pdg in pdg_to_decay and pdg_to_decay[particle.pdg]:
1734 one_decay = pdg_to_decay[particle.pdg].pop()
1735 self.add_decay_to_particle(i, one_decay)
1736 return self.add_decays(pdg_to_decay)
1737 return self
1738
1739
1740
1742
1743 to_remove = []
1744 if event_id is not None:
1745 to_remove.append(self[event_id])
1746
1747 if pdg_code:
1748 for particle in self:
1749 if particle.pid == pdg_code:
1750 to_remove.append(particle)
1751
1752 new_event = Event()
1753
1754 for tag in ['nexternal', 'ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']:
1755 setattr(new_event, tag, getattr(self, tag))
1756
1757 for particle in self:
1758 if isinstance(particle.mother1, Particle) and particle.mother1 in to_remove:
1759 to_remove.append(particle)
1760 if particle.status == 1:
1761 new_event.nexternal -= 1
1762 continue
1763 elif isinstance(particle.mother2, Particle) and particle.mother2 in to_remove:
1764 to_remove.append(particle)
1765 if particle.status == 1:
1766 new_event.nexternal -= 1
1767 continue
1768 else:
1769 new_event.append(Particle(particle))
1770
1771
1772
1773 for pos, particle in enumerate(new_event):
1774 particle.event_id = pos
1775 if particle in to_remove:
1776 particle.status = 1
1777 return new_event
1778
1779 - def get_decay(self, pdg_code=0, event_id=None):
1780
1781 to_start = []
1782 if event_id is not None:
1783 to_start.append(self[event_id])
1784
1785 elif pdg_code:
1786 for particle in self:
1787 if particle.pid == pdg_code:
1788 to_start.append(particle)
1789 break
1790
1791 new_event = Event()
1792
1793 for tag in ['ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']:
1794 setattr(new_event, tag, getattr(self, tag))
1795
1796
1797 old2new = {}
1798 new_decay_part = Particle(to_start[0])
1799 new_decay_part.mother1 = None
1800 new_decay_part.mother2 = None
1801 new_decay_part.status = -1
1802 old2new[new_decay_part.event_id] = len(old2new)
1803 new_event.append(new_decay_part)
1804
1805
1806
1807 for particle in self:
1808 if isinstance(particle.mother1, Particle) and particle.mother1.event_id in old2new\
1809 or isinstance(particle.mother2, Particle) and particle.mother2.event_id in old2new:
1810 old2new[particle.event_id] = len(old2new)
1811 new_event.append(Particle(particle))
1812
1813
1814
1815 nexternal = 0
1816 for pos, particle in enumerate(new_event):
1817 particle.event_id = pos
1818 if particle.mother1:
1819 particle.mother1 = new_event[old2new[particle.mother1.event_id]]
1820 if particle.mother2:
1821 particle.mother2 = new_event[old2new[particle.mother2.event_id]]
1822 if particle.status in [-1,1]:
1823 nexternal +=1
1824 new_event.nexternal = nexternal
1825
1826 return new_event
1827
1828 - def boost(self, filter=None):
1829 """modify the current event to boost it according to the current filter"""
1830 if filter is None:
1831 filter = lambda p: p.status==-1
1832
1833 if not isinstance(filter, FourMomentum):
1834 pboost = FourMomentum()
1835 for p in self:
1836 if list(filter(p)):
1837 pboost += p
1838 else:
1839 pboost = FourMomentum(pboost)
1840
1841
1842 pboost.px *=-1
1843 pboost.py *=-1
1844 pboost.pz *=-1
1845 for p in self:
1846 b= FourMomentum(p).boost(pboost)
1847 p.E, p.px, p.py, p.pz = b.E, b.px, b.py, b.pz
1848
1849 return self
1850
1852 """check various property of the events"""
1853
1854
1855 threshold = 1e-6
1856
1857
1858 E, px, py, pz = 0,0,0,0
1859 absE, abspx, abspy, abspz = 0,0,0,0
1860 for particle in self:
1861 coeff = 1
1862 if particle.status == -1:
1863 coeff = -1
1864 elif particle.status != 1:
1865 continue
1866 E += coeff * particle.E
1867 absE += abs(particle.E)
1868 px += coeff * particle.px
1869 py += coeff * particle.py
1870 pz += coeff * particle.pz
1871 abspx += abs(particle.px)
1872 abspy += abs(particle.py)
1873 abspz += abs(particle.pz)
1874
1875 fourmass = FourMomentum(particle).mass
1876
1877 if particle.mass and (abs(particle.mass) - fourmass)/ abs(particle.mass) > threshold:
1878 raise Exception( "Do not have correct mass lhe: %s momentum: %s (error at %s" % (particle.mass, fourmass, (abs(particle.mass) - fourmass)/ abs(particle.mass)))
1879
1880
1881 if E/absE > threshold:
1882 logger.critical(self)
1883 raise Exception("Do not conserve Energy %s, %s" % (E/absE, E))
1884 if px/abspx > threshold:
1885 logger.critical(self)
1886 raise Exception("Do not conserve Px %s, %s" % (px/abspx, px))
1887 if py/abspy > threshold:
1888 logger.critical(self)
1889 raise Exception("Do not conserve Py %s, %s" % (py/abspy, py))
1890 if pz/abspz > threshold:
1891 logger.critical(self)
1892 raise Exception("Do not conserve Pz %s, %s" % (pz/abspz, pz))
1893
1894
1895 self.check_color_structure()
1896
1897
1898
1899
1901 """read the line corresponding to global event line
1902 format of the line is:
1903 Nexternal IEVENT WEIGHT SCALE AEW AS
1904 """
1905 inputs = line.split()
1906 assert len(inputs) == 6
1907 self.nexternal=int(inputs[0])
1908 self.ievent=int(inputs[1])
1909 self.wgt=float(inputs[2])
1910 self.scale=float(inputs[3])
1911 self.aqed=float(inputs[4])
1912 self.aqcd=float(inputs[5])
1913
1915 """Return the unique tag identifying the SubProcesses for the generation.
1916 Usefull for program like MadSpin and Reweight module."""
1917
1918 initial, final, order = [], [], [[], []]
1919 for particle in self:
1920 if particle.status == -1:
1921 initial.append(particle.pid)
1922 order[0].append(particle.pid)
1923 elif particle.status == 1:
1924 final.append(particle.pid)
1925 order[1].append(particle.pid)
1926 initial.sort(), final.sort()
1927 tag = (tuple(initial), tuple(final))
1928 return tag, order
1929
1930 @staticmethod
1931 - def mass_shuffle(momenta, sqrts, new_mass, new_sqrts=None):
1932 """use the RAMBO method to shuffle the PS. initial sqrts is preserved."""
1933
1934 if not new_sqrts:
1935 new_sqrts = sqrts
1936
1937 oldm = [p.mass_sqr for p in momenta]
1938 newm = [m**2 for m in new_mass]
1939 tot_mom = sum(momenta, FourMomentum())
1940 if tot_mom.pt2 > 1e-5:
1941 boost_back = FourMomentum(tot_mom.mass,0,0,0).boost_to_restframe(tot_mom)
1942 for i,m in enumerate(momenta):
1943 momenta[i] = m.boost_to_restframe(tot_mom)
1944
1945
1946 f = lambda chi: new_sqrts - sum(math.sqrt(max(0, M + chi**2*(p.E**2-m)))
1947 for M,p,m in zip(newm, momenta,oldm))
1948
1949 df = lambda chi: -1* sum(chi*(p.E**2-m)/math.sqrt(max(0,(p.E**2-m)*chi**2+M))
1950 for M,p,m in zip(newm, momenta,oldm))
1951
1952 if sum(new_mass) > new_sqrts:
1953 return momenta, 0
1954 try:
1955 chi = misc.newtonmethod(f, df, 1.0, error=1e-7,maxiter=1000)
1956 except:
1957 return momenta, 0
1958
1959 new_momenta = []
1960 for i,p in enumerate(momenta):
1961 new_momenta.append(
1962 FourMomentum(math.sqrt(newm[i]+chi**2*(p.E**2-oldm[i])),
1963 chi*p.px, chi*p.py, chi*p.pz))
1964
1965
1966
1967
1968
1969
1970
1971 jac = chi**(3*len(momenta)-3)
1972 jac *= reduce(operator.mul,[p.E/k.E for p,k in zip(momenta, new_momenta)],1)
1973 jac *= sum(p.norm_sq/p.E for p in momenta)
1974 jac /= sum(k.norm_sq/k.E for k in new_momenta)
1975
1976
1977 if tot_mom.pt2 > 1e-5:
1978 for i,m in enumerate(new_momenta):
1979 new_momenta[i] = m.boost_to_restframe(boost_back)
1980 return new_momenta, jac
1981
1982
1983
1984
1986 """routine to rescale the mass via RAMBO method. no internal mass preserve.
1987 sqrts is preserve (RAMBO algo)
1988 """
1989
1990 old_momenta = []
1991 new_masses = []
1992 change_mass = False
1993 for part in self:
1994 if part.status == 1:
1995 old_momenta.append(FourMomentum(part))
1996 new_masses.append(new_param_card.get_value('mass', abs(part.pid)))
1997 if not misc.equal(part.mass, new_masses[-1], 4, zero_limit=10):
1998 change_mass = True
1999
2000 if not change_mass:
2001 return 1
2002
2003 sqrts = self.sqrts
2004
2005
2006 new_mom, jac = self.mass_shuffle(old_momenta, sqrts, new_masses)
2007
2008
2009 ind =0
2010 for part in self:
2011 if part.status==1:
2012 part.E, part.px, part.py, part.pz, part.mass = \
2013 new_mom[ind].E, new_mom[ind].px, new_mom[ind].py, new_mom[ind].pz,new_mom[ind].mass
2014 ind+=1
2015 return jac
2016
2018 """routine to rescale the momenta to change the invariant mass"""
2019
2020 old_momenta = []
2021 incoming = []
2022 masses = []
2023 for part in self:
2024 if part.status == -1:
2025 incoming.append(FourMomentum(part))
2026 if part.status == 1:
2027 old_momenta.append(FourMomentum(part))
2028 masses.append(part.mass)
2029
2030 p_init = FourMomentum()
2031 p_inits = []
2032 n_init = 0
2033 for p in incoming:
2034 n_init +=1
2035 p_init += p
2036 p_inits.append(p)
2037 old_sqrts = p_init.mass
2038
2039 new_mom, jac = self.mass_shuffle(old_momenta, old_sqrts, masses, new_sqrts=new_sqrts)
2040
2041
2042 ind =0
2043 for part in self:
2044 if part.status==1:
2045 part.E, part.px, part.py, part.pz, part.mass = \
2046 new_mom[ind].E, new_mom[ind].px, new_mom[ind].py, new_mom[ind].pz,new_mom[ind].mass
2047 ind+=1
2048
2049
2050 p_init = FourMomentum()
2051 for part in self:
2052 if part.status==1:
2053 p_init += part
2054 if n_init == 1:
2055 for part in self:
2056 if part.status == -1:
2057 part.E, part.px, part.py, part.pz = \
2058 p_init.E, p_init.px, p_init.py, p_init.pz
2059 elif n_init ==2:
2060 if not misc.equal(p_init.px, 0) or not misc.equal(p_init.py, 0):
2061 raise Exception
2062 if not misc.equal(p_inits[0].px, 0) or not misc.equal(p_inits[0].py, 0):
2063 raise Exception
2064
2065
2066
2067
2068 eta = math.log(2*p_inits[0].E/old_sqrts)
2069 new_p = [[new_sqrts/2*math.exp(eta), 0., 0., new_sqrts/2*math.exp(eta)],
2070 [new_sqrts/2*math.exp(-eta), 0., 0., -new_sqrts/2*math.exp(-eta)]]
2071
2072 ind=0
2073 for part in self:
2074 if part.status == -1:
2075 part.E, part.px, part.py, part.pz = new_p[ind]
2076 ind+=1
2077 if ind ==2:
2078 break
2079 else:
2080 raise Exception
2081
2082 return jac
2083
2084
2086 """return a list with the helicities in the order asked for"""
2087
2088
2089 order = [list(get_order[0]), list(get_order[1])]
2090 out = [9] *(len(order[0])+len(order[1]))
2091 for i, part in enumerate(self):
2092 if part.status == 1:
2093 try:
2094 ind = order[1].index(part.pid)
2095 except ValueError as error:
2096 if not allow_reversed:
2097 raise error
2098 else:
2099 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2100 try:
2101 return self.get_helicity(order, False)
2102 except ValueError:
2103 raise error
2104 position = len(order[0]) + ind
2105 order[1][ind] = 0
2106 elif part.status == -1:
2107 try:
2108 ind = order[0].index(part.pid)
2109 except ValueError as error:
2110 if not allow_reversed:
2111 raise error
2112 else:
2113 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2114 try:
2115 return self.get_helicity(order, False)
2116 except ValueError:
2117 raise error
2118
2119 position = ind
2120 order[0][ind] = 0
2121 else:
2122 continue
2123 out[position] = int(part.helicity)
2124 return out
2125
2126
2128 """check the validity of the color structure"""
2129
2130
2131 color_index = collections.defaultdict(int)
2132 for particle in self:
2133 if particle.status in [-1,1]:
2134 if particle.color1:
2135 color_index[particle.color1] +=1
2136 if -7 < particle.pdg < 0:
2137 raise Exception("anti-quark with color tag")
2138 if particle.color2:
2139 color_index[particle.color2] +=1
2140 if 7 > particle.pdg > 0:
2141 raise Exception("quark with anti-color tag")
2142
2143
2144 for key,value in color_index.items():
2145 if value > 2:
2146 print(self)
2147 print(key, value)
2148 raise Exception('Wrong color_flow')
2149
2150
2151
2152 check = []
2153 popup_index = []
2154 for particle in self:
2155 mothers = []
2156 childs = []
2157 if particle.mother1:
2158 mothers.append(particle.mother1)
2159 if particle.mother2 and particle.mother2 is not particle.mother1:
2160 mothers.append(particle.mother2)
2161 if not mothers:
2162 continue
2163 if (particle.mother1.event_id, particle.mother2.event_id) in check:
2164 continue
2165 check.append((particle.mother1.event_id, particle.mother2.event_id))
2166
2167 childs = [p for p in self if p.mother1 is particle.mother1 and \
2168 p.mother2 is particle.mother2]
2169
2170 mcolors = []
2171 manticolors = []
2172 for m in mothers:
2173 if m.color1:
2174 if m.color1 in manticolors:
2175 manticolors.remove(m.color1)
2176 else:
2177 mcolors.append(m.color1)
2178 if m.color2:
2179 if m.color2 in mcolors:
2180 mcolors.remove(m.color2)
2181 else:
2182 manticolors.append(m.color2)
2183 ccolors = []
2184 canticolors = []
2185 for m in childs:
2186 if m.color1:
2187 if m.color1 in canticolors:
2188 canticolors.remove(m.color1)
2189 else:
2190 ccolors.append(m.color1)
2191 if m.color2:
2192 if m.color2 in ccolors:
2193 ccolors.remove(m.color2)
2194 else:
2195 canticolors.append(m.color2)
2196 for index in mcolors[:]:
2197 if index in ccolors:
2198 mcolors.remove(index)
2199 ccolors.remove(index)
2200 for index in manticolors[:]:
2201 if index in canticolors:
2202 manticolors.remove(index)
2203 canticolors.remove(index)
2204
2205 if mcolors != []:
2206
2207 if len(canticolors) + len(mcolors) != 3:
2208 logger.critical(str(self))
2209 raise Exception("Wrong color flow for %s -> %s" % ([m.pid for m in mothers], [c.pid for c in childs]))
2210 else:
2211 popup_index += canticolors
2212 elif manticolors != []:
2213
2214 if len(ccolors) + len(manticolors) != 3:
2215 logger.critical(str(self))
2216 raise Exception("Wrong color flow for %s -> %s" % ([m.pid for m in mothers], [c.pid for c in childs]))
2217 else:
2218 popup_index += ccolors
2219
2220
2221 if len(popup_index) != len(set(popup_index)):
2222 logger.critical(self)
2223 raise Exception("Wrong color flow: identical poping-up index, %s" % (popup_index))
2224
2226 """two event are the same if they have the same momentum. other info are ignored"""
2227
2228 if other is None:
2229 return False
2230
2231 for i,p in enumerate(self):
2232 if p.E != other[i].E:
2233 return False
2234 elif p.pz != other[i].pz:
2235 return False
2236 elif p.px != other[i].px:
2237 return False
2238 elif p.py != other[i].py:
2239 return False
2240 return True
2241
2242
2244 """return a correctly formatted LHE event"""
2245
2246 out="""<event%(event_flag)s>
2247 %(scale)s
2248 %(particles)s
2249 %(comments)s
2250 %(tag)s
2251 %(reweight)s
2252 </event>
2253 """
2254 if event_id not in ['', None]:
2255 self.eventflag['event'] = str(event_id)
2256
2257 if self.eventflag:
2258 event_flag = ' %s' % ' '.join('%s="%s"' % (k,v) for (k,v) in self.eventflag.items())
2259 else:
2260 event_flag = ''
2261
2262 scale_str = "%2d %6d %+13.7e %14.8e %14.8e %14.8e" % \
2263 (self.nexternal,self.ievent,self.wgt,self.scale,self.aqed,self.aqcd)
2264
2265
2266 if self.reweight_data:
2267
2268 if set(self.reweight_data.keys()) != set(self.reweight_order):
2269 self.reweight_order += [k for k in self.reweight_data.keys() \
2270 if k not in self.reweight_order]
2271
2272 reweight_str = '<rwgt>\n%s\n</rwgt>' % '\n'.join(
2273 '<wgt id=\'%s\'> %+13.7e </wgt>' % (i, float(self.reweight_data[i]))
2274 for i in self.reweight_order if i in self.reweight_data)
2275 else:
2276 reweight_str = ''
2277
2278 tag_str = self.tag
2279 if hasattr(self, 'nloweight') and self.nloweight.modified:
2280 self.rewrite_nlo_weight()
2281 tag_str = self.tag
2282
2283 if self.matched_scale_data:
2284 tmp_scale = ' '.join(['pt_clust_%i=\"%s\"' % (i+1,v)
2285 for i,v in enumerate(self.matched_scale_data)
2286 if v!=-1])
2287 if tmp_scale:
2288 tag_str = "<scales %s></scales>%s" % (tmp_scale, self.tag)
2289
2290 if self.syscalc_data:
2291 keys= ['rscale', 'asrwt', ('pdfrwt', 'beam', '1'), ('pdfrwt', 'beam', '2'),
2292 'matchscale', 'totfact']
2293 sys_str = "<mgrwt>\n"
2294 template = """<%(key)s%(opts)s>%(values)s</%(key)s>\n"""
2295 for k in keys:
2296 if k not in self.syscalc_data:
2297 continue
2298 replace = {}
2299 replace['values'] = self.syscalc_data[k]
2300 if isinstance(k, (str, six.text_type)):
2301 replace['key'] = k
2302 replace['opts'] = ''
2303 else:
2304 replace['key'] = k[0]
2305 replace['opts'] = ' %s=\"%s\"' % (k[1],k[2])
2306 sys_str += template % replace
2307 sys_str += "</mgrwt>\n"
2308 reweight_str = sys_str + reweight_str
2309
2310 out = out % {'event_flag': event_flag,
2311 'scale': scale_str,
2312 'particles': '\n'.join([str(p) for p in self]),
2313 'tag': tag_str,
2314 'comments': self.comment,
2315 'reweight': reweight_str}
2316
2317 return re.sub('[\n]+', '\n', out)
2318
2319 - def get_momenta(self, get_order, allow_reversed=True):
2320 """return the momenta vector in the order asked for"""
2321
2322
2323 order = [list(get_order[0]), list(get_order[1])]
2324 out = [''] *(len(order[0])+len(order[1]))
2325 for i, part in enumerate(self):
2326 if part.status == 1:
2327 try:
2328 ind = order[1].index(part.pid)
2329 except ValueError as error:
2330 if not allow_reversed:
2331 raise error
2332 else:
2333 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2334 try:
2335 return self.get_momenta_str(order, False)
2336 except ValueError:
2337 raise error
2338 position = len(order[0]) + ind
2339 order[1][ind] = 0
2340 elif part.status == -1:
2341 try:
2342 ind = order[0].index(part.pid)
2343 except ValueError as error:
2344 if not allow_reversed:
2345 raise error
2346 else:
2347 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2348 try:
2349 return self.get_momenta_str(order, False)
2350 except ValueError:
2351 raise error
2352
2353 position = ind
2354 order[0][ind] = 0
2355 else:
2356 continue
2357
2358 out[position] = (part.E, part.px, part.py, part.pz)
2359
2360 return out
2361
2362
2375
2376
2378
2379 scale = 0
2380 for particle in self:
2381 if particle.status != 1:
2382 continue
2383 p=FourMomentum(particle)
2384 scale += math.sqrt(p.mass_sqr + p.pt**2)
2385
2386 return prefactor * scale
2387
2388
2390
2391 scale = 0
2392 for particle in self:
2393 if particle.status != 1:
2394 continue
2395 p = FourMomentum(particle)
2396 pt = p.pt
2397 if (pt>0):
2398 scale += p.E*pt/math.sqrt(pt**2+p.pz**2)
2399
2400 return prefactor * scale
2401
2402 @property
2405
2407
2408 scale = 0
2409 init = []
2410 for particle in self:
2411 if particle.status == -1:
2412 init.append(FourMomentum(particle))
2413 if len(init) == 1:
2414 return init[0].mass
2415 elif len(init)==2:
2416 return math.sqrt((init[0]+init[1])**2)
2417
2418
2419
2420
2422 """return the momenta str in the order asked for"""
2423
2424 out = self.get_momenta(get_order, allow_reversed)
2425
2426 format = '%.12f'
2427 format_line = ' '.join([format]*4) + ' \n'
2428 out = [format_line % one for one in out]
2429 out = ''.join(out).replace('e','d')
2430 return out
2431
2433 """a convenient object for 4-momenta operation"""
2434
2435 - def __init__(self, obj=0, px=0, py=0, pz=0, E=0):
2436 """initialize the four momenta"""
2437
2438 if obj == 0 and E:
2439 obj = E
2440
2441 if isinstance(obj, (FourMomentum, Particle)):
2442 px = obj.px
2443 py = obj.py
2444 pz = obj.pz
2445 E = obj.E
2446 elif isinstance(obj, (list, tuple)):
2447 assert len(obj) ==4
2448 E = obj[0]
2449 px = obj[1]
2450 py = obj[2]
2451 pz = obj[3]
2452 elif isinstance(obj, (str, six.text_type)):
2453 obj = [float(i) for i in obj.split()]
2454 assert len(obj) ==4
2455 E = obj[0]
2456 px = obj[1]
2457 py = obj[2]
2458 pz = obj[3]
2459 else:
2460 E =obj
2461
2462
2463 self.E = float(E)
2464 self.px = float(px)
2465 self.py = float(py)
2466 self.pz = float(pz)
2467
2468 @property
2470 """return the mass"""
2471 return math.sqrt(max(self.E**2 - self.px**2 - self.py**2 - self.pz**2,0))
2472
2473 @property
2475 """return the mass square"""
2476 return self.E**2 - self.px**2 - self.py**2 - self.pz**2
2477
2478 @property
2480 return math.sqrt(max(0, self.pt2))
2481
2482 @property
2484 norm = math.sqrt(self.px**2 + self.py**2+self.pz**2)
2485 return 0.5* math.log((norm - self.pz) / (norm + self.pz))
2486
2487 @property
2489 return 0.5* math.log((self.E +self.pz) / (self.E - self.pz))
2490
2491
2492 @property
2494 """ return the pt square """
2495
2496 return self.px**2 + self.py**2
2497
2498 @property
2500 """ return |\vec p| """
2501 return math.sqrt(self.px**2 + self.py**2 + self.pz**2)
2502
2503 @property
2505 """ return |\vec p|^2 """
2506 return self.px**2 + self.py**2 + self.pz**2
2507
2508 @property
2510 """return the mass square"""
2511 import math
2512 return math.atan(math.sqrt((self.px**2+self.py**2)/self.pz**2))
2513
2514
2516
2517 assert isinstance(obj, FourMomentum)
2518 new = FourMomentum(self.E+obj.E,
2519 self.px + obj.px,
2520 self.py + obj.py,
2521 self.pz + obj.pz)
2522 return new
2523
2525 """update the object with the sum"""
2526 self.E += obj.E
2527 self.px += obj.px
2528 self.py += obj.py
2529 self.pz += obj.pz
2530 return self
2531
2533
2534 assert isinstance(obj, FourMomentum)
2535 new = FourMomentum(self.E-obj.E,
2536 self.px - obj.px,
2537 self.py - obj.py,
2538 self.pz - obj.pz)
2539 return new
2540
2542 """update the object with the sum"""
2543 self.E -= obj.E
2544 self.px -= obj.px
2545 self.py -= obj.py
2546 self.pz -= obj.pz
2547 return self
2548
2550 if isinstance(obj, FourMomentum):
2551 return self.E*obj.E - self.px *obj.px - self.py * obj.py - self.pz * obj.pz
2552 elif isinstance(obj, (float, int)):
2553 return FourMomentum(obj*self.E,obj*self.px,obj*self.py,obj*self.pz )
2554 else:
2555 raise NotImplemented
2556 __rmul__ = __mul__
2557
2559 assert power in [1,2]
2560
2561 if power == 1:
2562 return FourMomentum(self)
2563 elif power == 2:
2564 return self.mass_sqr
2565
2567 return 'FourMomentum(%s,%s,%s,%s)' % (self.E, self.px, self.py,self.pz)
2568
2569 - def __str__(self, mode='python'):
2570 if mode == 'python':
2571 return self.__repr__()
2572 elif mode == 'fortran':
2573 return '%.10e %.10e %.10e %.10e' % self.get_tuple()
2574
2576 return (self.E, self.px, self.py,self.pz)
2577
2579 """mom 4-momenta is suppose to be given in the rest frame of this 4-momenta.
2580 the output is the 4-momenta in the frame of this 4-momenta
2581 function copied from HELAS routine.
2582 if the current momenta is (E,\vec{p}), in order to go to the rest frame
2583 of the current particle, mom should be (E, -\vec{p})
2584 """
2585
2586 pnorm = mom.px**2 + mom.py**2 + mom.pz**2
2587 if pnorm:
2588 s3product = self.px * mom.px + self.py * mom.py + self.pz * mom.pz
2589 mass = mom.mass
2590 lf = (self.E + (mom.E - mass) * s3product / pnorm ) / mass
2591 return FourMomentum(E=(self.E*mom.E+s3product)/mass,
2592 px=self.px + mom.px * lf,
2593 py=self.py + mom.py * lf,
2594 pz=self.pz + mom.pz * lf)
2595 else:
2596 return FourMomentum(mom)
2597
2598 - def zboost(self, pboost=None, E=0, pz=0):
2599 """Both momenta should be in the same frame.
2600 The boost perform correspond to the boost required to set pboost at
2601 rest (only z boost applied).
2602 """
2603 if isinstance(pboost, FourMomentum):
2604 E = pboost.E
2605 pz = pboost.pz
2606
2607
2608 gamma = E / math.sqrt(E**2-pz**2)
2609 gammabeta = pz / math.sqrt(E**2-pz**2)
2610
2611 out = FourMomentum([gamma*self.E - gammabeta*self.pz,
2612 self.px,
2613 self.py,
2614 gamma*self.pz - gammabeta*self.E])
2615
2616 if abs(out.pz) < 1e-6 * out.E:
2617 out.pz = 0
2618 return out
2619
2621 """apply the boost transformation such that pboost is at rest in the new frame.
2622 First apply a rotation to allign the pboost to the z axis and then use
2623 zboost routine (see above)
2624 """
2625
2626 if pboost.px == 0 == pboost.py:
2627 out = self.zboost(E=pboost.E,pz=pboost.pz)
2628 return out
2629
2630
2631
2632
2633
2634
2635
2636 p = math.sqrt( pboost.px**2 + pboost.py**2+ pboost.pz**2)
2637 cosF = pboost.pz / p
2638 sinF = math.sqrt(1-cosF**2)
2639 sinT = pboost.py/p/sinF
2640 cosT = pboost.px/p/sinF
2641
2642 out=FourMomentum([self.E,
2643 self.px*cosT*cosF + self.py*sinT*cosF-self.pz*sinF,
2644 -self.px*sinT+ self.py*cosT,
2645 self.px*cosT*sinF + self.py*sinT*sinF + self.pz*cosF
2646 ])
2647 out = out.zboost(E=pboost.E,pz=p)
2648 return out
2649
2654
2655 - def __init__(self, input, real_type=(1,11)):
2656 """ """
2657
2658 self.real_type = real_type
2659 if isinstance(input, (str, six.text_type)):
2660 self.parse(input)
2661
2662 - def __str__(self, mode='display'):
2663
2664 if mode == 'display':
2665 out = """ pwgt: %(pwgt)s
2666 born, real : %(born)s %(real)s
2667 pdgs : %(pdgs)s
2668 bjks : %(bjks)s
2669 scales**2, gs: %(scales2)s %(gs)s
2670 born/real related : %(born_related)s %(real_related)s
2671 type / nfks : %(type)s %(nfks)s
2672 to merge : %(to_merge_pdg)s in %(merge_new_pdg)s
2673 ref_wgt : %(ref_wgt)s""" % self.__dict__
2674 return out
2675 elif mode == 'formatted':
2676 format_var = []
2677 variable = []
2678
2679 def to_add_full(f, v, format_var, variable):
2680 """ function to add to the formatted output"""
2681 if isinstance(v, list):
2682 format_var += [f]*len(v)
2683 variable += v
2684 else:
2685 format_var.append(f)
2686 variable.append(v)
2687 to_add = lambda x,y: to_add_full(x,y, format_var, variable)
2688
2689 to_add('%.10e', [p*self.bias_wgt for p in self.pwgt])
2690 to_add('%.10e', self.born)
2691 to_add('%.10e', self.real)
2692 to_add('%i', self.nexternal)
2693 to_add('%i', self.pdgs)
2694 to_add('%i', self.qcdpower)
2695 to_add('%.10e', self.bjks)
2696 to_add('%.10e', self.scales2)
2697 to_add('%.10e', self.gs)
2698 to_add('%i', [self.born_related, self.real_related])
2699 to_add('%i' , [self.type, self.nfks])
2700 to_add('%i' , self.to_merge_pdg)
2701 to_add('%i', self.merge_new_pdg)
2702 to_add('%.10e', self.ref_wgt*self.bias_wgt)
2703 to_add('%.10e', self.bias_wgt)
2704 return ' '.join(format_var) % tuple(variable)
2705
2706
2707 - def parse(self, text, keep_bias=False):
2708 """parse the line and create the related object.
2709 keep bias allow to not systematically correct for the bias in the written information"""
2710
2711
2712
2713 data = text.split()
2714
2715
2716
2717
2718 self.pwgt = [float(f) for f in data[:3]]
2719
2720
2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735
2736
2737 self.born = float(data[3])
2738 self.real = float(data[4])
2739
2740
2741 self.nexternal = int(data[5])
2742
2743
2744 self.pdgs = [int(i) for i in data[6:6+self.nexternal]]
2745 flag = 6+self.nexternal
2746
2747
2748 self.qcdpower = int(data[flag])
2749
2750
2751 self.bjks = [float(f) for f in data[flag+1:flag+3]]
2752
2753
2754 self.scales2 = [float(f) for f in data[flag+3:flag+6]]
2755
2756
2757 self.gs = float(data[flag+6])
2758
2759
2760
2761
2762
2763
2764
2765 self.born_related = int(data[flag+7])
2766 self.real_related = int(data[flag+8])
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
2780
2781
2782
2783
2784
2785
2786
2787
2788
2789
2790 self.type = int(data[flag+9])
2791
2792
2793
2794
2795 self.nfks = int(data[flag+10])
2796
2797
2798
2799
2800
2801
2802
2803 self.to_merge_pdg = [int (f) for f in data[flag+11:flag+13]]
2804
2805
2806 self.merge_new_pdg = int(data[flag+13])
2807
2808
2809
2810
2811
2812 self.ref_wgt = float(data[flag+14])
2813
2814
2815
2816 try:
2817 self.bias_wgt = float(data[flag+15])
2818 except IndexError:
2819 self.bias_wgt = 1.0
2820
2821 if not keep_bias:
2822 self.ref_wgt /= self.bias_wgt
2823 self.pwgt = [p/self.bias_wgt for p in self.pwgt]
2824
2825
2826 if self.type in self.real_type:
2827 self.momenta_config = self.real_related
2828 else:
2829 self.momenta_config = self.born_related
2830
2833
2835
2836
2837 - def __init__(self, momenta, wgts, event, real_type=(1,11)):
2838
2839 list.__init__(self, momenta)
2840 assert self
2841 self.soft = False
2842 self.wgts = wgts
2843 self.pdgs = list(wgts[0].pdgs)
2844 self.event = event
2845 self.real_type = real_type
2846
2847 if wgts[0].momenta_config == wgts[0].born_related:
2848
2849 ind1, ind2 = [ind-1 for ind in wgts[0].to_merge_pdg]
2850 if ind1> ind2:
2851 ind1, ind2 = ind2, ind1
2852 if ind1 >= sum(1 for p in event if p.status==-1):
2853 new_p = self[ind1] + self[ind2]
2854 else:
2855 new_p = self[ind1] - self[ind2]
2856 self.pop(ind1)
2857 self.insert(ind1, new_p)
2858 self.pop(ind2)
2859 self.pdgs.pop(ind1)
2860 self.pdgs.insert(ind1, wgts[0].merge_new_pdg )
2861 self.pdgs.pop(ind2)
2862
2863
2864 elif any(w.type in self.real_type for w in wgts):
2865 if any(w.type not in self.real_type for w in wgts):
2866 raise Exception
2867
2868
2869
2870
2871
2872
2873
2874
2875
2876
2877
2878
2879 else:
2880 raise Exception
2881
2883 """check that the propagator associated to ij is not too light
2884 [related to soft-collinear singularity]"""
2885
2886 if threshold is None:
2887 threshold = 1e-8
2888
2889 if ind1> ind2:
2890 ind1, ind2 = ind2, ind1
2891 if ind1 >= nb_init:
2892 new_p = self[ind1] + self[ind2]
2893 else:
2894 new_p = self[ind1] - self[ind2]
2895
2896 inv_mass = new_p.mass_sqr
2897 if nb_init == 2:
2898 shat = (self[0]+self[1]).mass_sqr
2899 else:
2900 shat = self[0].mass_sqr
2901
2902
2903 if (abs(inv_mass)/shat < threshold):
2904 return True
2905 else:
2906 return False
2907
2908
2911
2913 """ return the tag and order for this basic event"""
2914 (initial, _), _ = self.event.get_tag_and_order()
2915 order = self.get_pdg_code()
2916
2917
2918 initial, out = order[:len(initial)], order[len(initial):]
2919 initial.sort()
2920 out.sort()
2921 return (tuple(initial), tuple(out)), order
2922
2923 - def get_momenta(self, get_order, allow_reversed=True):
2924 """return the momenta vector in the order asked for"""
2925
2926
2927 order = [list(get_order[0]), list(get_order[1])]
2928 out = [''] *(len(order[0])+len(order[1]))
2929 pdgs = self.get_pdg_code()
2930 for pos, part in enumerate(self):
2931 if pos < len(get_order[0]):
2932 try:
2933 ind = order[0].index(pdgs[pos])
2934 except ValueError as error:
2935 if not allow_reversed:
2936 raise error
2937 else:
2938 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2939 try:
2940 return self.get_momenta(order, False)
2941 except ValueError:
2942 raise error
2943
2944
2945 position = ind
2946 order[0][ind] = 0
2947 else:
2948 try:
2949 ind = order[1].index(pdgs[pos])
2950 except ValueError as error:
2951 if not allow_reversed:
2952 raise error
2953 else:
2954 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2955 try:
2956 return self.get_momenta(order, False)
2957 except ValueError:
2958 raise error
2959 position = len(order[0]) + ind
2960 order[1][ind] = 0
2961
2962 out[position] = (part.E, part.px, part.py, part.pz)
2963
2964 return out
2965
2966
2968 return [9] * len(self)
2969
2970 @property
2972 return self.event.aqcd
2973
2975
2976 scale = 0
2977 for particle in self:
2978 p = particle
2979 scale += math.sqrt(max(0, p.mass_sqr + p.pt**2))
2980
2981 return prefactor * scale
2982
2984
2985 scale = 0
2986 for particle in self:
2987 p = particle
2988 pt = p.pt
2989 if (pt>0):
2990 scale += p.E*pt/math.sqrt(pt**2+p.pz**2)
2991
2992 return prefactor * scale
2993
2994
2996
2997 scale = 0
2998 nb_init = 0
2999 for particle in event:
3000 if particle.status == -1:
3001 nb_init+=1
3002 if nb_init == 1:
3003 return self[0].mass
3004 elif nb_init==2:
3005 return math.sqrt((self[0]+self[1])**2)
3006
3007
3008
3009
3010 - def __init__(self, input, event, real_type=(1,11), threshold=None):
3011
3012 self.real_type = real_type
3013 self.event = event
3014 self.total_wgt = 0.
3015 self.nb_event = 0
3016 self.nb_wgts = 0
3017 self.threshold = threshold
3018 self.modified = False
3019
3020
3021 if isinstance(input, (str,six.text_type)):
3022 self.parse(input)
3023
3024
3025
3027 """create the object from the string information (see example below)"""
3028
3029
3030
3031
3032
3033
3034
3035
3036
3037
3038
3039
3040
3041
3042
3043
3044
3045
3046
3047
3048
3049 text = text.lower().replace('d','e')
3050 all_line = text.split('\n')
3051
3052 first_line =''
3053 while not first_line.strip():
3054 first_line = all_line.pop(0)
3055
3056 wgt, nb_wgt, nb_event, _ = first_line.split()
3057 self.total_wgt = float(wgt.replace('d','e'))
3058 nb_wgt, nb_event = int(nb_wgt), int(nb_event)
3059 self.nb_wgt, self.nb_event = nb_wgt, nb_event
3060
3061 momenta = []
3062 self.momenta = momenta
3063 wgts = []
3064 for line in all_line:
3065 data = line.split()
3066 if len(data) == 4:
3067 p = FourMomentum(data)
3068 momenta.append(p)
3069 elif len(data)>0:
3070 wgt = OneNLOWeight(line, real_type=self.real_type)
3071 wgts.append(wgt)
3072
3073 assert len(wgts) == int(nb_wgt)
3074
3075 get_weights_for_momenta = dict( (i,[]) for i in range(1,nb_event+1) )
3076 size_momenta = 0
3077 for wgt in wgts:
3078 if wgt.momenta_config in get_weights_for_momenta:
3079 get_weights_for_momenta[wgt.momenta_config].append(wgt)
3080 else:
3081 if size_momenta == 0: size_momenta = wgt.nexternal
3082 assert size_momenta == wgt.nexternal
3083 get_weights_for_momenta[wgt.momenta_config] = [wgt]
3084
3085 assert sum(len(c) for c in get_weights_for_momenta.values()) == int(nb_wgt), "%s != %s" % (sum(len(c) for c in get_weights_for_momenta.values()), nb_wgt)
3086
3087
3088 for key in range(1, nb_event+1):
3089 wgts = get_weights_for_momenta[key]
3090 if not wgts:
3091 continue
3092 if size_momenta == 0: size_momenta = wgts[0].nexternal
3093 p = momenta[size_momenta*(key-1):key*size_momenta]
3094 evt = self.BasicEvent(p, wgts, self.event, self.real_type)
3095 if len(evt) == size_momenta:
3096 for wgt in wgts:
3097 if not wgt.type in self.real_type:
3098 continue
3099 if evt.check_fks_singularity(wgt.to_merge_pdg[0]-1,
3100 wgt.to_merge_pdg[1]-1,
3101 nb_init=sum(1 for p in self.event if p.status==-1),
3102 threshold=self.threshold):
3103 get_weights_for_momenta[wgt.momenta_config].remove(wgt)
3104 get_weights_for_momenta[wgt.born_related].append(wgt)
3105 wgt.momenta_config = wgt.born_related
3106
3107 assert sum(len(c) for c in get_weights_for_momenta.values()) == int(nb_wgt), "%s != %s" % (sum(len(c) for c in get_weights_for_momenta.values()), nb_wgt)
3108
3109 self.cevents = []
3110 for key in range(1, nb_event+1):
3111 if key in get_weights_for_momenta:
3112 wgt = get_weights_for_momenta[key]
3113 if not wgt:
3114 continue
3115 pdg_to_event = {}
3116 for w in wgt:
3117 pdgs = w.pdgs
3118 if w.momenta_config == w.born_related:
3119 pdgs = list(pdgs)
3120 ind1, ind2 = [ind-1 for ind in w.to_merge_pdg]
3121 if ind1> ind2:
3122 ind1, ind2 = ind2, ind1
3123 pdgs.pop(ind1)
3124 pdgs.insert(ind1, w.merge_new_pdg )
3125 pdgs.pop(ind2)
3126 pdgs = tuple(pdgs)
3127 if pdgs not in pdg_to_event:
3128 p = momenta[size_momenta*(key-1):key*size_momenta]
3129 evt = self.BasicEvent(p, [w], self.event, self.real_type)
3130 self.cevents.append(evt)
3131 pdg_to_event[pdgs] = evt
3132 else:
3133 pdg_to_event[pdgs].wgts.append(w)
3134
3135 if __debug__:
3136 nb_wgt_check = 0
3137 for cevt in self.cevents:
3138 nb_wgt_check += len(cevt.wgts)
3139 assert nb_wgt_check == int(nb_wgt)
3140
3141
3142
3143 if '__main__' == __name__:
3144
3145 if False:
3146 lhe = EventFile('unweighted_events.lhe')
3147
3148 start = time.time()
3149 for event in lhe:
3150 pass
3151 s = time.time()
3152 print(s-start)
3153
3154
3155
3156
3157
3158
3159
3160
3161
3162
3163
3164 if False:
3165 start = time
3166 lhe = EventFile('unweighted_events.lhe.gz')
3167 output = open('output_events.lhe', 'w')
3168
3169 output.write(lhe.banner)
3170
3171 for event in lhe:
3172 for particle in event:
3173
3174 particle.mass = 0
3175 particle.vtim = 2
3176
3177
3178 output.write(str(event))
3179 output.write('</LesHouchesEvent>\n')
3180
3181
3182 if False:
3183 lhe = EventFile('unweighted_events.lhe.gz')
3184 import matplotlib.pyplot as plt
3185 import matplotlib.gridspec as gridspec
3186 nbins = 100
3187
3188 nb_pass = 0
3189 data = []
3190 for event in lhe:
3191 etaabs = 0
3192 etafinal = 0
3193 for particle in event:
3194 if particle.status==1:
3195 p = FourMomentum(particle)
3196 eta = p.pseudorapidity
3197 if abs(eta) > etaabs:
3198 etafinal = eta
3199 etaabs = abs(eta)
3200 if etaabs < 4:
3201 data.append(etafinal)
3202 nb_pass +=1
3203
3204
3205 print(nb_pass)
3206 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1])
3207 gs1.update(wspace=0, hspace=0)
3208 ax = plt.subplot(gs1[0])
3209
3210 n, bins, patches = ax.hist(data, nbins, histtype='step', label='original')
3211 ax_c = ax.twinx()
3212 ax_c.set_ylabel('MadGraph5_aMC@NLO')
3213 ax_c.yaxis.set_label_coords(1.01, 0.25)
3214 ax_c.set_yticks(ax.get_yticks())
3215 ax_c.set_yticklabels([])
3216 ax.set_xlim([-4,4])
3217 print("bin value:", n)
3218 print("start/end point of bins", bins)
3219 plt.axis('on')
3220 plt.xlabel('weight ratio')
3221 plt.show()
3222
3223
3224
3225 if False:
3226 lhe = EventFile('unweighted_events.lhe')
3227 import matplotlib.pyplot as plt
3228 import matplotlib.gridspec as gridspec
3229 nbins = 100
3230
3231
3232 mtau, wtau = 1.777, 4.027000e-13
3233 nb_pass = 0
3234 data, data2, data3 = [], [], []
3235 for event in lhe:
3236 nb_pass +=1
3237 if nb_pass > 10000:
3238 break
3239 tau1 = FourMomentum()
3240 tau2 = FourMomentum()
3241 for part in event:
3242 if part.pid in [-12,11,16]:
3243 momenta = FourMomentum(part)
3244 tau1 += momenta
3245 elif part.pid == 15:
3246 tau2 += FourMomentum(part)
3247
3248 if abs((mtau-tau2.mass())/wtau)<1e6 and tau2.mass() >1:
3249 data.append((tau1.mass()-mtau)/wtau)
3250 data2.append((tau2.mass()-mtau)/wtau)
3251 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1])
3252 gs1.update(wspace=0, hspace=0)
3253 ax = plt.subplot(gs1[0])
3254
3255 n, bins, patches = ax.hist(data2, nbins, histtype='step', label='original')
3256 n2, bins2, patches2 = ax.hist(data, bins=bins, histtype='step',label='reconstructed')
3257 import cmath
3258
3259 breit = lambda m : math.sqrt(4*math.pi)*1/(((m)**2-mtau**2)**2+(mtau*wtau)**2)*wtau
3260
3261 data3 = [breit(mtau + x*wtau)*wtau*16867622.6624*50 for x in bins]
3262
3263 ax.plot(bins, data3,label='breit-wigner')
3264
3265 ax.legend()
3266
3267 ax_c = ax.twinx()
3268 ax_c.set_ylabel('MadGraph5_aMC@NLO')
3269 ax_c.yaxis.set_label_coords(1.01, 0.25)
3270 ax_c.set_yticks(ax.get_yticks())
3271 ax_c.set_yticklabels([])
3272
3273 plt.title('invariant mass of tau LHE/reconstructed')
3274 plt.axis('on')
3275 ax.set_xticklabels([])
3276
3277 ax = plt.subplot(gs1[1])
3278 data4 = [n[i]/(data3[i]) for i in range(nbins)]
3279 ax.plot(bins, data4 + [0] , 'b')
3280 data4 = [n2[i]/(data3[i]) for i in range(nbins)]
3281 ax.plot(bins, data4 + [0] , 'g')
3282 ax.set_ylim([0,2])
3283
3284 tick = ax.get_yticks()
3285 ax.set_yticks(tick[:-1])
3286
3287
3288 plt.axis('on')
3289 plt.xlabel('(M - Mtau)/Wtau')
3290 plt.show()
3291