1 from __future__ import division
2 import collections
3 import random
4 import re
5 import math
6 import time
7 import os
8 import shutil
9
10 pjoin = os.path.join
11
12 if '__main__' == __name__:
13 import sys
14 sys.path.append('../../')
15 import misc
16 import logging
17 import gzip
18 logger = logging.getLogger("madgraph.lhe_parser")
19
20 -class Particle(object):
21 """ """
22 pattern=re.compile(r'''^\s*
23 (?P<pid>-?\d+)\s+ #PID
24 (?P<status>-?\d+)\s+ #status (1 for output particle)
25 (?P<mother1>-?\d+)\s+ #mother
26 (?P<mother2>-?\d+)\s+ #mother
27 (?P<color1>[+-e.\d]*)\s+ #color1
28 (?P<color2>[+-e.\d]*)\s+ #color2
29 (?P<px>[+-e.\d]*)\s+ #px
30 (?P<py>[+-e.\d]*)\s+ #py
31 (?P<pz>[+-e.\d]*)\s+ #pz
32 (?P<E>[+-e.\d]*)\s+ #E
33 (?P<mass>[+-e.\d]*)\s+ #mass
34 (?P<vtim>[+-e.\d]*)\s+ #displace vertex
35 (?P<helicity>[+-e.\d]*)\s* #helicity
36 ($|(?P<comment>\#[\d|D]*)) #comment/end of string
37 ''',66)
38
39
40
41 - def __init__(self, line=None, event=None):
42 """ """
43
44 if isinstance(line, Particle):
45 for key in line.__dict__:
46 setattr(self, key, getattr(line, key))
47 if event:
48 self.event = event
49 return
50
51 self.event = event
52 self.event_id = len(event)
53
54 self.pid = 0
55 self.status = 0
56 self.mother1 = None
57 self.mother2 = None
58 self.color1 = 0
59 self.color2 = None
60 self.px = 0
61 self.py = 0
62 self.pz = 0
63 self.E = 0
64 self.mass = 0
65 self.vtim = 0
66 self.helicity = 9
67 self.rwgt = 0
68 self.comment = ''
69
70 if line:
71 self.parse(line)
72
73 @property
75 "convenient alias"
76 return self.pid
77
78 - def parse(self, line):
79 """parse the line"""
80
81 obj = self.pattern.search(line)
82 if not obj:
83 raise Exception, 'the line\n%s\n is not a valid format for LHE particle' % line
84 for key, value in obj.groupdict().items():
85 if key not in ['comment','pid']:
86 setattr(self, key, float(value))
87 elif key in ['pid', 'mother1', 'mother2']:
88 setattr(self, key, int(value))
89 else:
90 self.comment = value
91
92
93
95 """string representing the particles"""
96 return " %8d %2d %4d %4d %4d %4d %+13.10e %+13.10e %+13.10e %14.10e %14.10e %10.4e %10.4e" \
97 % (self.pid,
98 self.status,
99 self.mother1.event_id+1 if self.mother1 else 0,
100 self.mother2.event_id+1 if self.mother2 else 0,
101 self.color1,
102 self.color2,
103 self.px,
104 self.py,
105 self.pz,
106 self.E,
107 self.mass,
108 self.vtim,
109 self.helicity)
110
111 - def __eq__(self, other):
112
113 if not isinstance(other, Particle):
114 return False
115 if self.pid == other.pid and \
116 self.status == other.status and \
117 self.mother1 == other.mother1 and \
118 self.mother2 == other.mother2 and \
119 self.color1 == other.color1 and \
120 self.color2 == other.color2 and \
121 self.px == other.px and \
122 self.py == other.py and \
123 self.pz == other.pz and \
124 self.E == other.E and \
125 self.mass == other.mass and \
126 self.vtim == other.vtim and \
127 self.helicity == other.helicity:
128 return True
129 return False
130
131 - def set_momentum(self, momentum):
132
133 self.E = momentum.E
134 self.px = momentum.px
135 self.py = momentum.py
136 self.pz = momentum.pz
137
138 - def add_decay(self, decay_event):
139 """associate to this particle the decay in the associate event"""
140
141 return self.event.add_decay_to_particle(self.event_id, decay_event)
142
143 - def __repr__(self):
144 return 'Particle("%s", event=%s)' % (str(self), self.event)
145
148 """A class to allow to read both gzip and not gzip file"""
149
150 - def __new__(self, path, mode='r', *args, **opt):
151
152 if path.endswith(".gz"):
153 try:
154 return gzip.GzipFile.__new__(EventFileGzip, path, mode, *args, **opt)
155 except IOError, error:
156 raise
157 except Exception, error:
158 if mode == 'r':
159 misc.gunzip(path)
160 return file.__new__(EventFileNoGzip, path[:-3], mode, *args, **opt)
161 else:
162 return file.__new__(EventFileNoGzip, path, mode, *args, **opt)
163
164 - def __init__(self, path, mode='r', *args, **opt):
165 """open file and read the banner [if in read mode]"""
166
167 super(EventFile, self).__init__(path, mode, *args, **opt)
168 self.banner = ''
169 if mode == 'r':
170 line = ''
171 while '</init>' not in line.lower():
172 try:
173 line = super(EventFile, self).next()
174 except StopIteration:
175 self.seek(0)
176 self.banner = ''
177 break
178 if "<event" in line.lower():
179 self.seek(0)
180 self.banner = ''
181 break
182
183 self.banner += line
184
194
195 @property
197 """return the cross-section of the file #from the banner"""
198 try:
199 return self._cross
200 except Exception:
201 pass
202
203 onebanner = self.get_banner()
204 self._cross = onebanner.get_cross()
205 return self._cross
206
208 if self.closed:
209 return 0
210 if hasattr(self,"len"):
211 return self.len
212
213 init_pos = self.tell()
214 self.seek(0)
215 nb_event=0
216 for _ in self:
217 nb_event +=1
218 self.len = nb_event
219 self.seek(init_pos)
220 return self.len
221
222
224 """get next event"""
225 text = ''
226 line = ''
227 mode = 0
228 while '</event>' not in line:
229 line = super(EventFile, self).next().lower()
230 if '<event' in line:
231 mode = 1
232 text = ''
233 if mode:
234 text += line
235 return Event(text)
236
238 """ scan once the file to return
239 - the list of the hightest weight (of size trunc_error*NB_EVENT
240 - the cross-section by type of process
241 - the total number of events in the file
242 """
243
244
245
246 self.seek(0)
247 all_wgt = []
248 cross = collections.defaultdict(int)
249 nb_event = 0
250 for event in self:
251 nb_event +=1
252 wgt = get_wgt(event)
253 cross['all'] += wgt
254 cross['abs'] += abs(wgt)
255 cross[event.ievent] += wgt
256 all_wgt.append(abs(wgt))
257
258 if nb_event % 20000 == 0:
259 all_wgt.sort()
260
261 nb_keep = max(20, int(nb_event*trunc_error*15))
262 all_wgt = all_wgt[-nb_keep:]
263
264
265 all_wgt.sort()
266
267 nb_keep = max(20, int(nb_event*trunc_error*10))
268 all_wgt = all_wgt[-nb_keep:]
269 self.seek(0)
270 return all_wgt, cross, nb_event
271
272
273 - def unweight(self, outputpath, get_wgt=None, max_wgt=0, trunc_error=0, event_target=0,
274 log_level=logging.INFO):
275 """unweight the current file according to wgt information wgt.
276 which can either be a fct of the event or a tag in the rwgt list.
277 max_wgt allow to do partial unweighting.
278 trunc_error allow for dynamical partial unweighting
279 event_target reweight for that many event with maximal trunc_error.
280 (stop to write event when target is reached)
281 """
282 if not get_wgt:
283 def weight(event):
284 return event.wgt
285 get_wgt = weight
286 elif isinstance(get_wgt, str):
287 unwgt_name =get_wgt
288 def get_wgt(event):
289 event.parse_reweight()
290 return event.reweight_data[unwgt_name]
291 else:
292 unwgt_name = get_wgt.func_name
293
294
295
296 if hasattr(self, "written_weight"):
297 written_weight = lambda x: math.copysign(self.written_weight,float(x))
298 else:
299 written_weight = lambda x: x
300
301 all_wgt, cross, nb_event = self.initialize_unweighting(get_wgt, trunc_error)
302
303
304 def max_wgt_for_trunc(trunc):
305 """find the weight with the maximal truncation."""
306
307 xsum = 0
308 i=1
309 while (xsum - all_wgt[-i] * (i-1) <= cross['abs'] * trunc):
310 max_wgt = all_wgt[-i]
311 xsum += all_wgt[-i]
312 i +=1
313 if i == len(all_wgt):
314 break
315
316 return max_wgt
317
318
319
320 if not max_wgt:
321 if trunc_error == 0 or len(all_wgt)<2 or event_target:
322 max_wgt = all_wgt[-1]
323 else:
324 max_wgt = max_wgt_for_trunc(trunc_error)
325
326
327 if self.banner:
328 try:
329 import madgraph
330 except:
331 import internal.banner as banner_module
332 else:
333 import madgraph.various.banner as banner_module
334 if not isinstance(self.banner, banner_module.Banner):
335 banner = self.get_banner()
336
337 banner.modify_init_cross(cross)
338 strategy = banner.get_lha_strategy()
339
340 if strategy >0:
341 banner.set_lha_strategy(4)
342 else:
343 banner.set_lha_strategy(-4)
344
345 banner["unweight"] = "unweighted by %s" % unwgt_name
346 else:
347 banner = self.banner
348
349
350
351 nb_try = 20
352 nb_keep = 0
353 for i in range(nb_try):
354 self.seek(0)
355 if event_target:
356 if i==0:
357 max_wgt = max_wgt_for_trunc(0)
358 else:
359
360 efficiency = nb_keep/nb_event
361 needed_efficiency = event_target/nb_event
362 last_max_wgt = max_wgt
363 needed_max_wgt = last_max_wgt * efficiency / needed_efficiency
364
365 min_max_wgt = max_wgt_for_trunc(trunc_error)
366 max_wgt = max(min_max_wgt, needed_max_wgt)
367 max_wgt = min(max_wgt, all_wgt[-1])
368 if max_wgt == last_max_wgt:
369 if nb_keep <= event_target:
370 logger.log(log_level+10,"fail to reach target %s", event_target)
371 break
372 else:
373 break
374
375
376 if outputpath:
377 outfile = EventFile(outputpath, "w")
378
379
380 if self.banner and outputpath:
381 banner.write(outfile, close_tag=False)
382
383
384 nb_keep = 0
385 trunc_cross = 0
386 for event in self:
387 r = random.random()
388 wgt = get_wgt(event)
389 if abs(wgt) < r * max_wgt:
390 continue
391 elif wgt > 0:
392 nb_keep += 1
393 event.wgt = written_weight(max(wgt, max_wgt))
394
395 if abs(wgt) > max_wgt:
396 trunc_cross += abs(wgt) - max_wgt
397 if event_target ==0 or nb_keep <= event_target:
398 if outputpath:
399 outfile.write(str(event))
400
401 elif wgt < 0:
402 nb_keep += 1
403 event.wgt = -1 * max(abs(wgt), max_wgt)
404 if abs(wgt) > max_wgt:
405 trunc_cross += abs(wgt) - max_wgt
406 if outputpath and (event_target ==0 or nb_keep <= event_target):
407 outfile.write(str(event))
408
409 if event_target and nb_keep > event_target:
410 if not outputpath:
411
412 continue
413 elif event_target and i != nb_try-1 and nb_keep >= event_target *1.05:
414 outfile.close()
415
416 continue
417 else:
418 outfile.write("</LesHouchesEvents>\n")
419 outfile.close()
420 break
421 elif event_target == 0:
422 if outputpath:
423 outfile.write("</LesHouchesEvents>\n")
424 outfile.close()
425 break
426 elif outputpath:
427 outfile.close()
428
429
430 else:
431
432 logger.log(log_level+10,"fail to reach target event %s (iteration=%s)", event_target,i)
433
434
435
436
437 if event_target:
438 nb_events_unweighted = nb_keep
439 nb_keep = min( event_target, nb_keep)
440 else:
441 nb_events_unweighted = nb_keep
442
443 logger.log(log_level, "write %i event (efficiency %.2g %%, truncation %.2g %%) after %i iteration(s)",
444 nb_keep, nb_events_unweighted/nb_event*100, trunc_cross/cross['abs']*100, i)
445
446
447 if nb_keep != event_target and hasattr(self, "written_weight"):
448 written_weight = lambda x: math.copysign(self.written_weight*event_target/nb_keep, float(x))
449 startfile = EventFile(outputpath)
450 tmpname = pjoin(os.path.dirname(outputpath), "wgtcorrected_"+ os.path.basename(outputpath))
451 outfile = EventFile(tmpname, "w")
452 outfile.write(startfile.banner)
453 for event in startfile:
454 event.wgt = written_weight(event.wgt)
455 outfile.write(str(event))
456 outfile.write("</LesHouchesEvents>\n")
457 startfile.close()
458 outfile.close()
459 shutil.move(tmpname, outputpath)
460
461
462 self.max_wgt = max_wgt
463 return nb_keep
464
466 """ apply one or more fct on all event. """
467
468 opt= {"print_step": 2000}
469 opt.update(opts)
470
471 nb_fct = len(fcts)
472 out = []
473 for i in range(nb_fct):
474 out.append([])
475 self.seek(0)
476 nb_event = 0
477 for event in self:
478 nb_event += 1
479 if opt["print_step"] and nb_event % opt["print_step"] == 0:
480 if hasattr(self,"len"):
481 logger.info("currently at %s/%s event" % (nb_event, self.len))
482 else:
483 logger.info("currently at %s event" % nb_event)
484 for i in range(nb_fct):
485 out[i].append(fcts[i](event))
486 if nb_fct == 1:
487 return out[0]
488 else:
489 return out
490
493 """A way to read/write a gzipped lhef event"""
494
496 """A way to read a standard event file"""
497
499 """a class to read simultaneously multiple file and read them in mixing them.
500 Unweighting can be done at the same time.
501 The number of events in each file need to be provide in advance
502 (if not provide the file is first read to find that number"""
503
506
508 """if trunc_error is define here then this allow
509 to only read all the files twice and not three times."""
510 self.files = []
511 self.banner = ''
512 self.initial_nb_events = []
513 self.total_event_in_files = 0
514 self.curr_nb_events = []
515 self.allcross = []
516 self.error = []
517 self.across = []
518 self.scales = []
519 if start_list:
520 for p in start_list:
521 self.add(p)
522 self._configure = False
523
524 - def add(self, path, cross, error, across):
525 """ add a file to the pool, across allow to reweight the sum of weight
526 in the file to the given cross-section
527 """
528
529 if across == 0:
530
531 return
532
533 obj = EventFile(path)
534 if len(self.files) == 0 and not self.banner:
535 self.banner = obj.banner
536 self.curr_nb_events.append(0)
537 self.initial_nb_events.append(0)
538 self.allcross.append(cross)
539 self.across.append(across)
540 self.error.append(error)
541 self.scales.append(1)
542 self.files.append(obj)
543 self._configure = False
544
547
549
550 if not self._configure:
551 self.configure()
552
553 remaining_event = self.total_event_in_files - sum(self.curr_nb_events)
554 if remaining_event == 0:
555 raise StopIteration
556
557 nb_event = random.randint(1, remaining_event)
558 sum_nb=0
559 for i, obj in enumerate(self.files):
560 sum_nb += self.initial_nb_events[i] - self.curr_nb_events[i]
561 if nb_event <= sum_nb:
562 self.curr_nb_events[i] += 1
563 event = obj.next()
564 event.sample_scale = self.scales[i]
565 return event
566 else:
567 raise Exception
568
569
571 """define the part of the init_banner"""
572
573 if not self.banner:
574 return
575
576
577 grouped_cross = {}
578 grouped_error = {}
579 for i,ff in enumerate(self.files):
580 filename = ff.name
581 Pdir = [P for P in filename.split(os.path.sep) if P.startswith('P')][-1]
582 group = Pdir.split("_")[0][1:]
583 if group in grouped_cross:
584 grouped_cross[group] += self.allcross[i]
585 grouped_error[group] += self.error[i]**2
586 else:
587 grouped_cross[group] = self.allcross[i]
588 grouped_error[group] = self.error[i]**2
589
590 nb_group = len(grouped_cross)
591
592
593 try:
594 run_card = self.banner.run_card
595 except:
596 run_card = self.banner.charge_card("run_card")
597
598 init_information = run_card.get_banner_init_information()
599 init_information["nprup"] = nb_group
600
601 if run_card["lhe_version"] < 3:
602 init_information["generator_info"] = ""
603 else:
604 init_information["generator_info"] = "<generator name='MadGraph5_aMC@NLO' version='2.2.1'>please cite 1405.0301 </generator>\n"
605
606
607 cross_info = "%(cross)e %(error)e %(wgt)e %(id)i"
608 init_information["cross_info"] = []
609 for id in grouped_cross:
610 conv = {"id": int(id), "cross": grouped_cross[id], "error": math.sqrt(grouped_error[id]),
611 "wgt": wgt}
612 init_information["cross_info"].append( cross_info % conv)
613 init_information["cross_info"] = '\n'.join(init_information["cross_info"])
614
615
616
617 template_init =\
618 """ %(idbmup1)i %(idbmup2)i %(ebmup1)e %(ebmup2)e %(pdfgup1)i %(pdfgup2)i %(pdfsup1)i %(pdfsup2)i -3 %(nprup)i
619 %(cross_info)s
620 %(generator_info)s
621 """
622
623 self.banner["init"] = template_init % init_information
624
625
626
628 """ scan once the file to return
629 - the list of the hightest weight (of size trunc_error*NB_EVENT
630 - the cross-section by type of process
631 - the total number of events in the files
632 In top of that it initialise the information for the next routine
633 to determine how to choose which file to read
634 """
635 self.seek(0)
636 all_wgt = []
637 total_event = 0
638 sum_cross = collections.defaultdict(int)
639 for i,f in enumerate(self.files):
640 nb_event = 0
641
642
643 cross = collections.defaultdict(int)
644 new_wgt =[]
645 for event in f:
646 nb_event += 1
647 total_event += 1
648 event.sample_scale = 1
649 wgt = getwgt(event)
650 cross['all'] += wgt
651 cross['abs'] += abs(wgt)
652 cross[event.ievent] += wgt
653 new_wgt.append(abs(wgt))
654
655 if nb_event % 20000 == 0:
656 new_wgt.sort()
657
658 nb_keep = max(20, int(nb_event*trunc_error*15))
659 new_wgt = new_wgt[-nb_keep:]
660 if nb_event == 0:
661 raise Exception
662
663 self.initial_nb_events[i] = nb_event
664 self.scales[i] = self.across[i]/cross['abs'] if self.across[i] else 1
665
666
667 for key in cross:
668 sum_cross[key] += cross[key]* self.scales[i]
669 all_wgt +=[self.scales[i] * w for w in new_wgt]
670 all_wgt.sort()
671 nb_keep = max(20, int(total_event*trunc_error*10))
672 all_wgt = all_wgt[-nb_keep:]
673
674 self.total_event_in_files = total_event
675
676 all_wgt.sort()
677
678 nb_keep = max(20, int(total_event*trunc_error*10))
679 all_wgt = all_wgt[-nb_keep:]
680 self.seek(0)
681 self._configure = True
682 return all_wgt, sum_cross, total_event
683
690
692
693 return len(self.files)
694
695 - def seek(self, pos):
696 """ """
697
698 if pos !=0:
699 raise Exception
700 for i in range(len(self)):
701 self.curr_nb_events[i] = 0
702 for f in self.files:
703 f.seek(pos)
704
705 - def unweight(self, outputpath, get_wgt, **opts):
706 """unweight the current file according to wgt information wgt.
707 which can either be a fct of the event or a tag in the rwgt list.
708 max_wgt allow to do partial unweighting.
709 trunc_error allow for dynamical partial unweighting
710 event_target reweight for that many event with maximal trunc_error.
711 (stop to write event when target is reached)
712 """
713
714 if isinstance(get_wgt, str):
715 unwgt_name =get_wgt
716 def get_wgt_multi(event):
717 event.parse_reweight()
718 return event.reweight_data[unwgt_name] * event.sample_scale
719 else:
720 unwgt_name = get_wgt.func_name
721 get_wgt_multi = lambda event: get_wgt(event) * event.sample_scale
722
723
724 if opts['event_target']:
725 new_wgt = sum(self.across)/opts['event_target']
726 self.define_init_banner(new_wgt)
727 self.written_weight = new_wgt
728
729 return super(MultiEventFile, self).unweight(outputpath, get_wgt_multi, **opts)
730
733 """Class storing a single event information (list of particles + global information)"""
734
735 warning_order = True
736
738 """The initialization of an empty Event (or one associate to a text file)"""
739 list.__init__(self)
740
741
742 self.nexternal = 0
743 self.ievent = 0
744 self.wgt = 0
745 self.aqcd = 0
746 self.scale = 0
747 self.aqed = 0
748 self.aqcd = 0
749
750 self.tag = ''
751 self.comment = ''
752 self.reweight_data = {}
753 self.matched_scale_data = None
754 if text:
755 self.parse(text)
756
757
759 """Take the input file and create the structured information"""
760 text = re.sub(r'</?event>', '', text)
761 status = 'first'
762 for line in text.split('\n'):
763 line = line.strip()
764 if not line:
765 continue
766 if line.startswith('#'):
767 self.comment += '%s\n' % line
768 continue
769 if "<event" in line:
770 continue
771
772 if 'first' == status:
773 if '<rwgt>' in line:
774 status = 'tag'
775
776 if 'first' == status:
777 self.assign_scale_line(line)
778 status = 'part'
779 continue
780
781 if '<' in line:
782 status = 'tag'
783
784 if 'part' == status:
785 self.append(Particle(line, event=self))
786 else:
787 self.tag += '%s\n' % line
788
789
790 for i,particle in enumerate(self):
791 if self.warning_order:
792 if i < particle.mother1 or i < particle.mother2:
793 logger.warning("Order of particle in the event did not agree with parent/child order. This might be problematic for some code.")
794 Event.warning_order = False
795
796 if particle.mother1:
797 try:
798 particle.mother1 = self[int(particle.mother1) -1]
799 except Exception:
800 logger.warning("WRONG MOTHER INFO %s", self)
801 particle.mother1 = 0
802 if particle.mother2:
803 try:
804 particle.mother2 = self[int(particle.mother2) -1]
805 except Exception:
806 logger.warning("WRONG MOTHER INFO %s", self)
807 particle.mother2 = 0
808
809
811 """Parse the re-weight information in order to return a dictionary
812 {key: value}. If no group is define group should be '' """
813 if self.reweight_data:
814 return self.reweight_data
815 self.reweight_data = {}
816 self.reweight_order = []
817 start, stop = self.tag.find('<rwgt>'), self.tag.find('</rwgt>')
818 if start != -1 != stop :
819 pattern = re.compile(r'''<\s*wgt id=(?:\'|\")(?P<id>[^\'\"]+)(?:\'|\")\s*>\s*(?P<val>[\ded+-.]*)\s*</wgt>''')
820 data = pattern.findall(self.tag)
821 try:
822 self.reweight_data = dict([(pid, float(value)) for (pid, value) in data
823 if not self.reweight_order.append(pid)])
824
825 except ValueError, error:
826 raise Exception, 'Event File has unvalid weight. %s' % error
827 self.tag = self.tag[:start] + self.tag[stop+7:]
828 return self.reweight_data
829
831 """Parse the line containing the starting scale for the shower"""
832
833 if self.matched_scale_data is not None:
834 return self.matched_scale_data
835
836 self.matched_scale_data = []
837
838
839 pattern = re.compile("<scales\s|</scales>")
840 data = re.split(pattern,self.tag)
841 if len(data) == 1:
842 return []
843 else:
844 tmp = {}
845 start,content, end = data
846 self.tag = "%s%s" % (start, end)
847 pattern = re.compile("pt_clust_(\d*)=\"([\de+-.]*)\"")
848 for id,value in pattern.findall(content):
849 tmp[int(id)] = float(value)
850
851 for i in range(1, len(tmp)+1):
852 self.matched_scale_data.append(tmp[i])
853
854 return self.matched_scale_data
855
856
857
858
859
860 - def add_decay_to_particle(self, position, decay_event):
861 """define the decay of the particle id by the event pass in argument"""
862
863 this_particle = self[position]
864
865 this_particle.status = 2
866 this_particle.helicity = 0
867
868
869 decay_particle = decay_event[0]
870 this_4mom = FourMomentum(this_particle)
871 nb_part = len(self)
872
873 thres = decay_particle.E*1e-10
874 assert max(decay_particle.px, decay_particle.py, decay_particle.pz) < thres,\
875 "not on rest particle %s %s %s %s" % (decay_particle.E, decay_particle.px,decay_particle.py,decay_particle.pz)
876
877 self.nexternal += decay_event.nexternal -1
878 old_scales = list(self.parse_matching_scale())
879 if old_scales:
880 self.matched_scale_data.pop(position-2)
881
882
883 for particle in decay_event[1:]:
884
885 new_particle = Particle(particle, self)
886 new_particle.event_id = len(self)
887 self.append(new_particle)
888 if old_scales:
889 self.matched_scale_data.append(old_scales[position-2])
890
891 new_momentum = this_4mom.boost(FourMomentum(new_particle))
892 new_particle.set_momentum(new_momentum)
893
894 for tag in ['mother1', 'mother2']:
895 mother = getattr(particle, tag)
896 if isinstance(mother, Particle):
897 mother_id = getattr(particle, tag).event_id
898 if mother_id == 0:
899 setattr(new_particle, tag, this_particle)
900 else:
901 setattr(new_particle, tag, self[nb_part + mother_id -1])
902 elif tag == "mother2" and isinstance(particle.mother1, Particle):
903 new_particle.mother2 = this_particle
904 else:
905 raise Exception, "Something weird happens. Please report it for investigation"
906
907
908 max_color=501
909 for particle in self[:nb_part]:
910 max_color=max(max_color, particle.color1, particle.color2)
911
912
913 color_mapping = {}
914 color_mapping[decay_particle.color1] = this_particle.color1
915 color_mapping[decay_particle.color2] = this_particle.color2
916 for particle in self[nb_part:]:
917 if particle.color1:
918 if particle.color1 not in color_mapping:
919 max_color +=1
920 color_mapping[particle.color1] = max_color
921 particle.color1 = max_color
922 else:
923 particle.color1 = color_mapping[particle.color1]
924 if particle.color2:
925 if particle.color2 not in color_mapping:
926 max_color +=1
927 color_mapping[particle.color2] = max_color
928 particle.color2 = max_color
929 else:
930 particle.color2 = color_mapping[particle.color2]
931
932
933
935
936 to_remove = []
937 if event_id is not None:
938 to_remove.append(self[event_id])
939
940 if pdg_code:
941 for particle in self:
942 if particle.pid == pdg_code:
943 to_remove.append(particle)
944
945 new_event = Event()
946
947 for tag in ['nexternal', 'ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']:
948 setattr(new_event, tag, getattr(self, tag))
949
950 for particle in self:
951 if isinstance(particle.mother1, Particle) and particle.mother1 in to_remove:
952 to_remove.append(particle)
953 if particle.status == 1:
954 new_event.nexternal -= 1
955 continue
956 elif isinstance(particle.mother2, Particle) and particle.mother2 in to_remove:
957 to_remove.append(particle)
958 if particle.status == 1:
959 new_event.nexternal -= 1
960 continue
961 else:
962 new_event.append(Particle(particle))
963
964
965
966 for pos, particle in enumerate(new_event):
967 particle.event_id = pos
968 if particle in to_remove:
969 particle.status = 1
970 return new_event
971
972 - def get_decay(self, pdg_code=0, event_id=None):
973
974 to_start = []
975 if event_id is not None:
976 to_start.append(self[event_id])
977
978 elif pdg_code:
979 for particle in self:
980 if particle.pid == pdg_code:
981 to_start.append(particle)
982 break
983
984 new_event = Event()
985
986 for tag in ['ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']:
987 setattr(new_event, tag, getattr(self, tag))
988
989
990 old2new = {}
991 new_decay_part = Particle(to_start[0])
992 new_decay_part.mother1 = None
993 new_decay_part.mother2 = None
994 new_decay_part.status = -1
995 old2new[new_decay_part.event_id] = len(old2new)
996 new_event.append(new_decay_part)
997
998
999
1000 for particle in self:
1001 if isinstance(particle.mother1, Particle) and particle.mother1.event_id in old2new\
1002 or isinstance(particle.mother2, Particle) and particle.mother2.event_id in old2new:
1003 old2new[particle.event_id] = len(old2new)
1004 new_event.append(Particle(particle))
1005
1006
1007
1008 nexternal = 0
1009 for pos, particle in enumerate(new_event):
1010 particle.event_id = pos
1011 if particle.mother1:
1012 particle.mother1 = new_event[old2new[particle.mother1.event_id]]
1013 if particle.mother2:
1014 particle.mother2 = new_event[old2new[particle.mother2.event_id]]
1015 if particle.status in [-1,1]:
1016 nexternal +=1
1017 new_event.nexternal = nexternal
1018
1019 return new_event
1020
1021
1023 """check various property of the events"""
1024
1025
1026 E, px, py, pz = 0,0,0,0
1027 absE, abspx, abspy, abspz = 0,0,0,0
1028 for particle in self:
1029 coeff = 1
1030 if particle.status == -1:
1031 coeff = -1
1032 elif particle.status != 1:
1033 continue
1034 E += coeff * particle.E
1035 absE += abs(particle.E)
1036 px += coeff * particle.px
1037 py += coeff * particle.py
1038 pz += coeff * particle.pz
1039 abspx += abs(particle.px)
1040 abspy += abs(particle.py)
1041 abspz += abs(particle.pz)
1042
1043 threshold = 5e-7
1044 if E/absE > threshold:
1045 logger.critical(self)
1046 raise Exception, "Do not conserve Energy %s, %s" % (E/absE, E)
1047 if px/abspx > threshold:
1048 logger.critical(self)
1049 raise Exception, "Do not conserve Px %s, %s" % (px/abspx, px)
1050 if py/abspy > threshold:
1051 logger.critical(self)
1052 raise Exception, "Do not conserve Py %s, %s" % (py/abspy, py)
1053 if pz/abspz > threshold:
1054 logger.critical(self)
1055 raise Exception, "Do not conserve Pz %s, %s" % (pz/abspz, pz)
1056
1057
1058 self.check_color_structure()
1059
1061 """read the line corresponding to global event line
1062 format of the line is:
1063 Nexternal IEVENT WEIGHT SCALE AEW AS
1064 """
1065 inputs = line.split()
1066 assert len(inputs) == 6
1067 self.nexternal=int(inputs[0])
1068 self.ievent=int(inputs[1])
1069 self.wgt=float(inputs[2])
1070 self.scale=float(inputs[3])
1071 self.aqed=float(inputs[4])
1072 self.aqcd=float(inputs[5])
1073
1075 """Return the unique tag identifying the SubProcesses for the generation.
1076 Usefull for program like MadSpin and Reweight module."""
1077
1078 initial, final, order = [], [], [[], []]
1079 for particle in self:
1080 if particle.status == -1:
1081 initial.append(particle.pid)
1082 order[0].append(particle.pid)
1083 elif particle.status == 1:
1084 final.append(particle.pid)
1085 order[1].append(particle.pid)
1086 initial.sort(), final.sort()
1087 tag = (tuple(initial), tuple(final))
1088 return tag, order
1089
1091 """return a list with the helicities in the order asked for"""
1092
1093
1094
1095
1096 order = [list(get_order[0]), list(get_order[1])]
1097 out = [9] *(len(order[0])+len(order[1]))
1098 for i, part in enumerate(self):
1099 if part.status == 1:
1100 try:
1101 ind = order[1].index(part.pid)
1102 except ValueError, error:
1103 if not allow_reversed:
1104 raise error
1105 else:
1106 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
1107 try:
1108 return self.get_helicity(order, False)
1109 except ValueError:
1110 raise error
1111 position = len(order[0]) + ind
1112 order[1][ind] = 0
1113 elif part.status == -1:
1114 try:
1115 ind = order[0].index(part.pid)
1116 except ValueError, error:
1117 if not allow_reversed:
1118 raise error
1119 else:
1120 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
1121 try:
1122 return self.get_helicity(order, False)
1123 except ValueError:
1124 raise error
1125
1126 position = ind
1127 order[0][ind] = 0
1128 else:
1129 continue
1130 out[position] = int(part.helicity)
1131 return out
1132
1133
1135 """check the validity of the color structure"""
1136
1137
1138 color_index = collections.defaultdict(int)
1139 for particle in self:
1140 if particle.status in [-1,1]:
1141 if particle.color1:
1142 color_index[particle.color1] +=1
1143 if particle.color2:
1144 color_index[particle.color2] +=1
1145
1146 for key,value in color_index.items():
1147 if value > 2:
1148 print self
1149 print key, value
1150 raise Exception, 'Wrong color_flow'
1151
1152
1153 check = []
1154 popup_index = []
1155 for particle in self:
1156 mothers = []
1157 childs = []
1158 if particle.mother1:
1159 mothers.append(particle.mother1)
1160 if particle.mother2 and particle.mother2 is not particle.mother1:
1161 mothers.append(particle.mother2)
1162 if not mothers:
1163 continue
1164 if (particle.mother1.event_id, particle.mother2.event_id) in check:
1165 continue
1166 check.append((particle.mother1.event_id, particle.mother2.event_id))
1167
1168 childs = [p for p in self if p.mother1 is particle.mother1 and \
1169 p.mother2 is particle.mother2]
1170
1171 mcolors = []
1172 manticolors = []
1173 for m in mothers:
1174 if m.color1:
1175 if m.color1 in manticolors:
1176 manticolors.remove(m.color1)
1177 else:
1178 mcolors.append(m.color1)
1179 if m.color2:
1180 if m.color2 in mcolors:
1181 mcolors.remove(m.color2)
1182 else:
1183 manticolors.append(m.color2)
1184 ccolors = []
1185 canticolors = []
1186 for m in childs:
1187 if m.color1:
1188 if m.color1 in canticolors:
1189 canticolors.remove(m.color1)
1190 else:
1191 ccolors.append(m.color1)
1192 if m.color2:
1193 if m.color2 in ccolors:
1194 ccolors.remove(m.color2)
1195 else:
1196 canticolors.append(m.color2)
1197 for index in mcolors[:]:
1198 if index in ccolors:
1199 mcolors.remove(index)
1200 ccolors.remove(index)
1201 for index in manticolors[:]:
1202 if index in canticolors:
1203 manticolors.remove(index)
1204 canticolors.remove(index)
1205
1206 if mcolors != []:
1207
1208 if len(canticolors) + len(mcolors) != 3:
1209 logger.critical(str(self))
1210 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs])
1211 else:
1212 popup_index += canticolors
1213 elif manticolors != []:
1214
1215 if len(ccolors) + len(manticolors) != 3:
1216 logger.critical(str(self))
1217 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs])
1218 else:
1219 popup_index += ccolors
1220
1221
1222 if len(popup_index) != len(set(popup_index)):
1223 logger.critical(self)
1224 raise Exception, "Wrong color flow: identical poping-up index, %s" % (popup_index)
1225
1227 """return a correctly formatted LHE event"""
1228
1229 out="""<event>
1230 %(scale)s
1231 %(particles)s
1232 %(comments)s
1233 %(tag)s
1234 %(reweight)s
1235 </event>
1236 """
1237
1238 scale_str = "%2d %6d %+13.7e %14.8e %14.8e %14.8e" % \
1239 (self.nexternal,self.ievent,self.wgt,self.scale,self.aqed,self.aqcd)
1240 if self.reweight_data:
1241
1242 if set(self.reweight_data.keys()) != set(self.reweight_order):
1243 self.reweight_order += [k for k in self.reweight_data.keys() \
1244 if k not in self.reweight_order]
1245
1246 reweight_str = '<rwgt>\n%s\n</rwgt>' % '\n'.join(
1247 '<wgt id=\'%s\'> %+13.7e </wgt>' % (i, float(self.reweight_data[i]))
1248 for i in self.reweight_order)
1249 else:
1250 reweight_str = ''
1251
1252 tag_str = self.tag
1253 if self.matched_scale_data:
1254 tag_str = "<scales %s></scales>%s" % (
1255 ' '.join(['pt_clust_%i=\"%s\"' % (i,v)
1256 for i,v in enumerate(self.matched_scale_data)]),
1257 self.tag)
1258
1259 out = out % {'scale': scale_str,
1260 'particles': '\n'.join([str(p) for p in self]),
1261 'tag': tag_str,
1262 'comments': self.comment,
1263 'reweight': reweight_str}
1264 return re.sub('[\n]+', '\n', out)
1265
1267 """return the momenta str in the order asked for"""
1268
1269
1270
1271 order = [list(get_order[0]), list(get_order[1])]
1272 out = [''] *(len(order[0])+len(order[1]))
1273 for i, part in enumerate(self):
1274 if part.status == 1:
1275 try:
1276 ind = order[1].index(part.pid)
1277 except ValueError, error:
1278 if not allow_reversed:
1279 raise error
1280 else:
1281 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
1282 try:
1283 return self.get_momenta_str(order, False)
1284 except ValueError:
1285 raise error
1286 position = len(order[0]) + ind
1287 order[1][ind] = 0
1288 elif part.status == -1:
1289 try:
1290 ind = order[0].index(part.pid)
1291 except ValueError, error:
1292 if not allow_reversed:
1293 raise error
1294 else:
1295 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
1296 try:
1297 return self.get_momenta_str(order, False)
1298 except ValueError:
1299 raise error
1300
1301 position = ind
1302 order[0][ind] = 0
1303 else:
1304 continue
1305 format = '%.12f'
1306 format_line = ' '.join([format]*4) + ' \n'
1307 out[position] = format_line % (part.E, part.px, part.py, part.pz)
1308
1309 out = ''.join(out).replace('e','d')
1310 return out
1311
1313 """A class to allow to read both gzip and not gzip file.
1314 containing only weight from pythia --generated by SysCalc"""
1315
1316 - def __new__(self, path, mode='r', *args, **opt):
1317 if path.endswith(".gz"):
1318 try:
1319 return gzip.GzipFile.__new__(WeightFileGzip, path, mode, *args, **opt)
1320 except IOError, error:
1321 raise
1322 except Exception, error:
1323 if mode == 'r':
1324 misc.gunzip(path)
1325 return file.__new__(WeightFileNoGzip, path[:-3], mode, *args, **opt)
1326 else:
1327 return file.__new__(WeightFileNoGzip, path, mode, *args, **opt)
1328
1329
1330 - def __init__(self, path, mode='r', *args, **opt):
1331 """open file and read the banner [if in read mode]"""
1332
1333 super(EventFile, self).__init__(path, mode, *args, **opt)
1334 self.banner = ''
1335 if mode == 'r':
1336 line = ''
1337 while '</header>' not in line.lower():
1338 try:
1339 line = super(EventFile, self).next()
1340 except StopIteration:
1341 self.seek(0)
1342 self.banner = ''
1343 break
1344 if "<event" in line.lower():
1345 self.seek(0)
1346 self.banner = ''
1347 break
1348
1349 self.banner += line
1350
1354
1357
1360 """a convenient object for 4-momenta operation"""
1361
1362 - def __init__(self, obj=0, px=0, py=0, pz=0, E=0):
1363 """initialize the four momenta"""
1364
1365 if obj is 0 and E:
1366 obj = E
1367
1368 if isinstance(obj, (FourMomentum, Particle)):
1369 px = obj.px
1370 py = obj.py
1371 pz = obj.pz
1372 E = obj.E
1373 else:
1374 E =obj
1375
1376
1377 self.E = E
1378 self.px = px
1379 self.py = py
1380 self.pz =pz
1381
1382 @property
1384 """return the mass"""
1385 return math.sqrt(self.E**2 - self.px**2 - self.py**2 - self.pz**2)
1386
1388 """return the mass square"""
1389 return self.E**2 - self.px**2 - self.py**2 - self.pz**2
1390
1391 @property
1393 return math.sqrt(max(0, self.pt2()))
1394
1395 @property
1397 norm = math.sqrt(self.px**2 + self.py**2+self.pz**2)
1398 return 0.5* math.log((norm - self.pz) / (norm + self.pz))
1399
1401 """ return the pt square """
1402
1403 return self.px**2 + self.py**2
1404
1406
1407 assert isinstance(obj, FourMomentum)
1408 new = FourMomentum(self.E+obj.E,
1409 self.px + obj.px,
1410 self.py + obj.py,
1411 self.pz + obj.pz)
1412 return new
1413
1415 """update the object with the sum"""
1416 self.E += obj.E
1417 self.px += obj.px
1418 self.py += obj.py
1419 self.pz += obj.pz
1420 return self
1421
1423 assert power in [1,2]
1424
1425 if power == 1:
1426 return FourMomentum(self)
1427 elif power == 2:
1428 return self.mass_sqr()
1429
1431 """mom 4-momenta is suppose to be given in the rest frame of this 4-momenta.
1432 the output is the 4-momenta in the frame of this 4-momenta
1433 function copied from HELAS routine."""
1434
1435
1436 pt = self.px**2 + self.py**2 + self.pz**2
1437 if pt:
1438 s3product = self.px * mom.px + self.py * mom.py + self.pz * mom.pz
1439 mass = self.mass
1440 lf = (mom.E + (self.E - mass) * s3product / pt ) / mass
1441 return FourMomentum(E=(self.E*mom.E+s3product)/mass,
1442 px=mom.px + self.px * lf,
1443 py=mom.py + self.py * lf,
1444 pz=mom.pz + self.pz * lf)
1445 else:
1446 return FourMomentum(mom)
1447
1448
1449 if '__main__' == __name__:
1450
1451
1452 if False:
1453 lhe = EventFile('unweighted_events.lhe.gz')
1454 output = open('output_events.lhe', 'w')
1455
1456 output.write(lhe.banner)
1457
1458 for event in lhe:
1459 for particle in event:
1460
1461 particle.mass = 0
1462 particle.vtim = 2
1463
1464
1465 output.write(str(event))
1466 output.write('</LesHouchesEvent>\n')
1467
1468
1469
1470
1471 if True:
1472 lhe = EventFile('unweighted_events.lhe.gz')
1473 import matplotlib.pyplot as plt
1474 import matplotlib.gridspec as gridspec
1475 nbins = 100
1476
1477 nb_pass = 0
1478 data = []
1479 for event in lhe:
1480 etaabs = 0
1481 etafinal = 0
1482 for particle in event:
1483 if particle.status==1:
1484 p = FourMomentum(particle)
1485 eta = p.pseudorapidity
1486 if abs(eta) > etaabs:
1487 etafinal = eta
1488 etaabs = abs(eta)
1489 if etaabs < 4:
1490 data.append(etafinal)
1491 nb_pass +=1
1492
1493
1494 print nb_pass
1495 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1])
1496 gs1.update(wspace=0, hspace=0)
1497 ax = plt.subplot(gs1[0])
1498
1499 n, bins, patches = ax.hist(data, nbins, histtype='step', label='original')
1500 ax_c = ax.twinx()
1501 ax_c.set_ylabel('MadGraph5_aMC@NLO')
1502 ax_c.yaxis.set_label_coords(1.01, 0.25)
1503 ax_c.set_yticks(ax.get_yticks())
1504 ax_c.set_yticklabels([])
1505 ax.set_xlim([-4,4])
1506 print "bin value:", n
1507 print "start/end point of bins", bins
1508 plt.axis('on')
1509 plt.xlabel('weight ratio')
1510 plt.show()
1511
1512
1513
1514 if False:
1515 lhe = EventFile('unweighted_events.lhe')
1516 import matplotlib.pyplot as plt
1517 import matplotlib.gridspec as gridspec
1518 nbins = 100
1519
1520
1521 mtau, wtau = 1.777, 4.027000e-13
1522 nb_pass = 0
1523 data, data2, data3 = [], [], []
1524 for event in lhe:
1525 nb_pass +=1
1526 if nb_pass > 10000:
1527 break
1528 tau1 = FourMomentum()
1529 tau2 = FourMomentum()
1530 for part in event:
1531 if part.pid in [-12,11,16]:
1532 momenta = FourMomentum(part)
1533 tau1 += momenta
1534 elif part.pid == 15:
1535 tau2 += FourMomentum(part)
1536
1537 if abs((mtau-tau2.mass())/wtau)<1e6 and tau2.mass() >1:
1538 data.append((tau1.mass()-mtau)/wtau)
1539 data2.append((tau2.mass()-mtau)/wtau)
1540 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1])
1541 gs1.update(wspace=0, hspace=0)
1542 ax = plt.subplot(gs1[0])
1543
1544 n, bins, patches = ax.hist(data2, nbins, histtype='step', label='original')
1545 n2, bins2, patches2 = ax.hist(data, bins=bins, histtype='step',label='reconstructed')
1546 import cmath
1547
1548 breit = lambda m : math.sqrt(4*math.pi)*1/(((m)**2-mtau**2)**2+(mtau*wtau)**2)*wtau
1549
1550 data3 = [breit(mtau + x*wtau)*wtau*16867622.6624*50 for x in bins]
1551
1552 ax.plot(bins, data3,label='breit-wigner')
1553
1554 ax.legend()
1555
1556 ax_c = ax.twinx()
1557 ax_c.set_ylabel('MadGraph5_aMC@NLO')
1558 ax_c.yaxis.set_label_coords(1.01, 0.25)
1559 ax_c.set_yticks(ax.get_yticks())
1560 ax_c.set_yticklabels([])
1561
1562 plt.title('invariant mass of tau LHE/reconstructed')
1563 plt.axis('on')
1564 ax.set_xticklabels([])
1565
1566 ax = plt.subplot(gs1[1])
1567 data4 = [n[i]/(data3[i]) for i in range(nbins)]
1568 ax.plot(bins, data4 + [0] , 'b')
1569 data4 = [n2[i]/(data3[i]) for i in range(nbins)]
1570 ax.plot(bins, data4 + [0] , 'g')
1571 ax.set_ylim([0,2])
1572
1573 tick = ax.get_yticks()
1574 ax.set_yticks(tick[:-1])
1575
1576
1577 plt.axis('on')
1578 plt.xlabel('(M - Mtau)/Wtau')
1579 plt.show()
1580