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