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 not path.endswith(".gz"):
154 return file.__new__(EventFileNoGzip, path, mode, *args, **opt)
155 elif mode == 'r' and not os.path.exists(path) and os.path.exists(path[:-3]):
156 return EventFile.__new__(EventFileNoGzip, path[:-3], mode, *args, **opt)
157 else:
158 try:
159 return gzip.GzipFile.__new__(EventFileGzip, path, mode, *args, **opt)
160 except IOError, error:
161 raise
162 except Exception, error:
163 if mode == 'r':
164 misc.gunzip(path)
165 return file.__new__(EventFileNoGzip, path[:-3], mode, *args, **opt)
166
167
168 - def __init__(self, path, mode='r', *args, **opt):
169 """open file and read the banner [if in read mode]"""
170
171 try:
172 super(EventFile, self).__init__(path, mode, *args, **opt)
173 except IOError:
174 if '.gz' in path and isinstance(self, EventFileNoGzip) and\
175 mode == 'r' and os.path.exists(path[:-3]):
176 super(EventFile, self).__init__(path[:-3], mode, *args, **opt)
177
178 self.banner = ''
179 if mode == 'r':
180 line = ''
181 while '</init>' not in line.lower():
182 try:
183 line = super(EventFile, self).next()
184 except StopIteration:
185 self.seek(0)
186 self.banner = ''
187 break
188 if "<event" in line.lower():
189 self.seek(0)
190 self.banner = ''
191 break
192
193 self.banner += line
194
204
205 @property
207 """return the cross-section of the file #from the banner"""
208 try:
209 return self._cross
210 except Exception:
211 pass
212
213 onebanner = self.get_banner()
214 self._cross = onebanner.get_cross()
215 return self._cross
216
218 if self.closed:
219 return 0
220 if hasattr(self,"len"):
221 return self.len
222
223 init_pos = self.tell()
224 self.seek(0)
225 nb_event=0
226 for _ in self:
227 nb_event +=1
228 self.len = nb_event
229 self.seek(init_pos)
230 return self.len
231
233 """get next event"""
234 text = ''
235 line = ''
236 mode = 0
237 while '</event>' not in line:
238 line = super(EventFile, self).next().lower()
239 if '<event' in line:
240 mode = 1
241 text = ''
242 if mode:
243 text += line
244 return Event(text)
245
247 """ scan once the file to return
248 - the list of the hightest weight (of size trunc_error*NB_EVENT
249 - the cross-section by type of process
250 - the total number of events in the file
251 """
252
253
254
255 self.seek(0)
256 all_wgt = []
257 cross = collections.defaultdict(int)
258 nb_event = 0
259 for event in self:
260 nb_event +=1
261 wgt = get_wgt(event)
262 cross['all'] += wgt
263 cross['abs'] += abs(wgt)
264 cross[event.ievent] += wgt
265 all_wgt.append(abs(wgt))
266
267 if nb_event % 20000 == 0:
268 all_wgt.sort()
269
270 nb_keep = max(20, int(nb_event*trunc_error*15))
271 all_wgt = all_wgt[-nb_keep:]
272
273
274 all_wgt.sort()
275
276 nb_keep = max(20, int(nb_event*trunc_error*10))
277 all_wgt = all_wgt[-nb_keep:]
278 self.seek(0)
279 return all_wgt, cross, nb_event
280
281
282 - def unweight(self, outputpath, get_wgt=None, max_wgt=0, trunc_error=0, event_target=0,
283 log_level=logging.INFO):
284 """unweight the current file according to wgt information wgt.
285 which can either be a fct of the event or a tag in the rwgt list.
286 max_wgt allow to do partial unweighting.
287 trunc_error allow for dynamical partial unweighting
288 event_target reweight for that many event with maximal trunc_error.
289 (stop to write event when target is reached)
290 """
291 if not get_wgt:
292 def weight(event):
293 return event.wgt
294 get_wgt = weight
295 unwgt_name = "central weight"
296 elif isinstance(get_wgt, str):
297 unwgt_name =get_wgt
298 def get_wgt(event):
299 event.parse_reweight()
300 return event.reweight_data[unwgt_name]
301 else:
302 unwgt_name = get_wgt.func_name
303
304
305
306 if hasattr(self, "written_weight"):
307 written_weight = lambda x: math.copysign(self.written_weight,float(x))
308 else:
309 written_weight = lambda x: x
310
311 all_wgt, cross, nb_event = self.initialize_unweighting(get_wgt, trunc_error)
312
313
314 def max_wgt_for_trunc(trunc):
315 """find the weight with the maximal truncation."""
316
317 xsum = 0
318 i=1
319 while (xsum - all_wgt[-i] * (i-1) <= cross['abs'] * trunc):
320 max_wgt = all_wgt[-i]
321 xsum += all_wgt[-i]
322 i +=1
323 if i == len(all_wgt):
324 break
325
326 return max_wgt
327
328
329
330 if not max_wgt:
331 if trunc_error == 0 or len(all_wgt)<2 or event_target:
332 max_wgt = all_wgt[-1]
333 else:
334 max_wgt = max_wgt_for_trunc(trunc_error)
335
336
337 if self.banner:
338 try:
339 import internal
340 except:
341 import madgraph.various.banner as banner_module
342 else:
343 import internal.banner as banner_module
344 if not isinstance(self.banner, banner_module.Banner):
345 banner = self.get_banner()
346
347 banner.modify_init_cross(cross)
348 strategy = banner.get_lha_strategy()
349
350 if strategy >0:
351 banner.set_lha_strategy(4)
352 else:
353 banner.set_lha_strategy(-4)
354
355 banner["unweight"] = "unweighted by %s" % unwgt_name
356 else:
357 banner = self.banner
358
359
360
361 nb_try = 20
362 nb_keep = 0
363 for i in range(nb_try):
364 self.seek(0)
365 if event_target:
366 if i==0:
367 max_wgt = max_wgt_for_trunc(0)
368 else:
369
370 efficiency = nb_keep/nb_event
371 needed_efficiency = event_target/nb_event
372 last_max_wgt = max_wgt
373 needed_max_wgt = last_max_wgt * efficiency / needed_efficiency
374
375 min_max_wgt = max_wgt_for_trunc(trunc_error)
376 max_wgt = max(min_max_wgt, needed_max_wgt)
377 max_wgt = min(max_wgt, all_wgt[-1])
378 if max_wgt == last_max_wgt:
379 if nb_keep <= event_target:
380 logger.log(log_level+10,"fail to reach target %s", event_target)
381 break
382 else:
383 break
384
385
386 if outputpath:
387 outfile = EventFile(outputpath, "w")
388
389
390 if self.banner and outputpath:
391 banner.write(outfile, close_tag=False)
392
393
394 nb_keep = 0
395 trunc_cross = 0
396 for event in self:
397 r = random.random()
398 wgt = get_wgt(event)
399 if abs(wgt) < r * max_wgt:
400 continue
401 elif wgt > 0:
402 nb_keep += 1
403 event.wgt = written_weight(max(wgt, max_wgt))
404 if abs(wgt) > max_wgt:
405 trunc_cross += abs(wgt) - max_wgt
406 if event_target ==0 or nb_keep <= event_target:
407 if outputpath:
408 outfile.write(str(event))
409
410 elif wgt < 0:
411 nb_keep += 1
412 event.wgt = -1* written_weight(max(abs(wgt), max_wgt))
413 if abs(wgt) > max_wgt:
414 trunc_cross += abs(wgt) - max_wgt
415 if outputpath and (event_target ==0 or nb_keep <= event_target):
416 outfile.write(str(event))
417
418 if event_target and nb_keep > event_target:
419 if not outputpath:
420
421 continue
422 elif event_target and i != nb_try-1 and nb_keep >= event_target *1.05:
423 outfile.close()
424
425 continue
426 else:
427 outfile.write("</LesHouchesEvents>\n")
428 outfile.close()
429 break
430 elif event_target == 0:
431 if outputpath:
432 outfile.write("</LesHouchesEvents>\n")
433 outfile.close()
434 break
435 elif outputpath:
436 outfile.close()
437
438
439 else:
440
441 logger.log(log_level+10,"fail to reach target event %s (iteration=%s)", event_target,i)
442
443
444
445
446 if event_target:
447 nb_events_unweighted = nb_keep
448 nb_keep = min( event_target, nb_keep)
449 else:
450 nb_events_unweighted = nb_keep
451
452 logger.log(log_level, "write %i event (efficiency %.2g %%, truncation %.2g %%) after %i iteration(s)",
453 nb_keep, nb_events_unweighted/nb_event*100, trunc_cross/cross['abs']*100, i)
454
455
456 if nb_keep != event_target and hasattr(self, "written_weight"):
457 written_weight = lambda x: math.copysign(self.written_weight*event_target/nb_keep, float(x))
458 startfile = EventFile(outputpath)
459 tmpname = pjoin(os.path.dirname(outputpath), "wgtcorrected_"+ os.path.basename(outputpath))
460 outfile = EventFile(tmpname, "w")
461 outfile.write(startfile.banner)
462 for event in startfile:
463 event.wgt = written_weight(event.wgt)
464 outfile.write(str(event))
465 outfile.write("</LesHouchesEvents>\n")
466 startfile.close()
467 outfile.close()
468 shutil.move(tmpname, outputpath)
469
470
471 self.max_wgt = max_wgt
472 return nb_keep
473
475 """ apply one or more fct on all event. """
476
477 opt= {"print_step": 2000, "maxevent":float("inf"),'no_output':False}
478 opt.update(opts)
479
480 nb_fct = len(fcts)
481 out = []
482 for i in range(nb_fct):
483 out.append([])
484 self.seek(0)
485 nb_event = 0
486 for event in self:
487 nb_event += 1
488 if opt["print_step"] and nb_event % opt["print_step"] == 0:
489 if hasattr(self,"len"):
490 logger.info("currently at %s/%s event" % (nb_event, self.len))
491 else:
492 logger.info("currently at %s event" % nb_event)
493 for i in range(nb_fct):
494 if opts['no_output']:
495 fcts[i](event)
496 else:
497 out[i].append(fcts[i](event))
498 if nb_event > opt['maxevent']:
499 break
500 if nb_fct == 1:
501 return out[0]
502 else:
503 return out
504
505
506 - def update_Hwu(self, hwu, fct, name='lhe', keep_wgt=True):
507
508 first=True
509 def add_to_Hwu(event):
510 """function to update the HwU on the flight"""
511
512 value = fct(event)
513
514
515 if first:
516
517 if isinstance(value, dict):
518 hwu.add_line(value.keys())
519 else:
520 hwu.add_line(name)
521 if keep_wgt is True:
522 hwu.add_line(['%s_%s' % (name, key)
523 for key in event.reweight_data])
524 first = False
525
526 if isinstance(value, dict):
527 hwu.addEvent(value)
528 else:
529 hwu.addEvent({name:value})
530 if keep_wgt:
531 event.parse_reweight()
532 if keep_wgt is True:
533 data = dict(('%s_%s' % (name, key),value)
534 for key in event.reweight_data)
535 hwu.addEvent(data)
536
537
538
539 self.apply_fct_on_event(add_to_Hwu, no_output=True)
540 return hwu
541
543 """take the lhe file and add the matchscale from the pythia_input file"""
544
545 if pythia_input:
546 def next_data():
547 for line in open(pythia_input):
548 if line.startswith('#'):
549 continue
550 data = line.split()
551 print (int(data[0]), data[-3], data[-2], data[-1])
552 yield (int(data[0]), data[-3], data[-2], data[-1])
553 else:
554 def next_data():
555 i=0
556 while 1:
557 yield [i,0,0,0]
558 i+=1
559 sys_iterator = next_data()
560
561 self.seek(0)
562 out = open(out_path,'w')
563
564 pdf_pattern = re.compile(r'''<init>(.*)</init>''', re.M+re.S)
565 init = pdf_pattern.findall(self.banner)[0].split('\n',2)[1]
566 id1, id2, _, _, _, _, pdf1,pdf2,_,_ = init.split()
567 id = [int(id1), int(id2)]
568 type = []
569 for i in range(2):
570 if abs(id[i]) == 2212:
571 if i > 0:
572 type.append(1)
573 else:
574 type.append(-1)
575 else:
576 type.append(0)
577 pdf = max(int(pdf1),int(pdf2))
578
579 out.write("<header>\n" + \
580 "<orgpdf>%i</orgpdf>\n" % pdf + \
581 "<beams> %s %s</beams>\n" % tuple(type) + \
582 "</header>\n")
583
584
585 nevt, smin, smax, scomp = sys_iterator.next()
586 for i, orig_event in enumerate(self):
587 if i < nevt:
588 continue
589 new_event = Event()
590 sys = orig_event.parse_syscalc_info()
591 new_event.syscalc_data = sys
592 if smin:
593 new_event.syscalc_data['matchscale'] = "%s %s %s" % (smin, scomp, smax)
594 out.write(str(new_event), nevt)
595 try:
596 nevt, smin, smax, scomp = sys_iterator.next()
597 except StopIteration:
598 break
599
600
601
602
603
604
605
606
607 -class EventFileGzip(EventFile, gzip.GzipFile):
608 """A way to read/write a gzipped lhef event"""
609
611 """A way to read a standard event file"""
612
614 """a class to read simultaneously multiple file and read them in mixing them.
615 Unweighting can be done at the same time.
616 The number of events in each file need to be provide in advance
617 (if not provide the file is first read to find that number"""
618
621
623 """if trunc_error is define here then this allow
624 to only read all the files twice and not three times."""
625 self.files = []
626 self.banner = ''
627 self.initial_nb_events = []
628 self.total_event_in_files = 0
629 self.curr_nb_events = []
630 self.allcross = []
631 self.error = []
632 self.across = []
633 self.scales = []
634 if start_list:
635 for p in start_list:
636 self.add(p)
637 self._configure = False
638
639 - def add(self, path, cross, error, across):
640 """ add a file to the pool, across allow to reweight the sum of weight
641 in the file to the given cross-section
642 """
643
644 if across == 0:
645
646 return
647
648 obj = EventFile(path)
649 if len(self.files) == 0 and not self.banner:
650 self.banner = obj.banner
651 self.curr_nb_events.append(0)
652 self.initial_nb_events.append(0)
653 self.allcross.append(cross)
654 self.across.append(across)
655 self.error.append(error)
656 self.scales.append(1)
657 self.files.append(obj)
658 self._configure = False
659
662
664
665 if not self._configure:
666 self.configure()
667
668 remaining_event = self.total_event_in_files - sum(self.curr_nb_events)
669 if remaining_event == 0:
670 raise StopIteration
671
672 nb_event = random.randint(1, remaining_event)
673 sum_nb=0
674 for i, obj in enumerate(self.files):
675 sum_nb += self.initial_nb_events[i] - self.curr_nb_events[i]
676 if nb_event <= sum_nb:
677 self.curr_nb_events[i] += 1
678 event = obj.next()
679 event.sample_scale = self.scales[i]
680 return event
681 else:
682 raise Exception
683
684
686 """define the part of the init_banner"""
687
688 if not self.banner:
689 return
690
691
692 grouped_cross = {}
693 grouped_error = {}
694 for i,ff in enumerate(self.files):
695 filename = ff.name
696 Pdir = [P for P in filename.split(os.path.sep) if P.startswith('P')][-1]
697 group = Pdir.split("_")[0][1:]
698 if group in grouped_cross:
699 grouped_cross[group] += self.allcross[i]
700 grouped_error[group] += self.error[i]**2
701 else:
702 grouped_cross[group] = self.allcross[i]
703 grouped_error[group] = self.error[i]**2
704
705 nb_group = len(grouped_cross)
706
707
708 try:
709 run_card = self.banner.run_card
710 except:
711 run_card = self.banner.charge_card("run_card")
712
713 init_information = run_card.get_banner_init_information()
714 init_information["nprup"] = nb_group
715
716 if run_card["lhe_version"] < 3:
717 init_information["generator_info"] = ""
718 else:
719 init_information["generator_info"] = "<generator name='MadGraph5_aMC@NLO' version='%s'>please cite 1405.0301 </generator>\n" \
720 % misc.get_pkg_info()['version']
721
722
723 cross_info = "%(cross)e %(error)e %(wgt)e %(id)i"
724 init_information["cross_info"] = []
725 for id in grouped_cross:
726 conv = {"id": int(id), "cross": grouped_cross[id], "error": math.sqrt(grouped_error[id]),
727 "wgt": wgt}
728 init_information["cross_info"].append( cross_info % conv)
729 init_information["cross_info"] = '\n'.join(init_information["cross_info"])
730
731
732
733 template_init =\
734 """ %(idbmup1)i %(idbmup2)i %(ebmup1)e %(ebmup2)e %(pdfgup1)i %(pdfgup2)i %(pdfsup1)i %(pdfsup2)i -3 %(nprup)i
735 %(cross_info)s
736 %(generator_info)s
737 """
738
739 self.banner["init"] = template_init % init_information
740
741
742
744 """ scan once the file to return
745 - the list of the hightest weight (of size trunc_error*NB_EVENT
746 - the cross-section by type of process
747 - the total number of events in the files
748 In top of that it initialise the information for the next routine
749 to determine how to choose which file to read
750 """
751 self.seek(0)
752 all_wgt = []
753 total_event = 0
754 sum_cross = collections.defaultdict(int)
755 for i,f in enumerate(self.files):
756 nb_event = 0
757
758
759 cross = collections.defaultdict(int)
760 new_wgt =[]
761 for event in f:
762 nb_event += 1
763 total_event += 1
764 event.sample_scale = 1
765 wgt = getwgt(event)
766 cross['all'] += wgt
767 cross['abs'] += abs(wgt)
768 cross[event.ievent] += wgt
769 new_wgt.append(abs(wgt))
770
771 if nb_event % 20000 == 0:
772 new_wgt.sort()
773
774 nb_keep = max(20, int(nb_event*trunc_error*15))
775 new_wgt = new_wgt[-nb_keep:]
776 if nb_event == 0:
777 raise Exception
778
779 self.initial_nb_events[i] = nb_event
780 self.scales[i] = self.across[i]/cross['abs'] if self.across[i] else 1
781
782
783 for key in cross:
784 sum_cross[key] += cross[key]* self.scales[i]
785 all_wgt +=[self.scales[i] * w for w in new_wgt]
786 all_wgt.sort()
787 nb_keep = max(20, int(total_event*trunc_error*10))
788 all_wgt = all_wgt[-nb_keep:]
789
790 self.total_event_in_files = total_event
791
792 all_wgt.sort()
793
794 nb_keep = max(20, int(total_event*trunc_error*10))
795 all_wgt = all_wgt[-nb_keep:]
796 self.seek(0)
797 self._configure = True
798 return all_wgt, sum_cross, total_event
799
806
808
809 return len(self.files)
810
811 - def seek(self, pos):
812 """ """
813
814 if pos !=0:
815 raise Exception
816 for i in range(len(self)):
817 self.curr_nb_events[i] = 0
818 for f in self.files:
819 f.seek(pos)
820
821 - def unweight(self, outputpath, get_wgt, **opts):
822 """unweight the current file according to wgt information wgt.
823 which can either be a fct of the event or a tag in the rwgt list.
824 max_wgt allow to do partial unweighting.
825 trunc_error allow for dynamical partial unweighting
826 event_target reweight for that many event with maximal trunc_error.
827 (stop to write event when target is reached)
828 """
829
830 if isinstance(get_wgt, str):
831 unwgt_name =get_wgt
832 def get_wgt_multi(event):
833 event.parse_reweight()
834 return event.reweight_data[unwgt_name] * event.sample_scale
835 else:
836 unwgt_name = get_wgt.func_name
837 get_wgt_multi = lambda event: get_wgt(event) * event.sample_scale
838
839
840
841 if 'event_target' in opts and opts['event_target']:
842 new_wgt = sum(self.across)/opts['event_target']
843 self.define_init_banner(new_wgt)
844 self.written_weight = new_wgt
845
846 return super(MultiEventFile, self).unweight(outputpath, get_wgt_multi, **opts)
847
850 """Class storing a single event information (list of particles + global information)"""
851
852 warning_order = True
853
855 """The initialization of an empty Event (or one associate to a text file)"""
856 list.__init__(self)
857
858
859 self.nexternal = 0
860 self.ievent = 0
861 self.wgt = 0
862 self.aqcd = 0
863 self.scale = 0
864 self.aqed = 0
865 self.aqcd = 0
866
867 self.tag = ''
868 self.eventflag = {}
869 self.comment = ''
870 self.reweight_data = {}
871 self.matched_scale_data = None
872 self.syscalc_data = {}
873 if text:
874 self.parse(text)
875
876
877
879 """Take the input file and create the structured information"""
880 text = re.sub(r'</?event>', '', text)
881 status = 'first'
882 for line in text.split('\n'):
883 line = line.strip()
884 if not line:
885 continue
886 if line.startswith('#'):
887 self.comment += '%s\n' % line
888 continue
889 if "<event" in line:
890 if '=' in line:
891 found = re.findall(r"""(\w*)=(?:(?:['"])([^'"]*)(?=['"])|(\S*))""",line)
892
893
894 self.eventflag = dict((n, a1) if a1 else (n,a2) for n,a1,a2 in found)
895
896 continue
897
898 if 'first' == status:
899 if '<rwgt>' in line:
900 status = 'tag'
901
902 if 'first' == status:
903 self.assign_scale_line(line)
904 status = 'part'
905 continue
906
907 if '<' in line:
908 status = 'tag'
909
910 if 'part' == status:
911 self.append(Particle(line, event=self))
912 else:
913 self.tag += '%s\n' % line
914
915 self.assign_mother()
916
918
919 for i,particle in enumerate(self):
920 if i < particle.mother1 or i < particle.mother2:
921 if self.warning_order:
922 logger.warning("Order of particle in the event did not agree with parent/child order. This might be problematic for some code.")
923 Event.warning_order = False
924 self.reorder_mother_child()
925 return self.assign_mother()
926
927 if particle.mother1:
928 try:
929 particle.mother1 = self[int(particle.mother1) -1]
930 except Exception:
931 logger.warning("WRONG MOTHER INFO %s", self)
932 particle.mother1 = 0
933 if particle.mother2:
934 try:
935 particle.mother2 = self[int(particle.mother2) -1]
936 except Exception:
937 logger.warning("WRONG MOTHER INFO %s", self)
938 particle.mother2 = 0
939
940
942 """check and correct the mother/child position.
943 only correct one order by call (but this is a recursive call)"""
944
945 tomove, position = None, None
946 for i,particle in enumerate(self):
947 if i < particle.mother1:
948
949 tomove, position = i, particle.mother1-1
950 break
951 if i < particle.mother2:
952 tomove, position = i, particle.mother2-1
953
954
955 if not tomove:
956 return
957
958
959 particle = self.pop(tomove)
960 self.insert(int(position), particle)
961
962
963 for i, particle in enumerate(self):
964 particle.event_id = i
965
966 m1, m2 = particle.mother1, particle.mother2
967 if m1 == tomove +1:
968 particle.mother1 = position+1
969 elif tomove < m1 <= position +1:
970 particle.mother1 -= 1
971 if m2 == tomove +1:
972 particle.mother2 = position+1
973 elif tomove < m2 <= position +1:
974 particle.mother2 -= 1
975
976 return self.reorder_mother_child()
977
978
979
980
981
982
984 """Parse the re-weight information in order to return a dictionary
985 {key: value}. If no group is define group should be '' """
986 if self.reweight_data:
987 return self.reweight_data
988 self.reweight_data = {}
989 self.reweight_order = []
990 start, stop = self.tag.find('<rwgt>'), self.tag.find('</rwgt>')
991 if start != -1 != stop :
992 pattern = re.compile(r'''<\s*wgt id=(?:\'|\")(?P<id>[^\'\"]+)(?:\'|\")\s*>\s*(?P<val>[\ded+-.]*)\s*</wgt>''')
993 data = pattern.findall(self.tag)
994 try:
995 self.reweight_data = dict([(pid, float(value)) for (pid, value) in data
996 if not self.reweight_order.append(pid)])
997
998 except ValueError, error:
999 raise Exception, 'Event File has unvalid weight. %s' % error
1000 self.tag = self.tag[:start] + self.tag[stop+7:]
1001 return self.reweight_data
1002
1004 """ """
1005 if hasattr(self, 'nloweight'):
1006 return self.nloweight
1007
1008 start, stop = self.tag.find('<mgrwgt>'), self.tag.find('</mgrwgt>')
1009 if start != -1 != stop :
1010
1011 text = self.tag[start+8:stop]
1012 self.nloweight = NLO_PARTIALWEIGHT(text, self)
1013
1014
1016 """Parse the line containing the starting scale for the shower"""
1017
1018 if self.matched_scale_data is not None:
1019 return self.matched_scale_data
1020
1021 self.matched_scale_data = []
1022
1023
1024 pattern = re.compile("<scales\s|</scales>")
1025 data = re.split(pattern,self.tag)
1026 if len(data) == 1:
1027 return []
1028 else:
1029 tmp = {}
1030 start,content, end = data
1031 self.tag = "%s%s" % (start, end)
1032 pattern = re.compile("pt_clust_(\d*)=\"([\de+-.]*)\"")
1033 for id,value in pattern.findall(content):
1034 tmp[int(id)] = float(value)
1035
1036 for i in range(1, len(tmp)+1):
1037 self.matched_scale_data.append(tmp[i])
1038
1039 return self.matched_scale_data
1040
1042 """ parse the flag for syscalc between <mgrwt></mgrwt>
1043 <mgrwt>
1044 <rscale> 3 0.26552898E+03</rscale>
1045 <asrwt>0</asrwt>
1046 <pdfrwt beam="1"> 1 21 0.14527945E+00 0.26552898E+03</pdfrwt>
1047 <pdfrwt beam="2"> 1 21 0.15249110E-01 0.26552898E+03</pdfrwt>
1048 <totfact> 0.10344054E+04</totfact>
1049 </mgrwt>
1050 """
1051 if self.syscalc_data:
1052 return self.syscalc_data
1053
1054 pattern = re.compile("<mgrwt>|</mgrwt>")
1055 pattern2 = re.compile("<(?P<tag>[\w]*)(?:\s*(\w*)=[\"'](.*)[\"']\s*|\s*)>(.*)</(?P=tag)>")
1056 data = re.split(pattern,self.tag)
1057 if len(data) == 1:
1058 return []
1059 else:
1060 tmp = {}
1061 start,content, end = data
1062 self.tag = "%s%s" % (start, end)
1063 for tag, key, keyval, tagval in pattern2.findall(content):
1064 if key:
1065 self.syscalc_data[(tag, key, keyval)] = tagval
1066 else:
1067 self.syscalc_data[tag] = tagval
1068 return self.syscalc_data
1069
1070
1071 - def add_decay_to_particle(self, position, decay_event):
1072 """define the decay of the particle id by the event pass in argument"""
1073
1074 this_particle = self[position]
1075
1076 this_particle.status = 2
1077 this_particle.helicity = 0
1078
1079
1080 decay_particle = decay_event[0]
1081 this_4mom = FourMomentum(this_particle)
1082 nb_part = len(self)
1083
1084 thres = decay_particle.E*1e-10
1085 assert max(decay_particle.px, decay_particle.py, decay_particle.pz) < thres,\
1086 "not on rest particle %s %s %s %s" % (decay_particle.E, decay_particle.px,decay_particle.py,decay_particle.pz)
1087
1088 self.nexternal += decay_event.nexternal -1
1089 old_scales = list(self.parse_matching_scale())
1090 if old_scales:
1091 jet_position = sum(1 for i in range(position) if self[i].status==1)
1092 self.matched_scale_data.pop(jet_position)
1093
1094
1095 for particle in decay_event[1:]:
1096
1097 new_particle = Particle(particle, self)
1098 new_particle.event_id = len(self)
1099 self.append(new_particle)
1100 if old_scales:
1101 self.matched_scale_data.append(old_scales[jet_position])
1102
1103 new_momentum = this_4mom.boost(FourMomentum(new_particle))
1104 new_particle.set_momentum(new_momentum)
1105
1106 for tag in ['mother1', 'mother2']:
1107 mother = getattr(particle, tag)
1108 if isinstance(mother, Particle):
1109 mother_id = getattr(particle, tag).event_id
1110 if mother_id == 0:
1111 setattr(new_particle, tag, this_particle)
1112 else:
1113 try:
1114 setattr(new_particle, tag, self[nb_part + mother_id -1])
1115 except Exception, error:
1116 print error
1117 misc.sprint( self)
1118 misc.sprint(nb_part + mother_id -1)
1119 misc.sprint(tag)
1120 misc.sprint(position, decay_event)
1121 misc.sprint(particle)
1122 misc.sprint(len(self), nb_part + mother_id -1)
1123 raise
1124 elif tag == "mother2" and isinstance(particle.mother1, Particle):
1125 new_particle.mother2 = this_particle
1126 else:
1127 raise Exception, "Something weird happens. Please report it for investigation"
1128
1129
1130 max_color=501
1131 for particle in self[:nb_part]:
1132 max_color=max(max_color, particle.color1, particle.color2)
1133
1134
1135 color_mapping = {}
1136 color_mapping[decay_particle.color1] = this_particle.color1
1137 color_mapping[decay_particle.color2] = this_particle.color2
1138 for particle in self[nb_part:]:
1139 if particle.color1:
1140 if particle.color1 not in color_mapping:
1141 max_color +=1
1142 color_mapping[particle.color1] = max_color
1143 particle.color1 = max_color
1144 else:
1145 particle.color1 = color_mapping[particle.color1]
1146 if particle.color2:
1147 if particle.color2 not in color_mapping:
1148 max_color +=1
1149 color_mapping[particle.color2] = max_color
1150 particle.color2 = max_color
1151 else:
1152 particle.color2 = color_mapping[particle.color2]
1153
1154
1155
1157
1158 to_remove = []
1159 if event_id is not None:
1160 to_remove.append(self[event_id])
1161
1162 if pdg_code:
1163 for particle in self:
1164 if particle.pid == pdg_code:
1165 to_remove.append(particle)
1166
1167 new_event = Event()
1168
1169 for tag in ['nexternal', 'ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']:
1170 setattr(new_event, tag, getattr(self, tag))
1171
1172 for particle in self:
1173 if isinstance(particle.mother1, Particle) and particle.mother1 in to_remove:
1174 to_remove.append(particle)
1175 if particle.status == 1:
1176 new_event.nexternal -= 1
1177 continue
1178 elif isinstance(particle.mother2, Particle) and particle.mother2 in to_remove:
1179 to_remove.append(particle)
1180 if particle.status == 1:
1181 new_event.nexternal -= 1
1182 continue
1183 else:
1184 new_event.append(Particle(particle))
1185
1186
1187
1188 for pos, particle in enumerate(new_event):
1189 particle.event_id = pos
1190 if particle in to_remove:
1191 particle.status = 1
1192 return new_event
1193
1194 - def get_decay(self, pdg_code=0, event_id=None):
1195
1196 to_start = []
1197 if event_id is not None:
1198 to_start.append(self[event_id])
1199
1200 elif pdg_code:
1201 for particle in self:
1202 if particle.pid == pdg_code:
1203 to_start.append(particle)
1204 break
1205
1206 new_event = Event()
1207
1208 for tag in ['ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']:
1209 setattr(new_event, tag, getattr(self, tag))
1210
1211
1212 old2new = {}
1213 new_decay_part = Particle(to_start[0])
1214 new_decay_part.mother1 = None
1215 new_decay_part.mother2 = None
1216 new_decay_part.status = -1
1217 old2new[new_decay_part.event_id] = len(old2new)
1218 new_event.append(new_decay_part)
1219
1220
1221
1222 for particle in self:
1223 if isinstance(particle.mother1, Particle) and particle.mother1.event_id in old2new\
1224 or isinstance(particle.mother2, Particle) and particle.mother2.event_id in old2new:
1225 old2new[particle.event_id] = len(old2new)
1226 new_event.append(Particle(particle))
1227
1228
1229
1230 nexternal = 0
1231 for pos, particle in enumerate(new_event):
1232 particle.event_id = pos
1233 if particle.mother1:
1234 particle.mother1 = new_event[old2new[particle.mother1.event_id]]
1235 if particle.mother2:
1236 particle.mother2 = new_event[old2new[particle.mother2.event_id]]
1237 if particle.status in [-1,1]:
1238 nexternal +=1
1239 new_event.nexternal = nexternal
1240
1241 return new_event
1242
1243
1245 """check various property of the events"""
1246
1247
1248 E, px, py, pz = 0,0,0,0
1249 absE, abspx, abspy, abspz = 0,0,0,0
1250 for particle in self:
1251 coeff = 1
1252 if particle.status == -1:
1253 coeff = -1
1254 elif particle.status != 1:
1255 continue
1256 E += coeff * particle.E
1257 absE += abs(particle.E)
1258 px += coeff * particle.px
1259 py += coeff * particle.py
1260 pz += coeff * particle.pz
1261 abspx += abs(particle.px)
1262 abspy += abs(particle.py)
1263 abspz += abs(particle.pz)
1264
1265 threshold = 5e-7
1266 if E/absE > threshold:
1267 logger.critical(self)
1268 raise Exception, "Do not conserve Energy %s, %s" % (E/absE, E)
1269 if px/abspx > threshold:
1270 logger.critical(self)
1271 raise Exception, "Do not conserve Px %s, %s" % (px/abspx, px)
1272 if py/abspy > threshold:
1273 logger.critical(self)
1274 raise Exception, "Do not conserve Py %s, %s" % (py/abspy, py)
1275 if pz/abspz > threshold:
1276 logger.critical(self)
1277 raise Exception, "Do not conserve Pz %s, %s" % (pz/abspz, pz)
1278
1279
1280 self.check_color_structure()
1281
1283 """read the line corresponding to global event line
1284 format of the line is:
1285 Nexternal IEVENT WEIGHT SCALE AEW AS
1286 """
1287 inputs = line.split()
1288 assert len(inputs) == 6
1289 self.nexternal=int(inputs[0])
1290 self.ievent=int(inputs[1])
1291 self.wgt=float(inputs[2])
1292 self.scale=float(inputs[3])
1293 self.aqed=float(inputs[4])
1294 self.aqcd=float(inputs[5])
1295
1297 """Return the unique tag identifying the SubProcesses for the generation.
1298 Usefull for program like MadSpin and Reweight module."""
1299
1300 initial, final, order = [], [], [[], []]
1301 for particle in self:
1302 if particle.status == -1:
1303 initial.append(particle.pid)
1304 order[0].append(particle.pid)
1305 elif particle.status == 1:
1306 final.append(particle.pid)
1307 order[1].append(particle.pid)
1308 initial.sort(), final.sort()
1309 tag = (tuple(initial), tuple(final))
1310 return tag, order
1311
1313 """return a list with the helicities in the order asked for"""
1314
1315
1316
1317
1318 order = [list(get_order[0]), list(get_order[1])]
1319 out = [9] *(len(order[0])+len(order[1]))
1320 for i, part in enumerate(self):
1321 if part.status == 1:
1322 try:
1323 ind = order[1].index(part.pid)
1324 except ValueError, error:
1325 if not allow_reversed:
1326 raise error
1327 else:
1328 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
1329 try:
1330 return self.get_helicity(order, False)
1331 except ValueError:
1332 raise error
1333 position = len(order[0]) + ind
1334 order[1][ind] = 0
1335 elif part.status == -1:
1336 try:
1337 ind = order[0].index(part.pid)
1338 except ValueError, error:
1339 if not allow_reversed:
1340 raise error
1341 else:
1342 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
1343 try:
1344 return self.get_helicity(order, False)
1345 except ValueError:
1346 raise error
1347
1348 position = ind
1349 order[0][ind] = 0
1350 else:
1351 continue
1352 out[position] = int(part.helicity)
1353 return out
1354
1355
1357 """check the validity of the color structure"""
1358
1359
1360 color_index = collections.defaultdict(int)
1361 for particle in self:
1362 if particle.status in [-1,1]:
1363 if particle.color1:
1364 color_index[particle.color1] +=1
1365 if -7 < particle.pdg < 0:
1366 raise Exception, "anti-quark with color tag"
1367 if particle.color2:
1368 color_index[particle.color2] +=1
1369 if 7 > particle.pdg > 0:
1370 raise Exception, "quark with anti-color tag"
1371
1372
1373 for key,value in color_index.items():
1374 if value > 2:
1375 print self
1376 print key, value
1377 raise Exception, 'Wrong color_flow'
1378
1379
1380
1381 check = []
1382 popup_index = []
1383 for particle in self:
1384 mothers = []
1385 childs = []
1386 if particle.mother1:
1387 mothers.append(particle.mother1)
1388 if particle.mother2 and particle.mother2 is not particle.mother1:
1389 mothers.append(particle.mother2)
1390 if not mothers:
1391 continue
1392 if (particle.mother1.event_id, particle.mother2.event_id) in check:
1393 continue
1394 check.append((particle.mother1.event_id, particle.mother2.event_id))
1395
1396 childs = [p for p in self if p.mother1 is particle.mother1 and \
1397 p.mother2 is particle.mother2]
1398
1399 mcolors = []
1400 manticolors = []
1401 for m in mothers:
1402 if m.color1:
1403 if m.color1 in manticolors:
1404 manticolors.remove(m.color1)
1405 else:
1406 mcolors.append(m.color1)
1407 if m.color2:
1408 if m.color2 in mcolors:
1409 mcolors.remove(m.color2)
1410 else:
1411 manticolors.append(m.color2)
1412 ccolors = []
1413 canticolors = []
1414 for m in childs:
1415 if m.color1:
1416 if m.color1 in canticolors:
1417 canticolors.remove(m.color1)
1418 else:
1419 ccolors.append(m.color1)
1420 if m.color2:
1421 if m.color2 in ccolors:
1422 ccolors.remove(m.color2)
1423 else:
1424 canticolors.append(m.color2)
1425 for index in mcolors[:]:
1426 if index in ccolors:
1427 mcolors.remove(index)
1428 ccolors.remove(index)
1429 for index in manticolors[:]:
1430 if index in canticolors:
1431 manticolors.remove(index)
1432 canticolors.remove(index)
1433
1434 if mcolors != []:
1435
1436 if len(canticolors) + len(mcolors) != 3:
1437 logger.critical(str(self))
1438 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs])
1439 else:
1440 popup_index += canticolors
1441 elif manticolors != []:
1442
1443 if len(ccolors) + len(manticolors) != 3:
1444 logger.critical(str(self))
1445 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs])
1446 else:
1447 popup_index += ccolors
1448
1449
1450 if len(popup_index) != len(set(popup_index)):
1451 logger.critical(self)
1452 raise Exception, "Wrong color flow: identical poping-up index, %s" % (popup_index)
1453
1455 """return a correctly formatted LHE event"""
1456
1457 out="""<event%(event_flag)s>
1458 %(scale)s
1459 %(particles)s
1460 %(comments)s
1461 %(tag)s
1462 %(reweight)s
1463 </event>
1464 """
1465 if event_id not in ['', None]:
1466 self.eventflag['event'] = str(event_id)
1467
1468 if self.eventflag:
1469 event_flag = ' %s' % ' '.join('%s="%s"' % (k,v) for (k,v) in self.eventflag.items())
1470 else:
1471 event_flag = ''
1472
1473 if self.nexternal:
1474 scale_str = "%2d %6d %+13.7e %14.8e %14.8e %14.8e" % \
1475 (self.nexternal,self.ievent,self.wgt,self.scale,self.aqed,self.aqcd)
1476 else:
1477 scale_str = ''
1478
1479 if self.reweight_data:
1480
1481 if set(self.reweight_data.keys()) != set(self.reweight_order):
1482 self.reweight_order += [k for k in self.reweight_data.keys() \
1483 if k not in self.reweight_order]
1484
1485 reweight_str = '<rwgt>\n%s\n</rwgt>' % '\n'.join(
1486 '<wgt id=\'%s\'> %+13.7e </wgt>' % (i, float(self.reweight_data[i]))
1487 for i in self.reweight_order)
1488 else:
1489 reweight_str = ''
1490
1491 tag_str = self.tag
1492 if self.matched_scale_data:
1493 tag_str = "<scales %s></scales>%s" % (
1494 ' '.join(['pt_clust_%i=\"%s\"' % (i,v)
1495 for i,v in enumerate(self.matched_scale_data)]),
1496 self.tag)
1497 if self.syscalc_data:
1498 keys= ['rscale', 'asrwt', ('pdfrwt', 'beam', '1'), ('pdfrwt', 'beam', '2'),
1499 'matchscale', 'totfact']
1500 sys_str = "<mgrwt>\n"
1501 template = """<%(key)s%(opts)s>%(values)s</%(key)s>\n"""
1502 for k in keys:
1503 if k not in self.syscalc_data:
1504 continue
1505 replace = {}
1506 replace['values'] = self.syscalc_data[k]
1507 if isinstance(k, str):
1508 replace['key'] = k
1509 replace['opts'] = ''
1510 else:
1511 replace['key'] = k[0]
1512 replace['opts'] = ' %s=\"%s\"' % (k[1],k[2])
1513 sys_str += template % replace
1514 sys_str += "</mgrwt>\n"
1515 reweight_str = sys_str + reweight_str
1516
1517 out = out % {'event_flag': event_flag,
1518 'scale': scale_str,
1519 'particles': '\n'.join([str(p) for p in self]),
1520 'tag': tag_str,
1521 'comments': self.comment,
1522 'reweight': reweight_str}
1523
1524 return re.sub('[\n]+', '\n', out)
1525
1526 - def get_momenta(self, get_order, allow_reversed=True):
1527 """return the momenta vector in the order asked for"""
1528
1529
1530 order = [list(get_order[0]), list(get_order[1])]
1531 out = [''] *(len(order[0])+len(order[1]))
1532 for i, part in enumerate(self):
1533 if part.status == 1:
1534 try:
1535 ind = order[1].index(part.pid)
1536 except ValueError, error:
1537 if not allow_reversed:
1538 raise error
1539 else:
1540 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
1541 try:
1542 return self.get_momenta_str(order, False)
1543 except ValueError:
1544 raise error
1545 position = len(order[0]) + ind
1546 order[1][ind] = 0
1547 elif part.status == -1:
1548 try:
1549 ind = order[0].index(part.pid)
1550 except ValueError, error:
1551 if not allow_reversed:
1552 raise error
1553 else:
1554 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
1555 try:
1556 return self.get_momenta_str(order, False)
1557 except ValueError:
1558 raise error
1559
1560 position = ind
1561 order[0][ind] = 0
1562 else:
1563 continue
1564
1565 out[position] = (part.E, part.px, part.py, part.pz)
1566
1567 return out
1568
1569
1570
1572
1573 scale = 0
1574 for particle in self:
1575 if particle.status != 1:
1576 continue
1577 scale += particle.mass**2 + particle.momentum.pt**2
1578
1579 return prefactor * scale
1580
1582 """return the momenta str in the order asked for"""
1583
1584 out = self.get_momenta(get_order, allow_reversed)
1585
1586 format = '%.12f'
1587 format_line = ' '.join([format]*4) + ' \n'
1588 out = [format_line % one for one in out]
1589 out = ''.join(out).replace('e','d')
1590 return out
1591
1593 """A class to allow to read both gzip and not gzip file.
1594 containing only weight from pythia --generated by SysCalc"""
1595
1596 - def __new__(self, path, mode='r', *args, **opt):
1597 if path.endswith(".gz"):
1598 try:
1599 return gzip.GzipFile.__new__(WeightFileGzip, path, mode, *args, **opt)
1600 except IOError, error:
1601 raise
1602 except Exception, error:
1603 if mode == 'r':
1604 misc.gunzip(path)
1605 return file.__new__(WeightFileNoGzip, path[:-3], mode, *args, **opt)
1606 else:
1607 return file.__new__(WeightFileNoGzip, path, mode, *args, **opt)
1608
1609
1610 - def __init__(self, path, mode='r', *args, **opt):
1611 """open file and read the banner [if in read mode]"""
1612
1613 super(EventFile, self).__init__(path, mode, *args, **opt)
1614 self.banner = ''
1615 if mode == 'r':
1616 line = ''
1617 while '</header>' not in line.lower():
1618 try:
1619 line = super(EventFile, self).next()
1620 except StopIteration:
1621 self.seek(0)
1622 self.banner = ''
1623 break
1624 if "<event" in line.lower():
1625 self.seek(0)
1626 self.banner = ''
1627 break
1628
1629 self.banner += line
1630
1634
1637
1640 """a convenient object for 4-momenta operation"""
1641
1642 - def __init__(self, obj=0, px=0, py=0, pz=0, E=0):
1643 """initialize the four momenta"""
1644
1645 if obj is 0 and E:
1646 obj = E
1647
1648 if isinstance(obj, (FourMomentum, Particle)):
1649 px = obj.px
1650 py = obj.py
1651 pz = obj.pz
1652 E = obj.E
1653 elif isinstance(obj, list):
1654 assert len(obj) ==4
1655 E = obj[0]
1656 px = obj[1]
1657 py = obj[2]
1658 pz = obj[3]
1659 elif isinstance(obj, str):
1660 obj = [float(i) for i in obj.split()]
1661 assert len(obj) ==4
1662 E = obj[0]
1663 px = obj[1]
1664 py = obj[2]
1665 pz = obj[3]
1666 else:
1667 E =obj
1668
1669
1670 self.E = float(E)
1671 self.px = float(px)
1672 self.py = float(py)
1673 self.pz = float(pz)
1674
1675 @property
1677 """return the mass"""
1678 return math.sqrt(self.E**2 - self.px**2 - self.py**2 - self.pz**2)
1679
1680 @property
1682 """return the mass square"""
1683 return self.E**2 - self.px**2 - self.py**2 - self.pz**2
1684
1685 @property
1687 return math.sqrt(max(0, self.pt2()))
1688
1689 @property
1691 norm = math.sqrt(self.px**2 + self.py**2+self.pz**2)
1692 return 0.5* math.log((norm - self.pz) / (norm + self.pz))
1693
1694 @property
1696 return 0.5* math.log((self.E +self.pz) / (self.E - self.pz))
1697
1698
1699
1701 """ return the pt square """
1702
1703 return self.px**2 + self.py**2
1704
1706
1707 assert isinstance(obj, FourMomentum)
1708 new = FourMomentum(self.E+obj.E,
1709 self.px + obj.px,
1710 self.py + obj.py,
1711 self.pz + obj.pz)
1712 return new
1713
1715 """update the object with the sum"""
1716 self.E += obj.E
1717 self.px += obj.px
1718 self.py += obj.py
1719 self.pz += obj.pz
1720 return self
1721
1723
1724 assert isinstance(obj, FourMomentum)
1725 new = FourMomentum(self.E-obj.E,
1726 self.px - obj.px,
1727 self.py - obj.py,
1728 self.pz - obj.pz)
1729 return new
1730
1732 """update the object with the sum"""
1733 self.E -= obj.E
1734 self.px -= obj.px
1735 self.py -= obj.py
1736 self.pz -= obj.pz
1737 return self
1738
1740 if isinstance(obj, FourMomentum):
1741 return self.E*obj.E - self.px *obj.px - self.py * obj.py - self.pz * obj.pz
1742 elif isinstance(obj, (float, int)):
1743 return FourMomentum(obj*self.E,obj*self.px,obj*self.py,obj*self.pz )
1744 else:
1745 raise NotImplemented
1746 __rmul__ = __mul__
1747
1749 assert power in [1,2]
1750
1751 if power == 1:
1752 return FourMomentum(self)
1753 elif power == 2:
1754 return self.mass_sqr()
1755
1757 return 'FourMomentum(%s,%s,%s,%s)' % (self.E, self.px, self.py,self.pz)
1758
1760 """mom 4-momenta is suppose to be given in the rest frame of this 4-momenta.
1761 the output is the 4-momenta in the frame of this 4-momenta
1762 function copied from HELAS routine."""
1763
1764
1765 pt = self.px**2 + self.py**2 + self.pz**2
1766 if pt:
1767 s3product = self.px * mom.px + self.py * mom.py + self.pz * mom.pz
1768 mass = self.mass
1769 lf = (mom.E + (self.E - mass) * s3product / pt ) / mass
1770 return FourMomentum(E=(self.E*mom.E+s3product)/mass,
1771 px=mom.px + self.px * lf,
1772 py=mom.py + self.py * lf,
1773 pz=mom.pz + self.pz * lf)
1774 else:
1775 return FourMomentum(mom)
1776
1779
1781 """ """
1782
1783 if isinstance(input, str):
1784 self.parse(input)
1785
1787 """parse the line and create the related object"""
1788
1789
1790
1791 data = text.split()
1792
1793
1794
1795 self.pwgt = [float(f) for f in data[:3]]
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813 self.born = float(data[3])
1814 self.real = float(data[4])
1815
1816 self.nexternal = int(data[5])
1817
1818 self.pdgs = [int(i) for i in data[6:6+self.nexternal]]
1819 flag = 6+self.nexternal
1820
1821 self.qcdpower = int(data[flag])
1822
1823 self.bjks = [float(f) for f in data[flag+1:flag+3]]
1824
1825 self.scales2 = [float(f) for f in data[flag+3:flag+6]]
1826
1827 self.gs = float(data[flag+6])
1828
1829
1830
1831
1832
1833
1834 self.born_related = int(data[flag+7])
1835 self.real_related = int(data[flag+8])
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858 self.type = int(data[flag+9])
1859
1860
1861
1862 self.nfks = int(data[flag+10])
1863
1864
1865
1866
1867
1868
1869 self.to_merge_pdg = [int (f) for f in data[flag+11:flag+13]]
1870
1871 self.merge_new_pdg = int(data[flag+13])
1872
1873
1874
1875
1876 self.ref_wgt = float(data[flag+14])
1877
1878
1879 if self.type in [1,11]:
1880 self.momenta_config = self.real_related
1881 else:
1882 self.momenta_config = self.born_related
1883
1886
1888
1889 - def __init__(self, momenta, wgts, event):
1890 list.__init__(self, momenta)
1891
1892 assert self
1893 self.wgts = wgts
1894 self.pdgs = list(wgts[0].pdgs)
1895 self.event = event
1896
1897 if wgts[0].momenta_config == wgts[0].born_related:
1898
1899 ind1, ind2 = [ind-1 for ind in wgts[0].to_merge_pdg]
1900 if ind1> ind2:
1901 ind1, ind2 = ind2, ind1
1902 if ind1 >= sum(1 for p in event if p.status==-1):
1903 new_p = self[ind1] + self[ind2]
1904 else:
1905 new_p = self[ind1] - self[ind2]
1906 self.pop(ind1)
1907 self.insert(ind1, new_p)
1908 self.pop(ind2)
1909 self.pdgs.pop(ind1)
1910 self.pdgs.insert(ind1, wgts[0].merge_new_pdg )
1911 self.pdgs.pop(ind2)
1912
1913 elif any(w.type in [1,11] for w in wgts):
1914 if any(w.type not in [1,11] for w in wgts):
1915 raise Exception
1916
1917 ind1, ind2 = [ind-1 for ind in wgts[0].to_merge_pdg]
1918 if ind1> ind2:
1919 ind1, ind2 = ind2, ind1
1920 if ind1 >= sum(1 for p in event if p.status==-1):
1921 new_p = self[ind1] + self[ind2]
1922 else:
1923 new_p = self[ind1] - self[ind2]
1924
1925 if __debug__:
1926 ptot = FourMomentum()
1927 for i in xrange(len(self)):
1928 if i <2:
1929 ptot += self[i]
1930 else:
1931 ptot -= self[i]
1932 if ptot.mass_sqr > 1e-16:
1933 misc.sprint(ptot, ptot.mass_sqr)
1934
1935 inv_mass = new_p.mass_sqr
1936 shat = (self[0]+self[1]).mass_sqr
1937 if (abs(inv_mass)/shat < 1e-6):
1938
1939 self.pop(ind1)
1940 self.insert(ind1, new_p)
1941 self.pop(ind2)
1942 self.pdgs.pop(ind1)
1943 self.pdgs.insert(ind1, wgts[0].merge_new_pdg )
1944 self.pdgs.pop(ind2)
1945
1946
1947 if __debug__:
1948 ptot = FourMomentum()
1949 for i in xrange(len(self)):
1950 if i <2:
1951 ptot += self[i]
1952 else:
1953 ptot -= self[i]
1954 if ptot.mass_sqr > 1e-16:
1955 misc.sprint(ptot, ptot.mass_sqr)
1956
1957
1960
1962 """ return the tag and order for this basic event"""
1963 (initial, _), _ = self.event.get_tag_and_order()
1964 order = self.get_pdg_code()
1965
1966
1967 initial, out = order[:len(initial)], order[len(initial):]
1968 initial.sort()
1969 out.sort()
1970 return (tuple(initial), tuple(out)), order
1971
1972 - def get_momenta(self, get_order, allow_reversed=True):
1973 """return the momenta vector in the order asked for"""
1974
1975
1976 order = [list(get_order[0]), list(get_order[1])]
1977 out = [''] *(len(order[0])+len(order[1]))
1978 pdgs = self.get_pdg_code()
1979 for pos, part in enumerate(self):
1980 if pos < len(get_order[0]):
1981 try:
1982 ind = order[0].index(pdgs[pos])
1983 except ValueError, error:
1984 if not allow_reversed:
1985 raise error
1986 else:
1987 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
1988 try:
1989 return self.get_momenta(order, False)
1990 except ValueError:
1991 raise error
1992
1993
1994 position = ind
1995 order[0][ind] = 0
1996 else:
1997 try:
1998 ind = order[1].index(pdgs[pos])
1999 except ValueError, error:
2000 if not allow_reversed:
2001 raise error
2002 else:
2003 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2004 try:
2005 return self.get_momenta(order, False)
2006 except ValueError:
2007 raise error
2008 position = len(order[0]) + ind
2009 order[1][ind] = 0
2010
2011 out[position] = (part.E, part.px, part.py, part.pz)
2012
2013 return out
2014
2015
2017 return [9] * len(self)
2018
2019 @property
2021 return self.event.aqcd
2022
2024
2025 self.event = event
2026 if isinstance(input, str):
2027 self.parse(input)
2028
2029
2031 """create the object from the string information (see example below)"""
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052 text = text.lower().replace('d','e')
2053 all_line = text.split('\n')
2054
2055 first_line =''
2056 while not first_line.strip():
2057 first_line = all_line.pop(0)
2058
2059 wgt, nb_wgt, nb_event, _ = first_line.split()
2060 nb_wgt, nb_event = int(nb_wgt), int(nb_event)
2061
2062 momenta = []
2063 wgts = []
2064 for line in all_line:
2065 data = line.split()
2066 if len(data) == 4:
2067 p = FourMomentum(data)
2068 momenta.append(p)
2069 elif len(data)>0:
2070 wgt = OneNLOWeight(line)
2071 wgts.append(wgt)
2072
2073 assert len(wgts) == int(nb_wgt)
2074
2075 get_weights_for_momenta = {}
2076 size_momenta = 0
2077 for wgt in wgts:
2078 if wgt.momenta_config in get_weights_for_momenta:
2079 get_weights_for_momenta[wgt.momenta_config].append(wgt)
2080 else:
2081 if size_momenta == 0: size_momenta = wgt.nexternal
2082 assert size_momenta == wgt.nexternal
2083 get_weights_for_momenta[wgt.momenta_config] = [wgt]
2084
2085 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)
2086
2087
2088
2089 self.cevents = []
2090 for key in range(1, nb_event+1):
2091 if key in get_weights_for_momenta:
2092 wgt = get_weights_for_momenta[key]
2093 evt = self.BasicEvent(momenta[:size_momenta], get_weights_for_momenta[key], self.event)
2094 self.cevents.append(evt)
2095 momenta = momenta[size_momenta:]
2096
2097 nb_wgt_check = 0
2098 for cevt in self.cevents:
2099 nb_wgt_check += len(cevt.wgts)
2100 assert nb_wgt_check == int(nb_wgt)
2101
2102
2103
2104 if '__main__' == __name__:
2105
2106
2107 if False:
2108 lhe = EventFile('unweighted_events.lhe.gz')
2109 output = open('output_events.lhe', 'w')
2110
2111 output.write(lhe.banner)
2112
2113 for event in lhe:
2114 for particle in event:
2115
2116 particle.mass = 0
2117 particle.vtim = 2
2118
2119
2120 output.write(str(event))
2121 output.write('</LesHouchesEvent>\n')
2122
2123
2124 if False:
2125 lhe = EventFile('unweighted_events.lhe.gz')
2126 import matplotlib.pyplot as plt
2127 import matplotlib.gridspec as gridspec
2128 nbins = 100
2129
2130 nb_pass = 0
2131 data = []
2132 for event in lhe:
2133 etaabs = 0
2134 etafinal = 0
2135 for particle in event:
2136 if particle.status==1:
2137 p = FourMomentum(particle)
2138 eta = p.pseudorapidity
2139 if abs(eta) > etaabs:
2140 etafinal = eta
2141 etaabs = abs(eta)
2142 if etaabs < 4:
2143 data.append(etafinal)
2144 nb_pass +=1
2145
2146
2147 print nb_pass
2148 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1])
2149 gs1.update(wspace=0, hspace=0)
2150 ax = plt.subplot(gs1[0])
2151
2152 n, bins, patches = ax.hist(data, nbins, histtype='step', label='original')
2153 ax_c = ax.twinx()
2154 ax_c.set_ylabel('MadGraph5_aMC@NLO')
2155 ax_c.yaxis.set_label_coords(1.01, 0.25)
2156 ax_c.set_yticks(ax.get_yticks())
2157 ax_c.set_yticklabels([])
2158 ax.set_xlim([-4,4])
2159 print "bin value:", n
2160 print "start/end point of bins", bins
2161 plt.axis('on')
2162 plt.xlabel('weight ratio')
2163 plt.show()
2164
2165
2166
2167 if False:
2168 lhe = EventFile('unweighted_events.lhe')
2169 import matplotlib.pyplot as plt
2170 import matplotlib.gridspec as gridspec
2171 nbins = 100
2172
2173
2174 mtau, wtau = 1.777, 4.027000e-13
2175 nb_pass = 0
2176 data, data2, data3 = [], [], []
2177 for event in lhe:
2178 nb_pass +=1
2179 if nb_pass > 10000:
2180 break
2181 tau1 = FourMomentum()
2182 tau2 = FourMomentum()
2183 for part in event:
2184 if part.pid in [-12,11,16]:
2185 momenta = FourMomentum(part)
2186 tau1 += momenta
2187 elif part.pid == 15:
2188 tau2 += FourMomentum(part)
2189
2190 if abs((mtau-tau2.mass())/wtau)<1e6 and tau2.mass() >1:
2191 data.append((tau1.mass()-mtau)/wtau)
2192 data2.append((tau2.mass()-mtau)/wtau)
2193 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1])
2194 gs1.update(wspace=0, hspace=0)
2195 ax = plt.subplot(gs1[0])
2196
2197 n, bins, patches = ax.hist(data2, nbins, histtype='step', label='original')
2198 n2, bins2, patches2 = ax.hist(data, bins=bins, histtype='step',label='reconstructed')
2199 import cmath
2200
2201 breit = lambda m : math.sqrt(4*math.pi)*1/(((m)**2-mtau**2)**2+(mtau*wtau)**2)*wtau
2202
2203 data3 = [breit(mtau + x*wtau)*wtau*16867622.6624*50 for x in bins]
2204
2205 ax.plot(bins, data3,label='breit-wigner')
2206
2207 ax.legend()
2208
2209 ax_c = ax.twinx()
2210 ax_c.set_ylabel('MadGraph5_aMC@NLO')
2211 ax_c.yaxis.set_label_coords(1.01, 0.25)
2212 ax_c.set_yticks(ax.get_yticks())
2213 ax_c.set_yticklabels([])
2214
2215 plt.title('invariant mass of tau LHE/reconstructed')
2216 plt.axis('on')
2217 ax.set_xticklabels([])
2218
2219 ax = plt.subplot(gs1[1])
2220 data4 = [n[i]/(data3[i]) for i in range(nbins)]
2221 ax.plot(bins, data4 + [0] , 'b')
2222 data4 = [n2[i]/(data3[i]) for i in range(nbins)]
2223 ax.plot(bins, data4 + [0] , 'g')
2224 ax.set_ylim([0,2])
2225
2226 tick = ax.get_yticks()
2227 ax.set_yticks(tick[:-1])
2228
2229
2230 plt.axis('on')
2231 plt.xlabel('(M - Mtau)/Wtau')
2232 plt.show()
2233