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
405 if abs(wgt) > max_wgt:
406 trunc_cross += abs(wgt) - max_wgt
407 if event_target ==0 or nb_keep <= event_target:
408 if outputpath:
409 outfile.write(str(event))
410
411 elif wgt < 0:
412 nb_keep += 1
413 event.wgt = -1 * max(abs(wgt), max_wgt)
414 if abs(wgt) > max_wgt:
415 trunc_cross += abs(wgt) - max_wgt
416 if outputpath and (event_target ==0 or nb_keep <= event_target):
417 outfile.write(str(event))
418
419 if event_target and nb_keep > event_target:
420 if not outputpath:
421
422 continue
423 elif event_target and i != nb_try-1 and nb_keep >= event_target *1.05:
424 outfile.close()
425
426 continue
427 else:
428 outfile.write("</LesHouchesEvents>\n")
429 outfile.close()
430 break
431 elif event_target == 0:
432 if outputpath:
433 outfile.write("</LesHouchesEvents>\n")
434 outfile.close()
435 break
436 elif outputpath:
437 outfile.close()
438
439
440 else:
441
442 logger.log(log_level+10,"fail to reach target event %s (iteration=%s)", event_target,i)
443
444
445
446
447 if event_target:
448 nb_events_unweighted = nb_keep
449 nb_keep = min( event_target, nb_keep)
450 else:
451 nb_events_unweighted = nb_keep
452
453 logger.log(log_level, "write %i event (efficiency %.2g %%, truncation %.2g %%) after %i iteration(s)",
454 nb_keep, nb_events_unweighted/nb_event*100, trunc_cross/cross['abs']*100, i)
455
456
457 if nb_keep != event_target and hasattr(self, "written_weight"):
458 written_weight = lambda x: math.copysign(self.written_weight*event_target/nb_keep, float(x))
459 startfile = EventFile(outputpath)
460 tmpname = pjoin(os.path.dirname(outputpath), "wgtcorrected_"+ os.path.basename(outputpath))
461 outfile = EventFile(tmpname, "w")
462 outfile.write(startfile.banner)
463 for event in startfile:
464 event.wgt = written_weight(event.wgt)
465 outfile.write(str(event))
466 outfile.write("</LesHouchesEvents>\n")
467 startfile.close()
468 outfile.close()
469 shutil.move(tmpname, outputpath)
470
471
472 self.max_wgt = max_wgt
473 return nb_keep
474
476 """ apply one or more fct on all event. """
477
478 opt= {"print_step": 2000, "maxevent":float("inf"),'no_output':False}
479 opt.update(opts)
480
481 nb_fct = len(fcts)
482 out = []
483 for i in range(nb_fct):
484 out.append([])
485 self.seek(0)
486 nb_event = 0
487 for event in self:
488 nb_event += 1
489 if opt["print_step"] and nb_event % opt["print_step"] == 0:
490 if hasattr(self,"len"):
491 logger.info("currently at %s/%s event" % (nb_event, self.len))
492 else:
493 logger.info("currently at %s event" % nb_event)
494 for i in range(nb_fct):
495 if opts['no_output']:
496 fcts[i](event)
497 else:
498 out[i].append(fcts[i](event))
499 if nb_event > opt['maxevent']:
500 break
501 if nb_fct == 1:
502 return out[0]
503 else:
504 return out
505
506
507 - def update_Hwu(self, hwu, fct, name='lhe', keep_wgt=True):
508
509 first=True
510 def add_to_Hwu(event):
511 """function to update the HwU on the flight"""
512
513 value = fct(event)
514
515
516 if first:
517
518 if isinstance(value, dict):
519 hwu.add_line(value.keys())
520 else:
521 hwu.add_line(name)
522 if keep_wgt is True:
523 hwu.add_line(['%s_%s' % (name, key)
524 for key in event.reweight_data])
525 first = False
526
527 if isinstance(value, dict):
528 hwu.addEvent(value)
529 else:
530 hwu.addEvent({name:value})
531 if keep_wgt:
532 event.parse_reweight()
533 if keep_wgt is True:
534 data = dict(('%s_%s' % (name, key),value)
535 for key in event.reweight_data)
536 hwu.addEvent(data)
537
538
539
540 self.apply_fct_on_event(add_to_Hwu, no_output=True)
541 return hwu
542
544 """take the lhe file and add the matchscale from the pythia_input file"""
545
546 if pythia_input:
547 def next_data():
548 for line in open(pythia_input):
549 if line.startswith('#'):
550 continue
551 data = line.split()
552 print (int(data[0]), data[-3], data[-2], data[-1])
553 yield (int(data[0]), data[-3], data[-2], data[-1])
554 else:
555 def next_data():
556 i=0
557 while 1:
558 yield [i,0,0,0]
559 i+=1
560 sys_iterator = next_data()
561
562 self.seek(0)
563 out = open(out_path,'w')
564
565 pdf_pattern = re.compile(r'''<init>(.*)</init>''', re.M+re.S)
566 init = pdf_pattern.findall(self.banner)[0].split('\n',2)[1]
567 id1, id2, _, _, _, _, pdf1,pdf2,_,_ = init.split()
568 id = [int(id1), int(id2)]
569 type = []
570 for i in range(2):
571 if abs(id[i]) == 2212:
572 if i > 0:
573 type.append(1)
574 else:
575 type.append(-1)
576 else:
577 type.append(0)
578 pdf = max(int(pdf1),int(pdf2))
579
580 out.write("<header>\n" + \
581 "<orgpdf>%i</orgpdf>\n" % pdf + \
582 "<beams> %s %s</beams>\n" % tuple(type) + \
583 "</header>\n")
584
585
586 nevt, smin, smax, scomp = sys_iterator.next()
587 for i, orig_event in enumerate(self):
588 if i < nevt:
589 continue
590 new_event = Event()
591 sys = orig_event.parse_syscalc_info()
592 new_event.syscalc_data = sys
593 if smin:
594 new_event.syscalc_data['matchscale'] = "%s %s %s" % (smin, scomp, smax)
595 out.write(str(new_event), nevt)
596 try:
597 nevt, smin, smax, scomp = sys_iterator.next()
598 except StopIteration:
599 break
600
601
602
603
604
605
606
607
608 -class EventFileGzip(EventFile, gzip.GzipFile):
609 """A way to read/write a gzipped lhef event"""
610
612 """A way to read a standard event file"""
613
615 """a class to read simultaneously multiple file and read them in mixing them.
616 Unweighting can be done at the same time.
617 The number of events in each file need to be provide in advance
618 (if not provide the file is first read to find that number"""
619
622
624 """if trunc_error is define here then this allow
625 to only read all the files twice and not three times."""
626 self.files = []
627 self.banner = ''
628 self.initial_nb_events = []
629 self.total_event_in_files = 0
630 self.curr_nb_events = []
631 self.allcross = []
632 self.error = []
633 self.across = []
634 self.scales = []
635 if start_list:
636 for p in start_list:
637 self.add(p)
638 self._configure = False
639
640 - def add(self, path, cross, error, across):
641 """ add a file to the pool, across allow to reweight the sum of weight
642 in the file to the given cross-section
643 """
644
645 if across == 0:
646
647 return
648
649 obj = EventFile(path)
650 if len(self.files) == 0 and not self.banner:
651 self.banner = obj.banner
652 self.curr_nb_events.append(0)
653 self.initial_nb_events.append(0)
654 self.allcross.append(cross)
655 self.across.append(across)
656 self.error.append(error)
657 self.scales.append(1)
658 self.files.append(obj)
659 self._configure = False
660
663
665
666 if not self._configure:
667 self.configure()
668
669 remaining_event = self.total_event_in_files - sum(self.curr_nb_events)
670 if remaining_event == 0:
671 raise StopIteration
672
673 nb_event = random.randint(1, remaining_event)
674 sum_nb=0
675 for i, obj in enumerate(self.files):
676 sum_nb += self.initial_nb_events[i] - self.curr_nb_events[i]
677 if nb_event <= sum_nb:
678 self.curr_nb_events[i] += 1
679 event = obj.next()
680 event.sample_scale = self.scales[i]
681 return event
682 else:
683 raise Exception
684
685
687 """define the part of the init_banner"""
688
689 if not self.banner:
690 return
691
692
693 grouped_cross = {}
694 grouped_error = {}
695 for i,ff in enumerate(self.files):
696 filename = ff.name
697 Pdir = [P for P in filename.split(os.path.sep) if P.startswith('P')][-1]
698 group = Pdir.split("_")[0][1:]
699 if group in grouped_cross:
700 grouped_cross[group] += self.allcross[i]
701 grouped_error[group] += self.error[i]**2
702 else:
703 grouped_cross[group] = self.allcross[i]
704 grouped_error[group] = self.error[i]**2
705
706 nb_group = len(grouped_cross)
707
708
709 try:
710 run_card = self.banner.run_card
711 except:
712 run_card = self.banner.charge_card("run_card")
713
714 init_information = run_card.get_banner_init_information()
715 init_information["nprup"] = nb_group
716
717 if run_card["lhe_version"] < 3:
718 init_information["generator_info"] = ""
719 else:
720 init_information["generator_info"] = "<generator name='MadGraph5_aMC@NLO' version='%s'>please cite 1405.0301 </generator>\n" \
721 % misc.get_pkg_info()['version']
722
723
724 cross_info = "%(cross)e %(error)e %(wgt)e %(id)i"
725 init_information["cross_info"] = []
726 for id in grouped_cross:
727 conv = {"id": int(id), "cross": grouped_cross[id], "error": math.sqrt(grouped_error[id]),
728 "wgt": wgt}
729 init_information["cross_info"].append( cross_info % conv)
730 init_information["cross_info"] = '\n'.join(init_information["cross_info"])
731
732
733
734 template_init =\
735 """ %(idbmup1)i %(idbmup2)i %(ebmup1)e %(ebmup2)e %(pdfgup1)i %(pdfgup2)i %(pdfsup1)i %(pdfsup2)i -3 %(nprup)i
736 %(cross_info)s
737 %(generator_info)s
738 """
739
740 self.banner["init"] = template_init % init_information
741
742
743
745 """ scan once the file to return
746 - the list of the hightest weight (of size trunc_error*NB_EVENT
747 - the cross-section by type of process
748 - the total number of events in the files
749 In top of that it initialise the information for the next routine
750 to determine how to choose which file to read
751 """
752 self.seek(0)
753 all_wgt = []
754 total_event = 0
755 sum_cross = collections.defaultdict(int)
756 for i,f in enumerate(self.files):
757 nb_event = 0
758
759
760 cross = collections.defaultdict(int)
761 new_wgt =[]
762 for event in f:
763 nb_event += 1
764 total_event += 1
765 event.sample_scale = 1
766 wgt = getwgt(event)
767 cross['all'] += wgt
768 cross['abs'] += abs(wgt)
769 cross[event.ievent] += wgt
770 new_wgt.append(abs(wgt))
771
772 if nb_event % 20000 == 0:
773 new_wgt.sort()
774
775 nb_keep = max(20, int(nb_event*trunc_error*15))
776 new_wgt = new_wgt[-nb_keep:]
777 if nb_event == 0:
778 raise Exception
779
780 self.initial_nb_events[i] = nb_event
781 self.scales[i] = self.across[i]/cross['abs'] if self.across[i] else 1
782
783
784 for key in cross:
785 sum_cross[key] += cross[key]* self.scales[i]
786 all_wgt +=[self.scales[i] * w for w in new_wgt]
787 all_wgt.sort()
788 nb_keep = max(20, int(total_event*trunc_error*10))
789 all_wgt = all_wgt[-nb_keep:]
790
791 self.total_event_in_files = total_event
792
793 all_wgt.sort()
794
795 nb_keep = max(20, int(total_event*trunc_error*10))
796 all_wgt = all_wgt[-nb_keep:]
797 self.seek(0)
798 self._configure = True
799 return all_wgt, sum_cross, total_event
800
807
809
810 return len(self.files)
811
812 - def seek(self, pos):
813 """ """
814
815 if pos !=0:
816 raise Exception
817 for i in range(len(self)):
818 self.curr_nb_events[i] = 0
819 for f in self.files:
820 f.seek(pos)
821
822 - def unweight(self, outputpath, get_wgt, **opts):
823 """unweight the current file according to wgt information wgt.
824 which can either be a fct of the event or a tag in the rwgt list.
825 max_wgt allow to do partial unweighting.
826 trunc_error allow for dynamical partial unweighting
827 event_target reweight for that many event with maximal trunc_error.
828 (stop to write event when target is reached)
829 """
830
831 if isinstance(get_wgt, str):
832 unwgt_name =get_wgt
833 def get_wgt_multi(event):
834 event.parse_reweight()
835 return event.reweight_data[unwgt_name] * event.sample_scale
836 else:
837 unwgt_name = get_wgt.func_name
838 get_wgt_multi = lambda event: get_wgt(event) * event.sample_scale
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
1695 """ return the pt square """
1696
1697 return self.px**2 + self.py**2
1698
1700
1701 assert isinstance(obj, FourMomentum)
1702 new = FourMomentum(self.E+obj.E,
1703 self.px + obj.px,
1704 self.py + obj.py,
1705 self.pz + obj.pz)
1706 return new
1707
1709 """update the object with the sum"""
1710 self.E += obj.E
1711 self.px += obj.px
1712 self.py += obj.py
1713 self.pz += obj.pz
1714 return self
1715
1717
1718 assert isinstance(obj, FourMomentum)
1719 new = FourMomentum(self.E-obj.E,
1720 self.px - obj.px,
1721 self.py - obj.py,
1722 self.pz - obj.pz)
1723 return new
1724
1726 """update the object with the sum"""
1727 self.E -= obj.E
1728 self.px -= obj.px
1729 self.py -= obj.py
1730 self.pz -= obj.pz
1731 return self
1732
1734 if isinstance(obj, FourMomentum):
1735 return self.E*obj.E - self.px *obj.px - self.py * obj.py - self.pz * obj.pz
1736 elif isinstance(obj, (float, int)):
1737 return FourMomentum(obj*self.E,obj*self.px,obj*self.py,obj*self.pz )
1738 else:
1739 raise NotImplemented
1740 __rmul__ = __mul__
1741
1743 assert power in [1,2]
1744
1745 if power == 1:
1746 return FourMomentum(self)
1747 elif power == 2:
1748 return self.mass_sqr()
1749
1751 return 'FourMomentum(%s,%s,%s,%s)' % (self.E, self.px, self.py,self.pz)
1752
1754 """mom 4-momenta is suppose to be given in the rest frame of this 4-momenta.
1755 the output is the 4-momenta in the frame of this 4-momenta
1756 function copied from HELAS routine."""
1757
1758
1759 pt = self.px**2 + self.py**2 + self.pz**2
1760 if pt:
1761 s3product = self.px * mom.px + self.py * mom.py + self.pz * mom.pz
1762 mass = self.mass
1763 lf = (mom.E + (self.E - mass) * s3product / pt ) / mass
1764 return FourMomentum(E=(self.E*mom.E+s3product)/mass,
1765 px=mom.px + self.px * lf,
1766 py=mom.py + self.py * lf,
1767 pz=mom.pz + self.pz * lf)
1768 else:
1769 return FourMomentum(mom)
1770
1773
1775 """ """
1776
1777 if isinstance(input, str):
1778 self.parse(input)
1779
1781 """parse the line and create the related object"""
1782
1783
1784
1785 data = text.split()
1786
1787
1788
1789 self.pwgt = [float(f) for f in data[:3]]
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807 self.born = float(data[3])
1808 self.real = float(data[4])
1809
1810 self.nexternal = int(data[5])
1811
1812 self.pdgs = [int(i) for i in data[6:6+self.nexternal]]
1813 flag = 6+self.nexternal
1814
1815 self.qcdpower = int(data[flag])
1816
1817 self.bjks = [float(f) for f in data[flag+1:flag+3]]
1818
1819 self.scales2 = [float(f) for f in data[flag+3:flag+6]]
1820
1821 self.gs = float(data[flag+6])
1822
1823
1824
1825
1826
1827
1828 self.born_related = int(data[flag+7])
1829 self.real_related = int(data[flag+8])
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852 self.type = int(data[flag+9])
1853
1854
1855
1856 self.nfks = int(data[flag+10])
1857
1858
1859
1860
1861
1862
1863 self.to_merge_pdg = [int (f) for f in data[flag+11:flag+13]]
1864
1865 self.merge_new_pdg = int(data[flag+13])
1866
1867
1868
1869
1870 self.ref_wgt = float(data[flag+14])
1871
1872
1873 if self.type in [1,11]:
1874 self.momenta_config = self.real_related
1875 else:
1876 self.momenta_config = self.born_related
1877
1880
1882
1883 - def __init__(self, momenta, wgts, event):
1884 list.__init__(self, momenta)
1885
1886 assert self
1887 self.wgts = wgts
1888 self.pdgs = list(wgts[0].pdgs)
1889 self.event = event
1890
1891 if wgts[0].momenta_config == wgts[0].born_related:
1892
1893 ind1, ind2 = [ind-1 for ind in wgts[0].to_merge_pdg]
1894 if ind1> ind2:
1895 ind1, ind2 = ind2, ind1
1896 if ind1 >= sum(1 for p in event if p.status==-1):
1897 new_p = self[ind1] + self[ind2]
1898 else:
1899 new_p = self[ind1] - self[ind2]
1900 self.pop(ind1)
1901 self.insert(ind1, new_p)
1902 self.pop(ind2)
1903 self.pdgs.pop(ind1)
1904 self.pdgs.insert(ind1, wgts[0].merge_new_pdg )
1905 self.pdgs.pop(ind2)
1906
1907 elif any(w.type in [1,11] for w in wgts):
1908 if any(w.type not in [1,11] for w in wgts):
1909 raise Exception
1910
1911 ind1, ind2 = [ind-1 for ind in wgts[0].to_merge_pdg]
1912 if ind1> ind2:
1913 ind1, ind2 = ind2, ind1
1914 if ind1 >= sum(1 for p in event if p.status==-1):
1915 new_p = self[ind1] + self[ind2]
1916 else:
1917 new_p = self[ind1] - self[ind2]
1918
1919 if __debug__:
1920 ptot = FourMomentum()
1921 for i in xrange(len(self)):
1922 if i <2:
1923 ptot += self[i]
1924 else:
1925 ptot -= self[i]
1926 if ptot.mass_sqr > 1e-16:
1927 misc.sprint(ptot, ptot.mass_sqr)
1928
1929 inv_mass = new_p.mass_sqr
1930 shat = (self[0]+self[1]).mass_sqr
1931 if (abs(inv_mass)/shat < 1e-6):
1932
1933 self.pop(ind1)
1934 self.insert(ind1, new_p)
1935 self.pop(ind2)
1936 self.pdgs.pop(ind1)
1937 self.pdgs.insert(ind1, wgts[0].merge_new_pdg )
1938 self.pdgs.pop(ind2)
1939
1940
1941 if __debug__:
1942 ptot = FourMomentum()
1943 for i in xrange(len(self)):
1944 if i <2:
1945 ptot += self[i]
1946 else:
1947 ptot -= self[i]
1948 if ptot.mass_sqr > 1e-16:
1949 misc.sprint(ptot, ptot.mass_sqr)
1950
1951
1954
1956 """ return the tag and order for this basic event"""
1957 (initial, _), _ = self.event.get_tag_and_order()
1958 order = self.get_pdg_code()
1959
1960
1961 initial, out = order[:len(initial)], order[len(initial):]
1962 initial.sort()
1963 out.sort()
1964 return (tuple(initial), tuple(out)), order
1965
1966 - def get_momenta(self, get_order, allow_reversed=True):
1967 """return the momenta vector in the order asked for"""
1968
1969
1970 order = [list(get_order[0]), list(get_order[1])]
1971 out = [''] *(len(order[0])+len(order[1]))
1972 pdgs = self.get_pdg_code()
1973 for pos, part in enumerate(self):
1974 if pos < len(get_order[0]):
1975 try:
1976 ind = order[0].index(pdgs[pos])
1977 except ValueError, error:
1978 if not allow_reversed:
1979 raise error
1980 else:
1981 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
1982 try:
1983 return self.get_momenta(order, False)
1984 except ValueError:
1985 raise error
1986
1987
1988 position = ind
1989 order[0][ind] = 0
1990 else:
1991 try:
1992 ind = order[1].index(pdgs[pos])
1993 except ValueError, error:
1994 if not allow_reversed:
1995 raise error
1996 else:
1997 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
1998 try:
1999 return self.get_momenta(order, False)
2000 except ValueError:
2001 raise error
2002 position = len(order[0]) + ind
2003 order[1][ind] = 0
2004
2005 out[position] = (part.E, part.px, part.py, part.pz)
2006
2007 return out
2008
2009
2011 return [9] * len(self)
2012
2013 @property
2015 return self.event.aqcd
2016
2018
2019 self.event = event
2020 if isinstance(input, str):
2021 self.parse(input)
2022
2023
2025 """create the object from the string information (see example below)"""
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
2045
2046 text = text.lower().replace('d','e')
2047 all_line = text.split('\n')
2048
2049 first_line =''
2050 while not first_line.strip():
2051 first_line = all_line.pop(0)
2052
2053 wgt, nb_wgt, nb_event, _ = first_line.split()
2054 nb_wgt, nb_event = int(nb_wgt), int(nb_event)
2055
2056 momenta = []
2057 wgts = []
2058 for line in all_line:
2059 data = line.split()
2060 if len(data) == 4:
2061 p = FourMomentum(data)
2062 momenta.append(p)
2063 elif len(data)>0:
2064 wgt = OneNLOWeight(line)
2065 wgts.append(wgt)
2066
2067 assert len(wgts) == int(nb_wgt)
2068
2069 get_weights_for_momenta = {}
2070 size_momenta = 0
2071 for wgt in wgts:
2072 if wgt.momenta_config in get_weights_for_momenta:
2073 get_weights_for_momenta[wgt.momenta_config].append(wgt)
2074 else:
2075 if size_momenta == 0: size_momenta = wgt.nexternal
2076 assert size_momenta == wgt.nexternal
2077 get_weights_for_momenta[wgt.momenta_config] = [wgt]
2078
2079 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)
2080
2081
2082
2083 self.cevents = []
2084 for key in range(1, nb_event+1):
2085 if key in get_weights_for_momenta:
2086 wgt = get_weights_for_momenta[key]
2087 evt = self.BasicEvent(momenta[:size_momenta], get_weights_for_momenta[key], self.event)
2088 self.cevents.append(evt)
2089 momenta = momenta[size_momenta:]
2090
2091 nb_wgt_check = 0
2092 for cevt in self.cevents:
2093 nb_wgt_check += len(cevt.wgts)
2094 assert nb_wgt_check == int(nb_wgt)
2095
2096
2097
2098 if '__main__' == __name__:
2099
2100
2101 if False:
2102 lhe = EventFile('unweighted_events.lhe.gz')
2103 output = open('output_events.lhe', 'w')
2104
2105 output.write(lhe.banner)
2106
2107 for event in lhe:
2108 for particle in event:
2109
2110 particle.mass = 0
2111 particle.vtim = 2
2112
2113
2114 output.write(str(event))
2115 output.write('</LesHouchesEvent>\n')
2116
2117
2118 if False:
2119 lhe = EventFile('unweighted_events.lhe.gz')
2120 import matplotlib.pyplot as plt
2121 import matplotlib.gridspec as gridspec
2122 nbins = 100
2123
2124 nb_pass = 0
2125 data = []
2126 for event in lhe:
2127 etaabs = 0
2128 etafinal = 0
2129 for particle in event:
2130 if particle.status==1:
2131 p = FourMomentum(particle)
2132 eta = p.pseudorapidity
2133 if abs(eta) > etaabs:
2134 etafinal = eta
2135 etaabs = abs(eta)
2136 if etaabs < 4:
2137 data.append(etafinal)
2138 nb_pass +=1
2139
2140
2141 print nb_pass
2142 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1])
2143 gs1.update(wspace=0, hspace=0)
2144 ax = plt.subplot(gs1[0])
2145
2146 n, bins, patches = ax.hist(data, nbins, histtype='step', label='original')
2147 ax_c = ax.twinx()
2148 ax_c.set_ylabel('MadGraph5_aMC@NLO')
2149 ax_c.yaxis.set_label_coords(1.01, 0.25)
2150 ax_c.set_yticks(ax.get_yticks())
2151 ax_c.set_yticklabels([])
2152 ax.set_xlim([-4,4])
2153 print "bin value:", n
2154 print "start/end point of bins", bins
2155 plt.axis('on')
2156 plt.xlabel('weight ratio')
2157 plt.show()
2158
2159
2160
2161 if False:
2162 lhe = EventFile('unweighted_events.lhe')
2163 import matplotlib.pyplot as plt
2164 import matplotlib.gridspec as gridspec
2165 nbins = 100
2166
2167
2168 mtau, wtau = 1.777, 4.027000e-13
2169 nb_pass = 0
2170 data, data2, data3 = [], [], []
2171 for event in lhe:
2172 nb_pass +=1
2173 if nb_pass > 10000:
2174 break
2175 tau1 = FourMomentum()
2176 tau2 = FourMomentum()
2177 for part in event:
2178 if part.pid in [-12,11,16]:
2179 momenta = FourMomentum(part)
2180 tau1 += momenta
2181 elif part.pid == 15:
2182 tau2 += FourMomentum(part)
2183
2184 if abs((mtau-tau2.mass())/wtau)<1e6 and tau2.mass() >1:
2185 data.append((tau1.mass()-mtau)/wtau)
2186 data2.append((tau2.mass()-mtau)/wtau)
2187 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1])
2188 gs1.update(wspace=0, hspace=0)
2189 ax = plt.subplot(gs1[0])
2190
2191 n, bins, patches = ax.hist(data2, nbins, histtype='step', label='original')
2192 n2, bins2, patches2 = ax.hist(data, bins=bins, histtype='step',label='reconstructed')
2193 import cmath
2194
2195 breit = lambda m : math.sqrt(4*math.pi)*1/(((m)**2-mtau**2)**2+(mtau*wtau)**2)*wtau
2196
2197 data3 = [breit(mtau + x*wtau)*wtau*16867622.6624*50 for x in bins]
2198
2199 ax.plot(bins, data3,label='breit-wigner')
2200
2201 ax.legend()
2202
2203 ax_c = ax.twinx()
2204 ax_c.set_ylabel('MadGraph5_aMC@NLO')
2205 ax_c.yaxis.set_label_coords(1.01, 0.25)
2206 ax_c.set_yticks(ax.get_yticks())
2207 ax_c.set_yticklabels([])
2208
2209 plt.title('invariant mass of tau LHE/reconstructed')
2210 plt.axis('on')
2211 ax.set_xticklabels([])
2212
2213 ax = plt.subplot(gs1[1])
2214 data4 = [n[i]/(data3[i]) for i in range(nbins)]
2215 ax.plot(bins, data4 + [0] , 'b')
2216 data4 = [n2[i]/(data3[i]) for i in range(nbins)]
2217 ax.plot(bins, data4 + [0] , 'g')
2218 ax.set_ylim([0,2])
2219
2220 tick = ax.get_yticks()
2221 ax.set_yticks(tick[:-1])
2222
2223
2224 plt.axis('on')
2225 plt.xlabel('(M - Mtau)/Wtau')
2226 plt.show()
2227