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