1 from __future__ import division
2 import collections
3 import random
4 import re
5 import operator
6 import numbers
7 import math
8 import time
9 import os
10 import shutil
11 import sys
12
13 pjoin = os.path.join
14
15 if '__main__' == __name__:
16 import sys
17 sys.path.append('../../')
18 import misc
19 import logging
20 import gzip
21 import banner as banner_mod
22 logger = logging.getLogger("madgraph.lhe_parser")
23
24 -class Particle(object):
25 """ """
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46 - def __init__(self, line=None, event=None):
47 """ """
48
49 if isinstance(line, Particle):
50 for key in line.__dict__:
51 setattr(self, key, getattr(line, key))
52 if event:
53 self.event = event
54 return
55 else:
56 try:
57 import madgraph.various.hepmc_parser as hepmc_parser
58 except Exception:
59 pass
60 else:
61 if isinstance(line, hepmc_parser.HEPMC_Particle):
62 self.event = event
63 self.event_id = len(event)
64 for key in ['pid', 'status', 'E','px','py','pz','mass']:
65 setattr(self, key, getattr(line, key))
66 self.mother1 = 1
67 self.mother2 = 1
68 self.color1 = 0
69 self.color2 = 0
70 self.vtim = 0
71 self.comment = ''
72 self.helicity = 9
73 self.rwgt = 0
74 return
75
76
77 self.event = event
78 self.event_id = len(event)
79
80 self.pid = 0
81 self.status = 0
82 self.mother1 = None
83 self.mother2 = None
84 self.color1 = 0
85 self.color2 = None
86 self.px = 0
87 self.py = 0
88 self.pz = 0
89 self.E = 0
90 self.mass = 0
91 self.vtim = 0
92 self.helicity = 9
93 self.rwgt = 0
94 self.comment = ''
95
96 if line:
97 self.parse(line)
98
99 @property
101 "convenient alias"
102 return self.pid
103
104 - def parse(self, line):
105 """parse the line"""
106
107 args = line.split()
108 keys = ['pid', 'status','mother1','mother2','color1', 'color2', 'px','py','pz','E',
109 'mass','vtim','helicity']
110
111 for key,value in zip(keys,args):
112 setattr(self, key, float(value))
113 self.pid = int(self.pid)
114
115 self.comment = ' '.join(args[len(keys):])
116 if self.comment.startswith(('|','#')):
117 self.comment = self.comment[1:]
118
119
120
121
123 """string representing the particles"""
124 return " %8d %2d %4d %4d %4d %4d %+13.10e %+13.10e %+13.10e %14.10e %14.10e %10.4e %10.4e" \
125 % (self.pid,
126 self.status,
127 (self.mother1 if isinstance(self.mother1, numbers.Number) else self.mother1.event_id+1) if self.mother1 else 0,
128 (self.mother2 if isinstance(self.mother2, numbers.Number) else self.mother2.event_id+1) if self.mother2 else 0,
129 self.color1,
130 self.color2,
131 self.px,
132 self.py,
133 self.pz,
134 self.E,
135 self.mass,
136 self.vtim,
137 self.helicity)
138
139 - def __eq__(self, other):
140
141 if not isinstance(other, Particle):
142 return False
143 if self.pid == other.pid and \
144 self.status == other.status and \
145 self.mother1 == other.mother1 and \
146 self.mother2 == other.mother2 and \
147 self.color1 == other.color1 and \
148 self.color2 == other.color2 and \
149 self.px == other.px and \
150 self.py == other.py and \
151 self.pz == other.pz and \
152 self.E == other.E and \
153 self.mass == other.mass and \
154 self.vtim == other.vtim and \
155 self.helicity == other.helicity:
156 return True
157 return False
158
159 - def set_momentum(self, momentum):
160
161 self.E = momentum.E
162 self.px = momentum.px
163 self.py = momentum.py
164 self.pz = momentum.pz
165
166 - def add_decay(self, decay_event):
167 """associate to this particle the decay in the associate event"""
168
169 return self.event.add_decay_to_particle(self.event_id, decay_event)
170
171
172 - def __repr__(self):
173 return 'Particle("%s", event=%s)' % (str(self), self.event)
174
177 """A class to allow to read both gzip and not gzip file"""
178
179 allow_empty_event = False
180
181 - def __new__(self, path, mode='r', *args, **opt):
182
183 if not path.endswith(".gz"):
184 return file.__new__(EventFileNoGzip, path, mode, *args, **opt)
185 elif mode == 'r' and not os.path.exists(path) and os.path.exists(path[:-3]):
186 return EventFile.__new__(EventFileNoGzip, path[:-3], mode, *args, **opt)
187 else:
188 try:
189 return gzip.GzipFile.__new__(EventFileGzip, path, mode, *args, **opt)
190 except IOError, error:
191 raise
192 except Exception, error:
193 if mode == 'r':
194 misc.gunzip(path)
195 return file.__new__(EventFileNoGzip, path[:-3], mode, *args, **opt)
196
197
198 - def __init__(self, path, mode='r', *args, **opt):
199 """open file and read the banner [if in read mode]"""
200
201 self.to_zip = False
202 if path.endswith('.gz') and mode == 'w' and\
203 isinstance(self, EventFileNoGzip):
204 path = path[:-3]
205 self.to_zip = True
206
207 self.parsing = True
208 self.eventgroup = False
209 try:
210 super(EventFile, self).__init__(path, mode, *args, **opt)
211 except IOError:
212 if '.gz' in path and isinstance(self, EventFileNoGzip) and\
213 mode == 'r' and os.path.exists(path[:-3]):
214 super(EventFile, self).__init__(path[:-3], mode, *args, **opt)
215 else:
216 raise
217
218 self.banner = ''
219 if mode == 'r':
220 line = ''
221 while '</init>' not in line.lower():
222 try:
223 line = super(EventFile, self).next()
224 except StopIteration:
225 self.seek(0)
226 self.banner = ''
227 break
228 if "<event" in line.lower():
229 self.seek(0)
230 self.banner = ''
231 break
232
233 self.banner += line
234
244
245 @property
247 """return the cross-section of the file #from the banner"""
248 try:
249 return self._cross
250 except Exception:
251 pass
252
253 onebanner = self.get_banner()
254 self._cross = onebanner.get_cross()
255 return self._cross
256
258 if self.closed:
259 return 0
260 if hasattr(self,"len"):
261 return self.len
262 self.seek(0)
263 nb_event=0
264 with misc.TMP_variable(self, 'parsing', False):
265 for _ in self:
266 nb_event +=1
267 self.len = nb_event
268 self.seek(0)
269 return self.len
270
272 """get next event"""
273
274 if not self.eventgroup:
275 text = []
276 line = ''
277 mode = 0
278 while '</event>' not in line:
279 line = super(EventFile, self).next()
280 if '<event' in line:
281 mode = 1
282 text = []
283 if mode:
284 text.append(line)
285 if self.parsing:
286 out = Event(text)
287 if len(out) == 0 and not self.allow_empty_event:
288 raise Exception
289 return out
290 else:
291 return text
292 else:
293 events = []
294 text = []
295 line = ''
296 mode = 0
297 while '</eventgroup>' not in line:
298 line = super(EventFile, self).next()
299 if '<eventgroup' in line:
300 events=[]
301 text = ''
302 elif '<event' in line:
303 text = []
304 mode = 1
305 elif '</event>' in line:
306 if self.parsing:
307 events.append(Event(text))
308 else:
309 events.append(text)
310 text = []
311 mode = 0
312 if mode:
313 text.append(line)
314 if len(events) == 0:
315 return self.next()
316 return events
317
318
320 """ scan once the file to return
321 - the list of the hightest weight (of size trunc_error*NB_EVENT
322 - the cross-section by type of process
323 - the total number of events in the file
324 """
325
326
327
328 self.seek(0)
329 all_wgt = []
330 cross = collections.defaultdict(int)
331 nb_event = 0
332 for event in self:
333 nb_event +=1
334 wgt = get_wgt(event)
335 cross['all'] += wgt
336 cross['abs'] += abs(wgt)
337 cross[event.ievent] += wgt
338 all_wgt.append(abs(wgt))
339
340 if nb_event % 20000 == 0:
341 all_wgt.sort()
342
343 nb_keep = max(20, int(nb_event*trunc_error*15))
344 all_wgt = all_wgt[-nb_keep:]
345
346
347 all_wgt.sort()
348
349 nb_keep = max(20, int(nb_event*trunc_error*10))
350 all_wgt = all_wgt[-nb_keep:]
351 self.seek(0)
352 return all_wgt, cross, nb_event
353
355 """ write a single events or a list of event
356 if self.eventgroup is ON, then add <eventgroup> around the lists of events
357 """
358 if isinstance(event, Event):
359 if self.eventgroup:
360 self.write('<eventgroup>\n%s\n</eventgroup>\n' % event)
361 else:
362 self.write(str(event))
363 elif isinstance(event, list):
364 if self.eventgroup:
365 self.write('<eventgroup>\n')
366 for evt in event:
367 self.write(str(evt))
368 if self.eventgroup:
369 self.write('</eventgroup>\n')
370
371 - def unweight(self, outputpath, get_wgt=None, max_wgt=0, trunc_error=0,
372 event_target=0, log_level=logging.INFO, normalization='average'):
373 """unweight the current file according to wgt information wgt.
374 which can either be a fct of the event or a tag in the rwgt list.
375 max_wgt allow to do partial unweighting.
376 trunc_error allow for dynamical partial unweighting
377 event_target reweight for that many event with maximal trunc_error.
378 (stop to write event when target is reached)
379 """
380 if not get_wgt:
381 def weight(event):
382 return event.wgt
383 get_wgt = weight
384 unwgt_name = "central weight"
385 elif isinstance(get_wgt, str):
386 unwgt_name =get_wgt
387 def get_wgt(event):
388 event.parse_reweight()
389 return event.reweight_data[unwgt_name]
390 else:
391 unwgt_name = get_wgt.func_name
392
393
394 if hasattr(self, "written_weight"):
395 written_weight = lambda x: math.copysign(self.written_weight,float(x))
396 else:
397 written_weight = lambda x: x
398
399 all_wgt, cross, nb_event = self.initialize_unweighting(get_wgt, trunc_error)
400
401
402 def max_wgt_for_trunc(trunc):
403 """find the weight with the maximal truncation."""
404
405 xsum = 0
406 i=1
407 while (xsum - all_wgt[-i] * (i-1) <= cross['abs'] * trunc):
408 max_wgt = all_wgt[-i]
409 xsum += all_wgt[-i]
410 i +=1
411 if i == len(all_wgt):
412 break
413
414 return max_wgt
415
416
417
418 if not max_wgt:
419 if trunc_error == 0 or len(all_wgt)<2 or event_target:
420 max_wgt = all_wgt[-1]
421 else:
422 max_wgt = max_wgt_for_trunc(trunc_error)
423
424
425 if self.banner:
426 try:
427 import internal
428 except:
429 import madgraph.various.banner as banner_module
430 else:
431 import internal.banner as banner_module
432 if not isinstance(self.banner, banner_module.Banner):
433 banner = self.get_banner()
434
435 banner.modify_init_cross(cross)
436
437 banner["unweight"] = "unweighted by %s" % unwgt_name
438 else:
439 banner = self.banner
440
441 curr_strategy = banner.get_lha_strategy()
442 if normalization in ['unit', 'sum']:
443 strategy = 3
444 else:
445 strategy = 4
446 if curr_strategy >0:
447 banner.set_lha_strategy(abs(strategy))
448 else:
449 banner.set_lha_strategy(-1*abs(strategy))
450
451
452 nb_try = 20
453 nb_keep = 0
454 for i in range(nb_try):
455 self.seek(0)
456 if event_target:
457 if i==0:
458 max_wgt = max_wgt_for_trunc(0)
459 else:
460
461 efficiency = nb_keep/nb_event
462 needed_efficiency = event_target/nb_event
463 last_max_wgt = max_wgt
464 needed_max_wgt = last_max_wgt * efficiency / needed_efficiency
465
466 min_max_wgt = max_wgt_for_trunc(trunc_error)
467 max_wgt = max(min_max_wgt, needed_max_wgt)
468 max_wgt = min(max_wgt, all_wgt[-1])
469 if max_wgt == last_max_wgt:
470 if nb_keep < event_target and log_level>=10:
471 logger.log(log_level+10,"fail to reach target %s", event_target)
472 break
473 else:
474 break
475
476
477 if outputpath:
478 outfile = EventFile(outputpath, "w")
479
480
481 if self.banner and outputpath:
482 banner.write(outfile, close_tag=False)
483
484
485 nb_keep = 0
486 trunc_cross = 0
487 for event in self:
488 r = random.random()
489 wgt = get_wgt(event)
490 if abs(wgt) < r * max_wgt:
491 continue
492 elif wgt > 0:
493 nb_keep += 1
494 event.wgt = written_weight(max(wgt, max_wgt))
495 if abs(wgt) > max_wgt:
496 trunc_cross += abs(wgt) - max_wgt
497 if event_target ==0 or nb_keep <= event_target:
498 if outputpath:
499 outfile.write(str(event))
500
501 elif wgt < 0:
502 nb_keep += 1
503 event.wgt = -1* written_weight(max(abs(wgt), max_wgt))
504 if abs(wgt) > max_wgt:
505 trunc_cross += abs(wgt) - max_wgt
506 if outputpath and (event_target ==0 or nb_keep <= event_target):
507 outfile.write(str(event))
508
509 if event_target and nb_keep > event_target:
510 if not outputpath:
511
512 continue
513 elif event_target and i != nb_try-1 and nb_keep >= event_target *1.05:
514 outfile.write("</LesHouchesEvents>\n")
515 outfile.close()
516
517 continue
518 else:
519 outfile.write("</LesHouchesEvents>\n")
520 outfile.close()
521 break
522 elif event_target == 0:
523 if outputpath:
524 outfile.write("</LesHouchesEvents>\n")
525 outfile.close()
526 break
527 elif outputpath:
528 outfile.write("</LesHouchesEvents>\n")
529 outfile.close()
530
531
532 else:
533
534 logger.log(log_level+10,"fail to reach target event %s (iteration=%s)", event_target,i)
535
536
537
538
539 if event_target:
540 nb_events_unweighted = nb_keep
541 nb_keep = min( event_target, nb_keep)
542 else:
543 nb_events_unweighted = nb_keep
544
545 logger.log(log_level, "write %i event (efficiency %.2g %%, truncation %.2g %%) after %i iteration(s)",
546 nb_keep, nb_events_unweighted/nb_event*100, trunc_cross/cross['abs']*100, i)
547
548
549 if nb_keep != event_target and hasattr(self, "written_weight") and strategy !=4:
550 written_weight = lambda x: math.copysign(self.written_weight*event_target/nb_keep, float(x))
551 startfile = EventFile(outputpath)
552 tmpname = pjoin(os.path.dirname(outputpath), "wgtcorrected_"+ os.path.basename(outputpath))
553 outfile = EventFile(tmpname, "w")
554 outfile.write(startfile.banner)
555 for event in startfile:
556 event.wgt = written_weight(event.wgt)
557 outfile.write(str(event))
558 outfile.write("</LesHouchesEvents>\n")
559 startfile.close()
560 outfile.close()
561 shutil.move(tmpname, outputpath)
562
563
564
565
566 self.max_wgt = max_wgt
567 return nb_keep
568
570 """ apply one or more fct on all event. """
571
572 opt= {"print_step": 5000, "maxevent":float("inf"),'no_output':False}
573 opt.update(opts)
574 start = time.time()
575 nb_fct = len(fcts)
576 out = []
577 for i in range(nb_fct):
578 out.append([])
579 self.seek(0)
580 nb_event = 0
581 for event in self:
582 nb_event += 1
583 if opt["print_step"] and (nb_event % opt["print_step"]) == 0:
584 if hasattr(self,"len"):
585 print("currently at %s/%s event [%is]" % (nb_event, self.len, time.time()-start))
586 else:
587 print("currently at %s event [%is]" % (nb_event, time.time()-start))
588 for i in range(nb_fct):
589 value = fcts[i](event)
590 if not opt['no_output']:
591 out[i].append(value)
592 if nb_event > opt['maxevent']:
593 break
594 if nb_fct == 1:
595 return out[0]
596 else:
597 return out
598
599 - def split(self, nb_event=0, partition=None, cwd=os.path.curdir, zip=False):
600 """split the file in multiple file. Do not change the weight!"""
601
602 nb_file = -1
603 for i, event in enumerate(self):
604 if (not (partition is None) and i==sum(partition[:nb_file+1])) or \
605 (partition is None and i % nb_event == 0):
606 if i:
607
608 current.write('</LesHouchesEvent>\n')
609 current.close()
610
611 nb_file +=1
612
613 if not partition is None and (nb_file+1>len(partition)):
614 return nb_file
615 if zip:
616 current = EventFile(pjoin(cwd,'%s_%s.lhe.gz' % (self.name, nb_file)),'w')
617 else:
618 current = open(pjoin(cwd,'%s_%s.lhe' % (self.name, nb_file)),'w')
619 current.write(self.banner)
620 current.write(str(event))
621 if i!=0:
622 current.write('</LesHouchesEvent>\n')
623 current.close()
624
625 return nb_file +1
626
627 - def update_HwU(self, hwu, fct, name='lhe', keep_wgt=False, maxevents=sys.maxint):
628 """take a HwU and add this event file for the function fct"""
629
630 if not isinstance(hwu, list):
631 hwu = [hwu]
632
633 class HwUUpdater(object):
634
635 def __init__(self, fct, keep_wgt):
636
637 self.fct = fct
638 self.first = True
639 self.keep_wgt = keep_wgt
640
641 def add(self, event):
642
643 value = self.fct(event)
644
645 if self.first:
646 for h in hwu:
647
648 if isinstance(value, dict):
649 h.add_line(value.keys())
650 else:
651
652 h.add_line(name)
653 if self.keep_wgt is True:
654 event.parse_reweight()
655 h.add_line(['%s_%s' % (name, key)
656 for key in event.reweight_data])
657 elif self.keep_wgt:
658 h.add_line(self.keep_wgt.values())
659 self.first = False
660
661 for h in hwu:
662 if isinstance(value, tuple):
663 h.addEvent(value[0], value[1])
664 else:
665 h.addEvent(value,{name:event.wgt})
666 if self.keep_wgt:
667 event.parse_reweight()
668 if self.keep_wgt is True:
669 data = dict(('%s_%s' % (name, key),event.reweight_data[key])
670 for key in event.reweight_data)
671 h.addEvent(value, data)
672 else:
673 data = dict(( value,event.reweight_data[key])
674 for key,value in self.keep_wgt.items())
675 h.addEvent(value, data)
676
677
678
679 self.apply_fct_on_event(HwUUpdater(fct,keep_wgt).add, no_output=True,maxevent=maxevents)
680 return hwu
681
683 """take the lhe file and add the matchscale from the pythia_input file"""
684
685 if pythia_input:
686 def next_data():
687 for line in open(pythia_input):
688 if line.startswith('#'):
689 continue
690 data = line.split()
691 print (int(data[0]), data[-3], data[-2], data[-1])
692 yield (int(data[0]), data[-3], data[-2], data[-1])
693 else:
694 def next_data():
695 i=0
696 while 1:
697 yield [i,0,0,0]
698 i+=1
699 sys_iterator = next_data()
700
701 self.seek(0)
702 out = open(out_path,'w')
703
704 pdf_pattern = re.compile(r'''<init>(.*)</init>''', re.M+re.S)
705 init = pdf_pattern.findall(self.banner)[0].split('\n',2)[1]
706 id1, id2, _, _, _, _, pdf1,pdf2,_,_ = init.split()
707 id = [int(id1), int(id2)]
708 type = []
709 for i in range(2):
710 if abs(id[i]) == 2212:
711 if i > 0:
712 type.append(1)
713 else:
714 type.append(-1)
715 else:
716 type.append(0)
717 pdf = max(int(pdf1),int(pdf2))
718
719 out.write("<header>\n" + \
720 "<orgpdf>%i</orgpdf>\n" % pdf + \
721 "<beams> %s %s</beams>\n" % tuple(type) + \
722 "</header>\n")
723
724
725 nevt, smin, smax, scomp = sys_iterator.next()
726 for i, orig_event in enumerate(self):
727 if i < nevt:
728 continue
729 new_event = Event()
730 sys = orig_event.parse_syscalc_info()
731 new_event.syscalc_data = sys
732 if smin:
733 new_event.syscalc_data['matchscale'] = "%s %s %s" % (smin, scomp, smax)
734 out.write(str(new_event), nevt)
735 try:
736 nevt, smin, smax, scomp = sys_iterator.next()
737 except StopIteration:
738 break
739
740 - def get_alphas(self, scale, lhapdf_config='lhapdf-config'):
741 """return the alphas value associated to a given scale"""
742
743 if hasattr(self, 'alpsrunner'):
744 return self.alpsrunner(scale)
745
746
747 banner = banner_mod.Banner(self.banner)
748 run_card = banner.charge_card('run_card')
749 use_runner = False
750 if abs(run_card['lpp1']) != 1 and abs(run_card['lpp2']) != 1:
751
752 use_runner = True
753 else:
754
755 lhapdf = misc.import_python_lhapdf(lhapdf_config)
756 if not lhapdf:
757 logger.warning('fail to link to lhapdf for the alphas-running. Use Two loop computation')
758 use_runner = True
759 try:
760 self.pdf = lhapdf.mkPDF(int(self.banner.run_card.get_lhapdf_id()))
761 except Exception:
762 logger.warning('fail to link to lhapdf for the alphas-running. Use Two loop computation')
763 use_runner = True
764
765 if not use_runner:
766 self.alpsrunner = lambda scale: self.pdf.alphasQ(scale)
767 else:
768 try:
769 from models.model_reader import Alphas_Runner
770 except ImportError:
771 root = os.path.dirname(__file__)
772 root_path = pjoin(root, os.pardir, os.pardir)
773 try:
774 import internal.madevent_interface as me_int
775 cmd = me_int.MadEventCmd(root_path,force_run=True)
776 except ImportError:
777 import internal.amcnlo_run_interface as me_int
778 cmd = me_int.Cmd(root_path,force_run=True)
779 if 'mg5_path' in cmd.options and cmd.options['mg5_path']:
780 sys.path.append(cmd.options['mg5_path'])
781 from models.model_reader import Alphas_Runner
782
783 if not hasattr(banner, 'param_card'):
784 param_card = banner.charge_card('param_card')
785 else:
786 param_card = banner.param_card
787
788 asmz = param_card.get_value('sminputs', 3, 0.13)
789 nloop =2
790 zmass = param_card.get_value('mass', 23, 91.188)
791 cmass = param_card.get_value('mass', 4, 1.4)
792 if cmass == 0:
793 cmass = 1.4
794 bmass = param_card.get_value('mass', 5, 4.7)
795 if bmass == 0:
796 bmass = 4.7
797 self.alpsrunner = Alphas_Runner(asmz, nloop, zmass, cmass, bmass)
798
799
800
801 return self.alpsrunner(scale)
802
803
804
805
806
807
808
809
810 -class EventFileGzip(EventFile, gzip.GzipFile):
811 """A way to read/write a gzipped lhef event"""
812
813
815 currpos = super(EventFileGzip, self).tell()
816 if not currpos:
817 currpos = self.size
818 return currpos
819
821 """A way to read a standard event file"""
822
823 - def close(self,*args, **opts):
828
830 """a class to read simultaneously multiple file and read them in mixing them.
831 Unweighting can be done at the same time.
832 The number of events in each file need to be provide in advance
833 (if not provide the file is first read to find that number"""
834
835 - def __new__(cls, start_list=[],parse=True):
837
838 - def __init__(self, start_list=[], parse=True):
839 """if trunc_error is define here then this allow
840 to only read all the files twice and not three times."""
841 self.eventgroup = False
842 self.files = []
843 self.parsefile = parse
844 self.banner = ''
845 self.initial_nb_events = []
846 self.total_event_in_files = 0
847 self.curr_nb_events = []
848 self.allcross = []
849 self.error = []
850 self.across = []
851 self.scales = []
852 if start_list:
853 if parse:
854 for p in start_list:
855 self.add(p)
856 else:
857 self.files = start_list
858 self._configure = False
859
860 - def close(self,*args,**opts):
863
864 - def add(self, path, cross, error, across, nb_event=0, scale=1):
865 """ add a file to the pool, across allow to reweight the sum of weight
866 in the file to the given cross-section
867 """
868
869 if across == 0:
870
871 return
872
873 obj = EventFile(path)
874 obj.eventgroup = self.eventgroup
875 if len(self.files) == 0 and not self.banner:
876 self.banner = obj.banner
877 self.curr_nb_events.append(0)
878 self.initial_nb_events.append(0)
879 self.allcross.append(cross)
880 self.across.append(across)
881 self.error.append(error)
882 self.scales.append(scale)
883 self.files.append(obj)
884 if nb_event:
885 obj.len = nb_event
886 self._configure = False
887 return obj
888
891
893
894 if not self._configure:
895 self.configure()
896
897 remaining_event = self.total_event_in_files - sum(self.curr_nb_events)
898 if remaining_event == 0:
899 raise StopIteration
900
901 nb_event = random.randint(1, remaining_event)
902 sum_nb=0
903 for i, obj in enumerate(self.files):
904 sum_nb += self.initial_nb_events[i] - self.curr_nb_events[i]
905 if nb_event <= sum_nb:
906 self.curr_nb_events[i] += 1
907 event = obj.next()
908 if not self.eventgroup:
909 event.sample_scale = self.scales[i]
910 else:
911 for evt in event:
912 evt.sample_scale = self.scales[i]
913 return event
914 else:
915 raise Exception
916
917
919 """define the part of the init_banner"""
920
921 if not self.banner:
922 return
923
924
925 grouped_cross = {}
926 grouped_error = {}
927 for i,ff in enumerate(self.files):
928 filename = ff.name
929 from_init = False
930 Pdir = [P for P in filename.split(os.path.sep) if P.startswith('P')]
931 if Pdir:
932 Pdir = Pdir[-1]
933 group = Pdir.split("_")[0][1:]
934 if not group.isdigit():
935 from_init = True
936 else:
937 from_init = True
938
939 if not from_init:
940 if group in grouped_cross:
941 grouped_cross[group] += self.allcross[i]
942 grouped_error[group] += self.error[i]**2
943 else:
944 grouped_cross[group] = self.allcross[i]
945 grouped_error[group] = self.error[i]**2
946 else:
947 ban = banner_mod.Banner(ff.banner)
948 for line in ban['init'].split('\n'):
949 splitline = line.split()
950 if len(splitline)==4:
951 cross, error, _, group = splitline
952 if int(group) in grouped_cross:
953 grouped_cross[group] += float(cross)
954 grouped_error[group] += float(error)**2
955 else:
956 grouped_cross[group] = float(cross)
957 grouped_error[group] = float(error)**2
958 nb_group = len(grouped_cross)
959
960
961 try:
962 run_card = self.banner.run_card
963 except:
964 run_card = self.banner.charge_card("run_card")
965
966 init_information = run_card.get_banner_init_information()
967
968 if proc_charac and proc_charac['ninitial'] == 1:
969
970 init_information = run_card.get_banner_init_information()
971 event = self.next()
972 init_information["idbmup1"] = event[0].pdg
973 init_information["ebmup1"] = event[0].mass
974 init_information["idbmup2"] = 0
975 init_information["ebmup2"] = 0
976 self.seek(0)
977 else:
978
979 if init_information["idbmup1"] in [0,9]:
980 event = self.next()
981 init_information["idbmup1"]= event[0].pdg
982 if init_information["idbmup2"] == 0:
983 init_information["idbmup2"]= event[1].pdg
984 self.seek(0)
985 if init_information["idbmup2"] in [0,9]:
986 event = self.next()
987 init_information["idbmup2"] = event[1].pdg
988 self.seek(0)
989
990 init_information["nprup"] = nb_group
991
992 if run_card["lhe_version"] < 3:
993 init_information["generator_info"] = ""
994 else:
995 init_information["generator_info"] = "<generator name='MadGraph5_aMC@NLO' version='%s'>please cite 1405.0301 </generator>\n" \
996 % misc.get_pkg_info()['version']
997
998
999 cross_info = "%(cross)e %(error)e %(wgt)e %(id)i"
1000 init_information["cross_info"] = []
1001 for id in grouped_cross:
1002 conv = {"id": int(id), "cross": grouped_cross[id], "error": math.sqrt(grouped_error[id]),
1003 "wgt": wgt}
1004 init_information["cross_info"].append( cross_info % conv)
1005 init_information["cross_info"] = '\n'.join(init_information["cross_info"])
1006 init_information['lha_stra'] = -1 * abs(lha_strategy)
1007
1008 template_init =\
1009 """ %(idbmup1)i %(idbmup2)i %(ebmup1)e %(ebmup2)e %(pdfgup1)i %(pdfgup2)i %(pdfsup1)i %(pdfsup2)i %(lha_stra)i %(nprup)i
1010 %(cross_info)s
1011 %(generator_info)s
1012 """
1013
1014 self.banner["init"] = template_init % init_information
1015
1016
1017
1019 """ scan once the file to return
1020 - the list of the hightest weight (of size trunc_error*NB_EVENT
1021 - the cross-section by type of process
1022 - the total number of events in the files
1023 In top of that it initialise the information for the next routine
1024 to determine how to choose which file to read
1025 """
1026 self.seek(0)
1027 all_wgt = []
1028 total_event = 0
1029 sum_cross = collections.defaultdict(int)
1030 for i,f in enumerate(self.files):
1031 nb_event = 0
1032
1033
1034 cross = collections.defaultdict(int)
1035 new_wgt =[]
1036 for event in f:
1037 nb_event += 1
1038 total_event += 1
1039 event.sample_scale = 1
1040 wgt = getwgt(event)
1041 cross['all'] += wgt
1042 cross['abs'] += abs(wgt)
1043 cross[event.ievent] += wgt
1044 new_wgt.append(abs(wgt))
1045
1046 if nb_event % 20000 == 0:
1047 new_wgt.sort()
1048
1049 nb_keep = max(20, int(nb_event*trunc_error*15))
1050 new_wgt = new_wgt[-nb_keep:]
1051 if nb_event == 0:
1052 raise Exception
1053
1054 self.initial_nb_events[i] = nb_event
1055 self.scales[i] = self.across[i]/cross['abs'] if self.across[i] else 1
1056
1057
1058 for key in cross:
1059 sum_cross[key] += cross[key]* self.scales[i]
1060 all_wgt +=[self.scales[i] * w for w in new_wgt]
1061 all_wgt.sort()
1062 nb_keep = max(20, int(total_event*trunc_error*10))
1063 all_wgt = all_wgt[-nb_keep:]
1064
1065 self.total_event_in_files = total_event
1066
1067 all_wgt.sort()
1068
1069 nb_keep = max(20, int(total_event*trunc_error*10))
1070 all_wgt = all_wgt[-nb_keep:]
1071 self.seek(0)
1072 self._configure = True
1073 return all_wgt, sum_cross, total_event
1074
1081
1083
1084 return len(self.files)
1085
1086 - def seek(self, pos):
1087 """ """
1088
1089 if pos !=0:
1090 raise Exception
1091 for i in range(len(self)):
1092 self.curr_nb_events[i] = 0
1093 for f in self.files:
1094 f.seek(pos)
1095
1096 - def unweight(self, outputpath, get_wgt, **opts):
1097 """unweight the current file according to wgt information wgt.
1098 which can either be a fct of the event or a tag in the rwgt list.
1099 max_wgt allow to do partial unweighting.
1100 trunc_error allow for dynamical partial unweighting
1101 event_target reweight for that many event with maximal trunc_error.
1102 (stop to write event when target is reached)
1103 """
1104
1105 if isinstance(get_wgt, str):
1106 unwgt_name =get_wgt
1107 def get_wgt_multi(event):
1108 event.parse_reweight()
1109 return event.reweight_data[unwgt_name] * event.sample_scale
1110 else:
1111 unwgt_name = get_wgt.func_name
1112 get_wgt_multi = lambda event: get_wgt(event) * event.sample_scale
1113
1114
1115 if 'proc_charac' in opts:
1116 if opts['proc_charac']:
1117 proc_charac = opts['proc_charac']
1118 else:
1119 proc_charac=None
1120 del opts['proc_charac']
1121 else:
1122 proc_charac = None
1123
1124 if 'event_target' in opts and opts['event_target']:
1125 if 'normalization' in opts:
1126 if opts['normalization'] == 'sum':
1127 new_wgt = sum(self.across)/opts['event_target']
1128 strategy = 3
1129 elif opts['normalization'] == 'average':
1130 strategy = 4
1131 new_wgt = sum(self.across)
1132 elif opts['normalization'] == 'unit':
1133 strategy =3
1134 new_wgt = 1.
1135 else:
1136 strategy = 4
1137 new_wgt = sum(self.across)
1138 self.define_init_banner(new_wgt, strategy, proc_charac=proc_charac)
1139 self.written_weight = new_wgt
1140 elif 'write_init' in opts and opts['write_init']:
1141 self.define_init_banner(0,0, proc_charac=proc_charac)
1142 del opts['write_init']
1143 return super(MultiEventFile, self).unweight(outputpath, get_wgt_multi, **opts)
1144
1145 - def write(self, path, random=False, banner=None, get_info=False):
1194
1203
1204
1205
1206 -class Event(list):
1207 """Class storing a single event information (list of particles + global information)"""
1208
1209 warning_order = True
1210
1212 """The initialization of an empty Event (or one associate to a text file)"""
1213 list.__init__(self)
1214
1215
1216 self.nexternal = 0
1217 self.ievent = 0
1218 self.wgt = 0
1219 self.aqcd = 0
1220 self.scale = 0
1221 self.aqed = 0
1222 self.aqcd = 0
1223
1224 self.tag = ''
1225 self.eventflag = {}
1226 self.comment = ''
1227 self.reweight_data = {}
1228 self.matched_scale_data = None
1229 self.syscalc_data = {}
1230 if text:
1231 self.parse(text)
1232
1233
1234
1236 """Take the input file and create the structured information"""
1237
1238 status = 'first'
1239 try:
1240 text = text.split('\n')
1241 except Exception:
1242 pass
1243 for line in text:
1244 line = line.strip()
1245 if not line:
1246 continue
1247 elif line[0] == '#':
1248 self.comment += '%s\n' % line
1249 continue
1250 elif line.startswith('<event'):
1251 if '=' in line:
1252 found = re.findall(r"""(\w*)=(?:(?:['"])([^'"]*)(?=['"])|(\S*))""",line)
1253
1254
1255 self.eventflag = dict((n, a1) if a1 else (n,a2) for n,a1,a2 in found)
1256
1257 continue
1258
1259 elif 'first' == status:
1260 if '<rwgt>' in line:
1261 status = 'tag'
1262 else:
1263 self.assign_scale_line(line)
1264 status = 'part'
1265 continue
1266 if '<' in line:
1267 status = 'tag'
1268
1269 if 'part' == status:
1270 part = Particle(line, event=self)
1271 if part.E != 0 or part.status==-1:
1272 self.append(part)
1273 elif self.nexternal:
1274 self.nexternal-=1
1275 else:
1276 if '</event>' in line:
1277 line = line.replace('</event>','',1)
1278 self.tag += '%s\n' % line
1279
1280 self.assign_mother()
1281
1282
1284 """convert the number in actual particle"""
1285
1286 if all(p.status != -1 for p in self):
1287 if not self.nexternal:
1288 return
1289 if self.warning_order:
1290 Event.warning_order = False
1291 logger.warning("Weird format for lhe format: no incoming particle... adding a fake one")
1292 raise Exception
1293 mother = Particle(event=self)
1294 mother.status = -1
1295 mother.pid = 0
1296 self.insert(0,mother)
1297 mother.color2 = 0
1298 mother.event_id = 0
1299 self.nexternal += 1
1300 for p in self[1:]:
1301 if isinstance(p.mother1, int) and p.mother1 > 1:
1302 p.mother1 += 1
1303 if isinstance(p.mother2, int) and p.mother2 > 1:
1304 p.mother2 += 1
1305 p.event_id += 1
1306
1307
1308
1309 for i,particle in enumerate(self):
1310 if i < particle.mother1 or i < particle.mother2:
1311 if self.warning_order:
1312 logger.warning("Order of particle in the event did not agree with parent/child order. This might be problematic for some code.")
1313 Event.warning_order = False
1314 self.reorder_mother_child()
1315 return self.assign_mother()
1316
1317 if particle.mother1:
1318 try:
1319 particle.mother1 = self[int(particle.mother1) -1]
1320 except Exception:
1321 logger.warning("WRONG MOTHER INFO %s", self)
1322 particle.mother1 = 0
1323 if particle.mother2:
1324 try:
1325 particle.mother2 = self[int(particle.mother2) -1]
1326 except Exception:
1327 logger.warning("WRONG MOTHER INFO %s", self)
1328 particle.mother2 = 0
1329
1331 """change all the weights by a given ratio"""
1332
1333 self.wgt *= ratio
1334 self.parse_reweight()
1335 for key in self.reweight_data:
1336 self.reweight_data[key] *= ratio
1337 return self.wgt
1338
1340 """check and correct the mother/child position.
1341 only correct one order by call (but this is a recursive call)"""
1342
1343 tomove, position = None, None
1344 for i,particle in enumerate(self):
1345 if i < particle.mother1:
1346
1347 tomove, position = i, particle.mother1-1
1348 break
1349 if i < particle.mother2:
1350 tomove, position = i, particle.mother2-1
1351
1352
1353 if not tomove:
1354 return
1355
1356
1357 particle = self.pop(tomove)
1358 self.insert(int(position), particle)
1359
1360
1361 for i, particle in enumerate(self):
1362 particle.event_id = i
1363
1364 m1, m2 = particle.mother1, particle.mother2
1365 if m1 == tomove +1:
1366 particle.mother1 = position+1
1367 elif tomove < m1 <= position +1:
1368 particle.mother1 -= 1
1369 if m2 == tomove +1:
1370 particle.mother2 = position+1
1371 elif tomove < m2 <= position +1:
1372 particle.mother2 -= 1
1373
1374 return self.reorder_mother_child()
1375
1376
1377
1378
1379
1380
1382 """Parse the re-weight information in order to return a dictionary
1383 {key: value}. If no group is define group should be '' """
1384 if self.reweight_data:
1385 return self.reweight_data
1386 self.reweight_data = {}
1387 self.reweight_order = []
1388 start, stop = self.tag.find('<rwgt>'), self.tag.find('</rwgt>')
1389 if start != -1 != stop :
1390 pattern = re.compile(r'''<\s*wgt id=(?:\'|\")(?P<id>[^\'\"]+)(?:\'|\")\s*>\s*(?P<val>[\ded+-.]*)\s*</wgt>''',re.I)
1391 data = pattern.findall(self.tag[start:stop])
1392 try:
1393 self.reweight_data = dict([(pid, float(value)) for (pid, value) in data
1394 if not self.reweight_order.append(pid)])
1395
1396 except ValueError, error:
1397 raise Exception, 'Event File has unvalid weight. %s' % error
1398 self.tag = self.tag[:start] + self.tag[stop+7:]
1399 return self.reweight_data
1400
1402 """ """
1403 if hasattr(self, 'nloweight'):
1404 return self.nloweight
1405
1406 start, stop = self.tag.find('<mgrwgt>'), self.tag.find('</mgrwgt>')
1407 if start != -1 != stop :
1408
1409 text = self.tag[start+8:stop]
1410 self.nloweight = NLO_PARTIALWEIGHT(text, self, real_type=real_type,
1411 threshold=threshold)
1412 return self.nloweight
1413
1415 """get the string associate to the weight"""
1416
1417 text="""<mgrwgt>
1418 %(total_wgt).10e %(nb_wgt)i %(nb_event)i 0
1419 %(event)s
1420 %(wgt)s
1421 </mgrwgt>"""
1422
1423
1424 if not wgt:
1425 if not hasattr(self, 'nloweight'):
1426 return
1427 wgt = self.nloweight
1428
1429 data = {'total_wgt': wgt.total_wgt,
1430 'nb_wgt': wgt.nb_wgt,
1431 'nb_event': wgt.nb_event,
1432 'event': '\n'.join(p.__str__(mode='fortran') for p in wgt.momenta),
1433 'wgt':'\n'.join(w.__str__(mode='formatted')
1434 for e in wgt.cevents for w in e.wgts)}
1435
1436 data['total_wgt'] = sum([w.ref_wgt for e in wgt.cevents for w in e.wgts])
1437 start, stop = self.tag.find('<mgrwgt>'), self.tag.find('</mgrwgt>')
1438
1439 self.tag = self.tag[:start] + text % data + self.tag[stop+9:]
1440
1441
1443 """ """
1444
1445
1446 if hasattr(self, 'loweight'):
1447 return self.loweight
1448
1449 if not hasattr(Event, 'loweight_pattern'):
1450 Event.loweight_pattern = re.compile('''<rscale>\s*(?P<nqcd>\d+)\s+(?P<ren_scale>[\d.e+-]+)\s*</rscale>\s*\n\s*
1451 <asrwt>\s*(?P<asrwt>[\s\d.+-e]+)\s*</asrwt>\s*\n\s*
1452 <pdfrwt\s+beam=["']?1["']?\>\s*(?P<beam1>[\s\d.e+-]*)\s*</pdfrwt>\s*\n\s*
1453 <pdfrwt\s+beam=["']?2["']?\>\s*(?P<beam2>[\s\d.e+-]*)\s*</pdfrwt>\s*\n\s*
1454 <totfact>\s*(?P<totfact>[\d.e+-]*)\s*</totfact>
1455 ''',re.X+re.I+re.M)
1456
1457 start, stop = self.tag.find('<mgrwt>'), self.tag.find('</mgrwt>')
1458
1459 if start != -1 != stop :
1460 text = self.tag[start+8:stop]
1461
1462 info = Event.loweight_pattern.search(text)
1463 if not info:
1464 raise Exception, '%s not parsed'% text
1465 self.loweight={}
1466 self.loweight['n_qcd'] = int(info.group('nqcd'))
1467 self.loweight['ren_scale'] = float(info.group('ren_scale'))
1468 self.loweight['asrwt'] =[float(x) for x in info.group('asrwt').split()[1:]]
1469 self.loweight['tot_fact'] = float(info.group('totfact'))
1470
1471 args = info.group('beam1').split()
1472 npdf = int(args[0])
1473 self.loweight['n_pdfrw1'] = npdf
1474 self.loweight['pdf_pdg_code1'] = [int(i) for i in args[1:1+npdf]]
1475 self.loweight['pdf_x1'] = [float(i) for i in args[1+npdf:1+2*npdf]]
1476 self.loweight['pdf_q1'] = [float(i) for i in args[1+2*npdf:1+3*npdf]]
1477 args = info.group('beam2').split()
1478 npdf = int(args[0])
1479 self.loweight['n_pdfrw2'] = npdf
1480 self.loweight['pdf_pdg_code2'] = [int(i) for i in args[1:1+npdf]]
1481 self.loweight['pdf_x2'] = [float(i) for i in args[1+npdf:1+2*npdf]]
1482 self.loweight['pdf_q2'] = [float(i) for i in args[1+2*npdf:1+3*npdf]]
1483
1484 else:
1485 return None
1486 return self.loweight
1487
1488
1490 """Parse the line containing the starting scale for the shower"""
1491
1492 if self.matched_scale_data is not None:
1493 return self.matched_scale_data
1494
1495 self.matched_scale_data = []
1496
1497
1498 pattern = re.compile("<scales\s|</scales>")
1499 data = re.split(pattern,self.tag)
1500 if len(data) == 1:
1501 return []
1502 else:
1503 tmp = {}
1504 start,content, end = data
1505 self.tag = "%s%s" % (start, end)
1506 pattern = re.compile("pt_clust_(\d*)=\"([\de+-.]*)\"")
1507 for id,value in pattern.findall(content):
1508 tmp[int(id)] = float(value)
1509 for i in range(1, len(self)+1):
1510 if i in tmp:
1511 self.matched_scale_data.append(tmp[i])
1512 else:
1513 self.matched_scale_data.append(-1)
1514 return self.matched_scale_data
1515
1517 """ parse the flag for syscalc between <mgrwt></mgrwt>
1518 <mgrwt>
1519 <rscale> 3 0.26552898E+03</rscale>
1520 <asrwt>0</asrwt>
1521 <pdfrwt beam="1"> 1 21 0.14527945E+00 0.26552898E+03</pdfrwt>
1522 <pdfrwt beam="2"> 1 21 0.15249110E-01 0.26552898E+03</pdfrwt>
1523 <totfact> 0.10344054E+04</totfact>
1524 </mgrwt>
1525 """
1526 if self.syscalc_data:
1527 return self.syscalc_data
1528
1529 pattern = re.compile("<mgrwt>|</mgrwt>")
1530 pattern2 = re.compile("<(?P<tag>[\w]*)(?:\s*(\w*)=[\"'](.*)[\"']\s*|\s*)>(.*)</(?P=tag)>")
1531 data = re.split(pattern,self.tag)
1532 if len(data) == 1:
1533 return []
1534 else:
1535 tmp = {}
1536 start,content, end = data
1537 self.tag = "%s%s" % (start, end)
1538 for tag, key, keyval, tagval in pattern2.findall(content):
1539 if key:
1540 self.syscalc_data[(tag, key, keyval)] = tagval
1541 else:
1542 self.syscalc_data[tag] = tagval
1543 return self.syscalc_data
1544
1545
1546 - def add_decay_to_particle(self, position, decay_event):
1547 """define the decay of the particle id by the event pass in argument"""
1548
1549 this_particle = self[position]
1550
1551 this_particle.status = 2
1552 this_particle.helicity = 0
1553
1554
1555 decay_particle = decay_event[0]
1556 this_4mom = FourMomentum(this_particle)
1557 nb_part = len(self)
1558
1559 thres = decay_particle.E*1e-10
1560 assert max(decay_particle.px, decay_particle.py, decay_particle.pz) < thres,\
1561 "not on rest particle %s %s %s %s" % (decay_particle.E, decay_particle.px,decay_particle.py,decay_particle.pz)
1562
1563 self.nexternal += decay_event.nexternal -1
1564 old_scales = list(self.parse_matching_scale())
1565 if old_scales:
1566 jet_position = sum(1 for i in range(position) if self[i].status==1)
1567 initial_pos = sum(1 for i in range(position) if self[i].status==-1)
1568 self.matched_scale_data.pop(initial_pos+jet_position)
1569
1570
1571 for particle in decay_event[1:]:
1572
1573 new_particle = Particle(particle, self)
1574 new_particle.event_id = len(self)
1575 self.append(new_particle)
1576 if old_scales:
1577 self.matched_scale_data.append(old_scales[initial_pos+jet_position])
1578
1579 new_momentum = FourMomentum(new_particle).boost(this_4mom)
1580 new_particle.set_momentum(new_momentum)
1581
1582 for tag in ['mother1', 'mother2']:
1583 mother = getattr(particle, tag)
1584 if isinstance(mother, Particle):
1585 mother_id = getattr(particle, tag).event_id
1586 if mother_id == 0:
1587 setattr(new_particle, tag, this_particle)
1588 else:
1589 try:
1590 setattr(new_particle, tag, self[nb_part + mother_id -1])
1591 except Exception, error:
1592 print error
1593 misc.sprint( self)
1594 misc.sprint(nb_part + mother_id -1)
1595 misc.sprint(tag)
1596 misc.sprint(position, decay_event)
1597 misc.sprint(particle)
1598 misc.sprint(len(self), nb_part + mother_id -1)
1599 raise
1600 elif tag == "mother2" and isinstance(particle.mother1, Particle):
1601 new_particle.mother2 = this_particle
1602 else:
1603 raise Exception, "Something weird happens. Please report it for investigation"
1604
1605
1606 max_color=501
1607 for particle in self[:nb_part]:
1608 max_color=max(max_color, particle.color1, particle.color2)
1609
1610
1611 color_mapping = {}
1612 color_mapping[decay_particle.color1] = this_particle.color1
1613 color_mapping[decay_particle.color2] = this_particle.color2
1614 for particle in self[nb_part:]:
1615 if particle.color1:
1616 if particle.color1 not in color_mapping:
1617 max_color +=1
1618 color_mapping[particle.color1] = max_color
1619 particle.color1 = max_color
1620 else:
1621 particle.color1 = color_mapping[particle.color1]
1622 if particle.color2:
1623 if particle.color2 not in color_mapping:
1624 max_color +=1
1625 color_mapping[particle.color2] = max_color
1626 particle.color2 = max_color
1627 else:
1628 particle.color2 = color_mapping[particle.color2]
1629
1631 """use auto-recursion"""
1632
1633 pdg_to_decay = dict(pdg_to_decay)
1634
1635 for i,particle in enumerate(self):
1636 if particle.status != 1:
1637 continue
1638 if particle.pdg in pdg_to_decay and pdg_to_decay[particle.pdg]:
1639 one_decay = pdg_to_decay[particle.pdg].pop()
1640 self.add_decay_to_particle(i, one_decay)
1641 return self.add_decays(pdg_to_decay)
1642 return self
1643
1644
1645
1647
1648 to_remove = []
1649 if event_id is not None:
1650 to_remove.append(self[event_id])
1651
1652 if pdg_code:
1653 for particle in self:
1654 if particle.pid == pdg_code:
1655 to_remove.append(particle)
1656
1657 new_event = Event()
1658
1659 for tag in ['nexternal', 'ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']:
1660 setattr(new_event, tag, getattr(self, tag))
1661
1662 for particle in self:
1663 if isinstance(particle.mother1, Particle) and particle.mother1 in to_remove:
1664 to_remove.append(particle)
1665 if particle.status == 1:
1666 new_event.nexternal -= 1
1667 continue
1668 elif isinstance(particle.mother2, Particle) and particle.mother2 in to_remove:
1669 to_remove.append(particle)
1670 if particle.status == 1:
1671 new_event.nexternal -= 1
1672 continue
1673 else:
1674 new_event.append(Particle(particle))
1675
1676
1677
1678 for pos, particle in enumerate(new_event):
1679 particle.event_id = pos
1680 if particle in to_remove:
1681 particle.status = 1
1682 return new_event
1683
1684 - def get_decay(self, pdg_code=0, event_id=None):
1685
1686 to_start = []
1687 if event_id is not None:
1688 to_start.append(self[event_id])
1689
1690 elif pdg_code:
1691 for particle in self:
1692 if particle.pid == pdg_code:
1693 to_start.append(particle)
1694 break
1695
1696 new_event = Event()
1697
1698 for tag in ['ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']:
1699 setattr(new_event, tag, getattr(self, tag))
1700
1701
1702 old2new = {}
1703 new_decay_part = Particle(to_start[0])
1704 new_decay_part.mother1 = None
1705 new_decay_part.mother2 = None
1706 new_decay_part.status = -1
1707 old2new[new_decay_part.event_id] = len(old2new)
1708 new_event.append(new_decay_part)
1709
1710
1711
1712 for particle in self:
1713 if isinstance(particle.mother1, Particle) and particle.mother1.event_id in old2new\
1714 or isinstance(particle.mother2, Particle) and particle.mother2.event_id in old2new:
1715 old2new[particle.event_id] = len(old2new)
1716 new_event.append(Particle(particle))
1717
1718
1719
1720 nexternal = 0
1721 for pos, particle in enumerate(new_event):
1722 particle.event_id = pos
1723 if particle.mother1:
1724 particle.mother1 = new_event[old2new[particle.mother1.event_id]]
1725 if particle.mother2:
1726 particle.mother2 = new_event[old2new[particle.mother2.event_id]]
1727 if particle.status in [-1,1]:
1728 nexternal +=1
1729 new_event.nexternal = nexternal
1730
1731 return new_event
1732
1733 - def boost(self, filter=None):
1734 """modify the current event to boost it according to the current filter"""
1735 if filter is None:
1736 filter = lambda p: p.status==-1
1737
1738 if not isinstance(filter, FourMomentum):
1739 pboost = FourMomentum()
1740 for p in self:
1741 if filter(p):
1742 pboost += p
1743 else:
1744 pboost = FourMomentum(pboost)
1745
1746
1747 pboost.px *=-1
1748 pboost.py *=-1
1749 pboost.pz *=-1
1750 for p in self:
1751 b= FourMomentum(p).boost(pboost)
1752 p.E, p.px, p.py, p.pz = b.E, b.px, b.py, b.pz
1753
1754 return self
1755
1757 """check various property of the events"""
1758
1759
1760 threshold = 1e-6
1761
1762
1763 E, px, py, pz = 0,0,0,0
1764 absE, abspx, abspy, abspz = 0,0,0,0
1765 for particle in self:
1766 coeff = 1
1767 if particle.status == -1:
1768 coeff = -1
1769 elif particle.status != 1:
1770 continue
1771 E += coeff * particle.E
1772 absE += abs(particle.E)
1773 px += coeff * particle.px
1774 py += coeff * particle.py
1775 pz += coeff * particle.pz
1776 abspx += abs(particle.px)
1777 abspy += abs(particle.py)
1778 abspz += abs(particle.pz)
1779
1780 fourmass = FourMomentum(particle).mass
1781
1782 if particle.mass and (abs(particle.mass) - fourmass)/ abs(particle.mass) > threshold:
1783 raise Exception, "Do not have correct mass lhe: %s momentum: %s (error at %s" % (particle.mass, fourmass, (abs(particle.mass) - fourmass)/ abs(particle.mass))
1784
1785
1786
1787 if E/absE > threshold:
1788 logger.critical(self)
1789 raise Exception, "Do not conserve Energy %s, %s" % (E/absE, E)
1790 if px/abspx > threshold:
1791 logger.critical(self)
1792 raise Exception, "Do not conserve Px %s, %s" % (px/abspx, px)
1793 if py/abspy > threshold:
1794 logger.critical(self)
1795 raise Exception, "Do not conserve Py %s, %s" % (py/abspy, py)
1796 if pz/abspz > threshold:
1797 logger.critical(self)
1798 raise Exception, "Do not conserve Pz %s, %s" % (pz/abspz, pz)
1799
1800
1801 self.check_color_structure()
1802
1803
1804
1805
1807 """read the line corresponding to global event line
1808 format of the line is:
1809 Nexternal IEVENT WEIGHT SCALE AEW AS
1810 """
1811 inputs = line.split()
1812 assert len(inputs) == 6
1813 self.nexternal=int(inputs[0])
1814 self.ievent=int(inputs[1])
1815 self.wgt=float(inputs[2])
1816 self.scale=float(inputs[3])
1817 self.aqed=float(inputs[4])
1818 self.aqcd=float(inputs[5])
1819
1821 """Return the unique tag identifying the SubProcesses for the generation.
1822 Usefull for program like MadSpin and Reweight module."""
1823
1824 initial, final, order = [], [], [[], []]
1825 for particle in self:
1826 if particle.status == -1:
1827 initial.append(particle.pid)
1828 order[0].append(particle.pid)
1829 elif particle.status == 1:
1830 final.append(particle.pid)
1831 order[1].append(particle.pid)
1832 initial.sort(), final.sort()
1833 tag = (tuple(initial), tuple(final))
1834 return tag, order
1835
1836 @staticmethod
1837 - def mass_shuffle(momenta, sqrts, new_mass, new_sqrts=None):
1838 """use the RAMBO method to shuffle the PS. initial sqrts is preserved."""
1839
1840 if not new_sqrts:
1841 new_sqrts = sqrts
1842
1843 oldm = [p.mass_sqr for p in momenta]
1844 newm = [m**2 for m in new_mass]
1845 tot_mom = sum(momenta, FourMomentum())
1846 if tot_mom.pt2 > 1e-5:
1847 boost_back = FourMomentum(tot_mom.mass,0,0,0).boost_to_restframe(tot_mom)
1848 for i,m in enumerate(momenta):
1849 momenta[i] = m.boost_to_restframe(tot_mom)
1850
1851
1852 f = lambda chi: new_sqrts - sum(math.sqrt(max(0, M + chi**2*(p.E**2-m)))
1853 for M,p,m in zip(newm, momenta,oldm))
1854
1855 df = lambda chi: -1* sum(chi*(p.E**2-m)/math.sqrt(max(0,(p.E**2-m)*chi**2+M))
1856 for M,p,m in zip(newm, momenta,oldm))
1857
1858 if sum(new_mass) > new_sqrts:
1859 return momenta, 0
1860 try:
1861 chi = misc.newtonmethod(f, df, 1.0, error=1e-7,maxiter=1000)
1862 except:
1863 return momenta, 0
1864
1865 new_momenta = []
1866 for i,p in enumerate(momenta):
1867 new_momenta.append(
1868 FourMomentum(math.sqrt(newm[i]+chi**2*(p.E**2-oldm[i])),
1869 chi*p.px, chi*p.py, chi*p.pz))
1870
1871
1872
1873
1874
1875
1876
1877 jac = chi**(3*len(momenta)-3)
1878 jac *= reduce(operator.mul,[p.E/k.E for p,k in zip(momenta, new_momenta)],1)
1879 jac *= sum(p.norm_sq/p.E for p in momenta)
1880 jac /= sum(k.norm_sq/k.E for k in new_momenta)
1881
1882
1883 if tot_mom.pt2 > 1e-5:
1884 for i,m in enumerate(new_momenta):
1885 new_momenta[i] = m.boost_to_restframe(boost_back)
1886 return new_momenta, jac
1887
1888
1889
1890
1892 """routine to rescale the mass via RAMBO method. no internal mass preserve.
1893 sqrts is preserve (RAMBO algo)
1894 """
1895
1896 old_momenta = []
1897 new_masses = []
1898 change_mass = False
1899 for part in self:
1900 if part.status == 1:
1901 old_momenta.append(FourMomentum(part))
1902 new_masses.append(new_param_card.get_value('mass', abs(part.pid)))
1903 if not misc.equal(part.mass, new_masses[-1], 4, zero_limit=10):
1904 change_mass = True
1905
1906 if not change_mass:
1907 return 1
1908
1909 sqrts = self.sqrts
1910
1911
1912 new_mom, jac = self.mass_shuffle(old_momenta, sqrts, new_masses)
1913
1914
1915 ind =0
1916 for part in self:
1917 if part.status==1:
1918 part.E, part.px, part.py, part.pz, part.mass = \
1919 new_mom[ind].E, new_mom[ind].px, new_mom[ind].py, new_mom[ind].pz,new_mom[ind].mass
1920 ind+=1
1921 return jac
1922
1924 """routine to rescale the momenta to change the invariant mass"""
1925
1926 old_momenta = []
1927 incoming = []
1928 masses = []
1929 for part in self:
1930 if part.status == -1:
1931 incoming.append(FourMomentum(part))
1932 if part.status == 1:
1933 old_momenta.append(FourMomentum(part))
1934 masses.append(part.mass)
1935
1936 p_init = FourMomentum()
1937 p_inits = []
1938 n_init = 0
1939 for p in incoming:
1940 n_init +=1
1941 p_init += p
1942 p_inits.append(p)
1943 old_sqrts = p_init.mass
1944
1945 new_mom, jac = self.mass_shuffle(old_momenta, old_sqrts, masses, new_sqrts=new_sqrts)
1946
1947
1948 ind =0
1949 for part in self:
1950 if part.status==1:
1951 part.E, part.px, part.py, part.pz, part.mass = \
1952 new_mom[ind].E, new_mom[ind].px, new_mom[ind].py, new_mom[ind].pz,new_mom[ind].mass
1953 ind+=1
1954
1955
1956 p_init = FourMomentum()
1957 for part in self:
1958 if part.status==1:
1959 p_init += part
1960 if n_init == 1:
1961 for part in self:
1962 if part.status == -1:
1963 part.E, part.px, part.py, part.pz = \
1964 p_init.E, p_init.px, p_init.py, p_init.pz
1965 elif n_init ==2:
1966 if not misc.equal(p_init.px, 0) or not misc.equal(p_init.py, 0):
1967 raise Exception
1968 if not misc.equal(p_inits[0].px, 0) or not misc.equal(p_inits[0].py, 0):
1969 raise Exception
1970
1971
1972
1973
1974 eta = math.log(2*p_inits[0].E/old_sqrts)
1975 new_p = [[new_sqrts/2*math.exp(eta), 0., 0., new_sqrts/2*math.exp(eta)],
1976 [new_sqrts/2*math.exp(-eta), 0., 0., -new_sqrts/2*math.exp(-eta)]]
1977
1978 ind=0
1979 for part in self:
1980 if part.status == -1:
1981 part.E, part.px, part.py, part.pz = new_p[ind]
1982 ind+=1
1983 if ind ==2:
1984 break
1985 else:
1986 raise Exception
1987
1988 return jac
1989
1990
1992 """return a list with the helicities in the order asked for"""
1993
1994
1995 order = [list(get_order[0]), list(get_order[1])]
1996 out = [9] *(len(order[0])+len(order[1]))
1997 for i, part in enumerate(self):
1998 if part.status == 1:
1999 try:
2000 ind = order[1].index(part.pid)
2001 except ValueError, error:
2002 if not allow_reversed:
2003 raise error
2004 else:
2005 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2006 try:
2007 return self.get_helicity(order, False)
2008 except ValueError:
2009 raise error
2010 position = len(order[0]) + ind
2011 order[1][ind] = 0
2012 elif part.status == -1:
2013 try:
2014 ind = order[0].index(part.pid)
2015 except ValueError, error:
2016 if not allow_reversed:
2017 raise error
2018 else:
2019 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2020 try:
2021 return self.get_helicity(order, False)
2022 except ValueError:
2023 raise error
2024
2025 position = ind
2026 order[0][ind] = 0
2027 else:
2028 continue
2029 out[position] = int(part.helicity)
2030 return out
2031
2032
2034 """check the validity of the color structure"""
2035
2036
2037 color_index = collections.defaultdict(int)
2038 for particle in self:
2039 if particle.status in [-1,1]:
2040 if particle.color1:
2041 color_index[particle.color1] +=1
2042 if -7 < particle.pdg < 0:
2043 raise Exception, "anti-quark with color tag"
2044 if particle.color2:
2045 color_index[particle.color2] +=1
2046 if 7 > particle.pdg > 0:
2047 raise Exception, "quark with anti-color tag"
2048
2049
2050 for key,value in color_index.items():
2051 if value > 2:
2052 print self
2053 print key, value
2054 raise Exception, 'Wrong color_flow'
2055
2056
2057
2058 check = []
2059 popup_index = []
2060 for particle in self:
2061 mothers = []
2062 childs = []
2063 if particle.mother1:
2064 mothers.append(particle.mother1)
2065 if particle.mother2 and particle.mother2 is not particle.mother1:
2066 mothers.append(particle.mother2)
2067 if not mothers:
2068 continue
2069 if (particle.mother1.event_id, particle.mother2.event_id) in check:
2070 continue
2071 check.append((particle.mother1.event_id, particle.mother2.event_id))
2072
2073 childs = [p for p in self if p.mother1 is particle.mother1 and \
2074 p.mother2 is particle.mother2]
2075
2076 mcolors = []
2077 manticolors = []
2078 for m in mothers:
2079 if m.color1:
2080 if m.color1 in manticolors:
2081 manticolors.remove(m.color1)
2082 else:
2083 mcolors.append(m.color1)
2084 if m.color2:
2085 if m.color2 in mcolors:
2086 mcolors.remove(m.color2)
2087 else:
2088 manticolors.append(m.color2)
2089 ccolors = []
2090 canticolors = []
2091 for m in childs:
2092 if m.color1:
2093 if m.color1 in canticolors:
2094 canticolors.remove(m.color1)
2095 else:
2096 ccolors.append(m.color1)
2097 if m.color2:
2098 if m.color2 in ccolors:
2099 ccolors.remove(m.color2)
2100 else:
2101 canticolors.append(m.color2)
2102 for index in mcolors[:]:
2103 if index in ccolors:
2104 mcolors.remove(index)
2105 ccolors.remove(index)
2106 for index in manticolors[:]:
2107 if index in canticolors:
2108 manticolors.remove(index)
2109 canticolors.remove(index)
2110
2111 if mcolors != []:
2112
2113 if len(canticolors) + len(mcolors) != 3:
2114 logger.critical(str(self))
2115 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs])
2116 else:
2117 popup_index += canticolors
2118 elif manticolors != []:
2119
2120 if len(ccolors) + len(manticolors) != 3:
2121 logger.critical(str(self))
2122 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs])
2123 else:
2124 popup_index += ccolors
2125
2126
2127 if len(popup_index) != len(set(popup_index)):
2128 logger.critical(self)
2129 raise Exception, "Wrong color flow: identical poping-up index, %s" % (popup_index)
2130
2132 """two event are the same if they have the same momentum. other info are ignored"""
2133
2134 if other is None:
2135 return False
2136
2137 for i,p in enumerate(self):
2138 if p.E != other[i].E:
2139 return False
2140 elif p.pz != other[i].pz:
2141 return False
2142 elif p.px != other[i].px:
2143 return False
2144 elif p.py != other[i].py:
2145 return False
2146 return True
2147
2148
2150 """return a correctly formatted LHE event"""
2151
2152 out="""<event%(event_flag)s>
2153 %(scale)s
2154 %(particles)s
2155 %(comments)s
2156 %(tag)s
2157 %(reweight)s
2158 </event>
2159 """
2160 if event_id not in ['', None]:
2161 self.eventflag['event'] = str(event_id)
2162
2163 if self.eventflag:
2164 event_flag = ' %s' % ' '.join('%s="%s"' % (k,v) for (k,v) in self.eventflag.items())
2165 else:
2166 event_flag = ''
2167
2168 scale_str = "%2d %6d %+13.7e %14.8e %14.8e %14.8e" % \
2169 (self.nexternal,self.ievent,self.wgt,self.scale,self.aqed,self.aqcd)
2170
2171
2172 if self.reweight_data:
2173
2174 if set(self.reweight_data.keys()) != set(self.reweight_order):
2175 self.reweight_order += [k for k in self.reweight_data.keys() \
2176 if k not in self.reweight_order]
2177
2178 reweight_str = '<rwgt>\n%s\n</rwgt>' % '\n'.join(
2179 '<wgt id=\'%s\'> %+13.7e </wgt>' % (i, float(self.reweight_data[i]))
2180 for i in self.reweight_order if i in self.reweight_data)
2181 else:
2182 reweight_str = ''
2183
2184 tag_str = self.tag
2185 if hasattr(self, 'nloweight') and self.nloweight.modified:
2186 self.rewrite_nlo_weight()
2187 tag_str = self.tag
2188
2189 if self.matched_scale_data:
2190 tmp_scale = ' '.join(['pt_clust_%i=\"%s\"' % (i+1,v)
2191 for i,v in enumerate(self.matched_scale_data)
2192 if v!=-1])
2193 if tmp_scale:
2194 tag_str = "<scales %s></scales>%s" % (tmp_scale, self.tag)
2195
2196 if self.syscalc_data:
2197 keys= ['rscale', 'asrwt', ('pdfrwt', 'beam', '1'), ('pdfrwt', 'beam', '2'),
2198 'matchscale', 'totfact']
2199 sys_str = "<mgrwt>\n"
2200 template = """<%(key)s%(opts)s>%(values)s</%(key)s>\n"""
2201 for k in keys:
2202 if k not in self.syscalc_data:
2203 continue
2204 replace = {}
2205 replace['values'] = self.syscalc_data[k]
2206 if isinstance(k, str):
2207 replace['key'] = k
2208 replace['opts'] = ''
2209 else:
2210 replace['key'] = k[0]
2211 replace['opts'] = ' %s=\"%s\"' % (k[1],k[2])
2212 sys_str += template % replace
2213 sys_str += "</mgrwt>\n"
2214 reweight_str = sys_str + reweight_str
2215
2216 out = out % {'event_flag': event_flag,
2217 'scale': scale_str,
2218 'particles': '\n'.join([str(p) for p in self]),
2219 'tag': tag_str,
2220 'comments': self.comment,
2221 'reweight': reweight_str}
2222
2223 return re.sub('[\n]+', '\n', out)
2224
2225 - def get_momenta(self, get_order, allow_reversed=True):
2226 """return the momenta vector in the order asked for"""
2227
2228
2229 order = [list(get_order[0]), list(get_order[1])]
2230 out = [''] *(len(order[0])+len(order[1]))
2231 for i, part in enumerate(self):
2232 if part.status == 1:
2233 try:
2234 ind = order[1].index(part.pid)
2235 except ValueError, error:
2236 if not allow_reversed:
2237 raise error
2238 else:
2239 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2240 try:
2241 return self.get_momenta_str(order, False)
2242 except ValueError:
2243 raise error
2244 position = len(order[0]) + ind
2245 order[1][ind] = 0
2246 elif part.status == -1:
2247 try:
2248 ind = order[0].index(part.pid)
2249 except ValueError, error:
2250 if not allow_reversed:
2251 raise error
2252 else:
2253 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2254 try:
2255 return self.get_momenta_str(order, False)
2256 except ValueError:
2257 raise error
2258
2259 position = ind
2260 order[0][ind] = 0
2261 else:
2262 continue
2263
2264 out[position] = (part.E, part.px, part.py, part.pz)
2265
2266 return out
2267
2268
2281
2282
2284
2285 scale = 0
2286 for particle in self:
2287 if particle.status != 1:
2288 continue
2289 p=FourMomentum(particle)
2290 scale += math.sqrt(p.mass_sqr + p.pt**2)
2291
2292 return prefactor * scale
2293
2294
2296
2297 scale = 0
2298 for particle in self:
2299 if particle.status != 1:
2300 continue
2301 p = FourMomentum(particle)
2302 pt = p.pt
2303 if (pt>0):
2304 scale += p.E*pt/math.sqrt(pt**2+p.pz**2)
2305
2306 return prefactor * scale
2307
2308 @property
2311
2313
2314 scale = 0
2315 init = []
2316 for particle in self:
2317 if particle.status == -1:
2318 init.append(FourMomentum(particle))
2319 if len(init) == 1:
2320 return init[0].mass
2321 elif len(init)==2:
2322 return math.sqrt((init[0]+init[1])**2)
2323
2324
2325
2326
2328 """return the momenta str in the order asked for"""
2329
2330 out = self.get_momenta(get_order, allow_reversed)
2331
2332 format = '%.12f'
2333 format_line = ' '.join([format]*4) + ' \n'
2334 out = [format_line % one for one in out]
2335 out = ''.join(out).replace('e','d')
2336 return out
2337
2339 """A class to allow to read both gzip and not gzip file.
2340 containing only weight from pythia --generated by SysCalc"""
2341
2342 - def __new__(self, path, mode='r', *args, **opt):
2343 if path.endswith(".gz"):
2344 try:
2345 return gzip.GzipFile.__new__(WeightFileGzip, path, mode, *args, **opt)
2346 except IOError, error:
2347 raise
2348 except Exception, error:
2349 if mode == 'r':
2350 misc.gunzip(path)
2351 return file.__new__(WeightFileNoGzip, path[:-3], mode, *args, **opt)
2352 else:
2353 return file.__new__(WeightFileNoGzip, path, mode, *args, **opt)
2354
2355
2356 - def __init__(self, path, mode='r', *args, **opt):
2357 """open file and read the banner [if in read mode]"""
2358
2359 super(EventFile, self).__init__(path, mode, *args, **opt)
2360 self.banner = ''
2361 if mode == 'r':
2362 line = ''
2363 while '</header>' not in line.lower():
2364 try:
2365 line = super(EventFile, self).next()
2366 except StopIteration:
2367 self.seek(0)
2368 self.banner = ''
2369 break
2370 if "<event" in line.lower():
2371 self.seek(0)
2372 self.banner = ''
2373 break
2374
2375 self.banner += line
2376
2380
2383
2386 """a convenient object for 4-momenta operation"""
2387
2388 - def __init__(self, obj=0, px=0, py=0, pz=0, E=0):
2389 """initialize the four momenta"""
2390
2391 if obj is 0 and E:
2392 obj = E
2393
2394 if isinstance(obj, (FourMomentum, Particle)):
2395 px = obj.px
2396 py = obj.py
2397 pz = obj.pz
2398 E = obj.E
2399 elif isinstance(obj, (list, tuple)):
2400 assert len(obj) ==4
2401 E = obj[0]
2402 px = obj[1]
2403 py = obj[2]
2404 pz = obj[3]
2405 elif isinstance(obj, str):
2406 obj = [float(i) for i in obj.split()]
2407 assert len(obj) ==4
2408 E = obj[0]
2409 px = obj[1]
2410 py = obj[2]
2411 pz = obj[3]
2412 else:
2413 E =obj
2414
2415
2416 self.E = float(E)
2417 self.px = float(px)
2418 self.py = float(py)
2419 self.pz = float(pz)
2420
2421 @property
2423 """return the mass"""
2424 return math.sqrt(max(self.E**2 - self.px**2 - self.py**2 - self.pz**2,0))
2425
2426 @property
2428 """return the mass square"""
2429 return self.E**2 - self.px**2 - self.py**2 - self.pz**2
2430
2431 @property
2433 return math.sqrt(max(0, self.pt2))
2434
2435 @property
2437 norm = math.sqrt(self.px**2 + self.py**2+self.pz**2)
2438 return 0.5* math.log((norm - self.pz) / (norm + self.pz))
2439
2440 @property
2442 return 0.5* math.log((self.E +self.pz) / (self.E - self.pz))
2443
2444
2445 @property
2447 """ return the pt square """
2448
2449 return self.px**2 + self.py**2
2450
2451 @property
2453 """ return |\vec p| """
2454 return math.sqrt(self.px**2 + self.py**2 + self.pz**2)
2455
2456 @property
2458 """ return |\vec p|^2 """
2459 return self.px**2 + self.py**2 + self.pz**2
2460
2461 @property
2463 """return the mass square"""
2464 import math
2465 return math.atan(math.sqrt((self.px**2+self.py**2)/self.pz**2))
2466
2467
2469
2470 assert isinstance(obj, FourMomentum)
2471 new = FourMomentum(self.E+obj.E,
2472 self.px + obj.px,
2473 self.py + obj.py,
2474 self.pz + obj.pz)
2475 return new
2476
2478 """update the object with the sum"""
2479 self.E += obj.E
2480 self.px += obj.px
2481 self.py += obj.py
2482 self.pz += obj.pz
2483 return self
2484
2486
2487 assert isinstance(obj, FourMomentum)
2488 new = FourMomentum(self.E-obj.E,
2489 self.px - obj.px,
2490 self.py - obj.py,
2491 self.pz - obj.pz)
2492 return new
2493
2495 """update the object with the sum"""
2496 self.E -= obj.E
2497 self.px -= obj.px
2498 self.py -= obj.py
2499 self.pz -= obj.pz
2500 return self
2501
2503 if isinstance(obj, FourMomentum):
2504 return self.E*obj.E - self.px *obj.px - self.py * obj.py - self.pz * obj.pz
2505 elif isinstance(obj, (float, int)):
2506 return FourMomentum(obj*self.E,obj*self.px,obj*self.py,obj*self.pz )
2507 else:
2508 raise NotImplemented
2509 __rmul__ = __mul__
2510
2512 assert power in [1,2]
2513
2514 if power == 1:
2515 return FourMomentum(self)
2516 elif power == 2:
2517 return self.mass_sqr
2518
2520 return 'FourMomentum(%s,%s,%s,%s)' % (self.E, self.px, self.py,self.pz)
2521
2522 - def __str__(self, mode='python'):
2523 if mode == 'python':
2524 return self.__repr__()
2525 elif mode == 'fortran':
2526 return '%.10e %.10e %.10e %.10e' % self.get_tuple()
2527
2529 return (self.E, self.px, self.py,self.pz)
2530
2532 """mom 4-momenta is suppose to be given in the rest frame of this 4-momenta.
2533 the output is the 4-momenta in the frame of this 4-momenta
2534 function copied from HELAS routine.
2535 if the current momenta is (E,\vec{p}), in order to go to the rest frame
2536 of the current particle, mom should be (E, -\vec{p})
2537 """
2538
2539 pnorm = mom.px**2 + mom.py**2 + mom.pz**2
2540 if pnorm:
2541 s3product = self.px * mom.px + self.py * mom.py + self.pz * mom.pz
2542 mass = mom.mass
2543 lf = (self.E + (mom.E - mass) * s3product / pnorm ) / mass
2544 return FourMomentum(E=(self.E*mom.E+s3product)/mass,
2545 px=self.px + mom.px * lf,
2546 py=self.py + mom.py * lf,
2547 pz=self.pz + mom.pz * lf)
2548 else:
2549 return FourMomentum(mom)
2550
2551 - def zboost(self, pboost=None, E=0, pz=0):
2552 """Both momenta should be in the same frame.
2553 The boost perform correspond to the boost required to set pboost at
2554 rest (only z boost applied).
2555 """
2556 if isinstance(pboost, FourMomentum):
2557 E = pboost.E
2558 pz = pboost.pz
2559
2560
2561 gamma = E / math.sqrt(E**2-pz**2)
2562 gammabeta = pz / math.sqrt(E**2-pz**2)
2563
2564 out = FourMomentum([gamma*self.E - gammabeta*self.pz,
2565 self.px,
2566 self.py,
2567 gamma*self.pz - gammabeta*self.E])
2568
2569 if abs(out.pz) < 1e-6 * out.E:
2570 out.pz = 0
2571 return out
2572
2574 """apply the boost transformation such that pboost is at rest in the new frame.
2575 First apply a rotation to allign the pboost to the z axis and then use
2576 zboost routine (see above)
2577 """
2578
2579 if pboost.px == 0 == pboost.py:
2580 out = self.zboost(E=pboost.E,pz=pboost.pz)
2581 return out
2582
2583
2584
2585
2586
2587
2588
2589 p = math.sqrt( pboost.px**2 + pboost.py**2+ pboost.pz**2)
2590 cosF = pboost.pz / p
2591 sinF = math.sqrt(1-cosF**2)
2592 sinT = pboost.py/p/sinF
2593 cosT = pboost.px/p/sinF
2594
2595 out=FourMomentum([self.E,
2596 self.px*cosT*cosF + self.py*sinT*cosF-self.pz*sinF,
2597 -self.px*sinT+ self.py*cosT,
2598 self.px*cosT*sinF + self.py*sinT*sinF + self.pz*cosF
2599 ])
2600 out = out.zboost(E=pboost.E,pz=p)
2601 return out
2602
2607
2608 - def __init__(self, input, real_type=(1,11)):
2609 """ """
2610
2611 self.real_type = real_type
2612 if isinstance(input, str):
2613 self.parse(input)
2614
2615 - def __str__(self, mode='display'):
2616
2617 if mode == 'display':
2618 out = """ pwgt: %(pwgt)s
2619 born, real : %(born)s %(real)s
2620 pdgs : %(pdgs)s
2621 bjks : %(bjks)s
2622 scales**2, gs: %(scales2)s %(gs)s
2623 born/real related : %(born_related)s %(real_related)s
2624 type / nfks : %(type)s %(nfks)s
2625 to merge : %(to_merge_pdg)s in %(merge_new_pdg)s
2626 ref_wgt : %(ref_wgt)s""" % self.__dict__
2627 return out
2628 elif mode == 'formatted':
2629 format_var = []
2630 variable = []
2631
2632 def to_add_full(f, v, format_var, variable):
2633 """ function to add to the formatted output"""
2634 if isinstance(v, list):
2635 format_var += [f]*len(v)
2636 variable += v
2637 else:
2638 format_var.append(f)
2639 variable.append(v)
2640 to_add = lambda x,y: to_add_full(x,y, format_var, variable)
2641
2642 to_add('%.10e', [p*self.bias_wgt for p in self.pwgt])
2643 to_add('%.10e', self.born)
2644 to_add('%.10e', self.real)
2645 to_add('%i', self.nexternal)
2646 to_add('%i', self.pdgs)
2647 to_add('%i', self.qcdpower)
2648 to_add('%.10e', self.bjks)
2649 to_add('%.10e', self.scales2)
2650 to_add('%.10e', self.gs)
2651 to_add('%i', [self.born_related, self.real_related])
2652 to_add('%i' , [self.type, self.nfks])
2653 to_add('%i' , self.to_merge_pdg)
2654 to_add('%i', self.merge_new_pdg)
2655 to_add('%.10e', self.ref_wgt*self.bias_wgt)
2656 to_add('%.10e', self.bias_wgt)
2657 return ' '.join(format_var) % tuple(variable)
2658
2659
2660 - def parse(self, text, keep_bias=False):
2661 """parse the line and create the related object.
2662 keep bias allow to not systematically correct for the bias in the written information"""
2663
2664
2665
2666 data = text.split()
2667
2668
2669
2670
2671 self.pwgt = [float(f) for f in data[:3]]
2672
2673
2674
2675
2676
2677
2678
2679
2680
2681
2682
2683
2684
2685
2686
2687
2688
2689
2690 self.born = float(data[3])
2691 self.real = float(data[4])
2692
2693
2694 self.nexternal = int(data[5])
2695
2696
2697 self.pdgs = [int(i) for i in data[6:6+self.nexternal]]
2698 flag = 6+self.nexternal
2699
2700
2701 self.qcdpower = int(data[flag])
2702
2703
2704 self.bjks = [float(f) for f in data[flag+1:flag+3]]
2705
2706
2707 self.scales2 = [float(f) for f in data[flag+3:flag+6]]
2708
2709
2710 self.gs = float(data[flag+6])
2711
2712
2713
2714
2715
2716
2717
2718 self.born_related = int(data[flag+7])
2719 self.real_related = int(data[flag+8])
2720
2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
2742
2743 self.type = int(data[flag+9])
2744
2745
2746
2747
2748 self.nfks = int(data[flag+10])
2749
2750
2751
2752
2753
2754
2755
2756 self.to_merge_pdg = [int (f) for f in data[flag+11:flag+13]]
2757
2758
2759 self.merge_new_pdg = int(data[flag+13])
2760
2761
2762
2763
2764
2765 self.ref_wgt = float(data[flag+14])
2766
2767
2768
2769 try:
2770 self.bias_wgt = float(data[flag+15])
2771 except IndexError:
2772 self.bias_wgt = 1.0
2773
2774 if not keep_bias:
2775 self.ref_wgt /= self.bias_wgt
2776 self.pwgt = [p/self.bias_wgt for p in self.pwgt]
2777
2778
2779 if self.type in self.real_type:
2780 self.momenta_config = self.real_related
2781 else:
2782 self.momenta_config = self.born_related
2783
2786
2788
2789
2790 - def __init__(self, momenta, wgts, event, real_type=(1,11)):
2791
2792 list.__init__(self, momenta)
2793 assert self
2794 self.soft = False
2795 self.wgts = wgts
2796 self.pdgs = list(wgts[0].pdgs)
2797 self.event = event
2798 self.real_type = real_type
2799
2800 if wgts[0].momenta_config == wgts[0].born_related:
2801
2802 ind1, ind2 = [ind-1 for ind in wgts[0].to_merge_pdg]
2803 if ind1> ind2:
2804 ind1, ind2 = ind2, ind1
2805 if ind1 >= sum(1 for p in event if p.status==-1):
2806 new_p = self[ind1] + self[ind2]
2807 else:
2808 new_p = self[ind1] - self[ind2]
2809 self.pop(ind1)
2810 self.insert(ind1, new_p)
2811 self.pop(ind2)
2812 self.pdgs.pop(ind1)
2813 self.pdgs.insert(ind1, wgts[0].merge_new_pdg )
2814 self.pdgs.pop(ind2)
2815
2816
2817 elif any(w.type in self.real_type for w in wgts):
2818 if any(w.type not in self.real_type for w in wgts):
2819 raise Exception
2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
2831
2832 else:
2833 raise Exception
2834
2836 """check that the propagator associated to ij is not too light
2837 [related to soft-collinear singularity]"""
2838
2839 if threshold is None:
2840 threshold = 1e-8
2841
2842 if ind1> ind2:
2843 ind1, ind2 = ind2, ind1
2844 if ind1 >= nb_init:
2845 new_p = self[ind1] + self[ind2]
2846 else:
2847 new_p = self[ind1] - self[ind2]
2848
2849 inv_mass = new_p.mass_sqr
2850 if nb_init == 2:
2851 shat = (self[0]+self[1]).mass_sqr
2852 else:
2853 shat = self[0].mass_sqr
2854
2855
2856 if (abs(inv_mass)/shat < threshold):
2857 return True
2858 else:
2859 return False
2860
2861
2864
2866 """ return the tag and order for this basic event"""
2867 (initial, _), _ = self.event.get_tag_and_order()
2868 order = self.get_pdg_code()
2869
2870
2871 initial, out = order[:len(initial)], order[len(initial):]
2872 initial.sort()
2873 out.sort()
2874 return (tuple(initial), tuple(out)), order
2875
2876 - def get_momenta(self, get_order, allow_reversed=True):
2877 """return the momenta vector in the order asked for"""
2878
2879
2880 order = [list(get_order[0]), list(get_order[1])]
2881 out = [''] *(len(order[0])+len(order[1]))
2882 pdgs = self.get_pdg_code()
2883 for pos, part in enumerate(self):
2884 if pos < len(get_order[0]):
2885 try:
2886 ind = order[0].index(pdgs[pos])
2887 except ValueError, error:
2888 if not allow_reversed:
2889 raise error
2890 else:
2891 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2892 try:
2893 return self.get_momenta(order, False)
2894 except ValueError:
2895 raise error
2896
2897
2898 position = ind
2899 order[0][ind] = 0
2900 else:
2901 try:
2902 ind = order[1].index(pdgs[pos])
2903 except ValueError, error:
2904 if not allow_reversed:
2905 raise error
2906 else:
2907 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2908 try:
2909 return self.get_momenta(order, False)
2910 except ValueError:
2911 raise error
2912 position = len(order[0]) + ind
2913 order[1][ind] = 0
2914
2915 out[position] = (part.E, part.px, part.py, part.pz)
2916
2917 return out
2918
2919
2921 return [9] * len(self)
2922
2923 @property
2925 return self.event.aqcd
2926
2928
2929 scale = 0
2930 for particle in self:
2931 p = particle
2932 scale += math.sqrt(max(0, p.mass_sqr + p.pt**2))
2933
2934 return prefactor * scale
2935
2937
2938 scale = 0
2939 for particle in self:
2940 p = particle
2941 pt = p.pt
2942 if (pt>0):
2943 scale += p.E*pt/math.sqrt(pt**2+p.pz**2)
2944
2945 return prefactor * scale
2946
2947
2949
2950 scale = 0
2951 nb_init = 0
2952 for particle in event:
2953 if particle.status == -1:
2954 nb_init+=1
2955 if nb_init == 1:
2956 return self[0].mass
2957 elif nb_init==2:
2958 return math.sqrt((self[0]+self[1])**2)
2959
2960
2961
2962
2963 - def __init__(self, input, event, real_type=(1,11), threshold=None):
2964
2965 self.real_type = real_type
2966 self.event = event
2967 self.total_wgt = 0.
2968 self.nb_event = 0
2969 self.nb_wgts = 0
2970 self.threshold = threshold
2971 self.modified = False
2972
2973
2974 if isinstance(input, str):
2975 self.parse(input)
2976
2977
2978
2980 """create the object from the string information (see example below)"""
2981
2982
2983
2984
2985
2986
2987
2988
2989
2990
2991
2992
2993
2994
2995
2996
2997
2998
2999
3000
3001
3002 text = text.lower().replace('d','e')
3003 all_line = text.split('\n')
3004
3005 first_line =''
3006 while not first_line.strip():
3007 first_line = all_line.pop(0)
3008
3009 wgt, nb_wgt, nb_event, _ = first_line.split()
3010 self.total_wgt = float(wgt.replace('d','e'))
3011 nb_wgt, nb_event = int(nb_wgt), int(nb_event)
3012 self.nb_wgt, self.nb_event = nb_wgt, nb_event
3013
3014 momenta = []
3015 self.momenta = momenta
3016 wgts = []
3017 for line in all_line:
3018 data = line.split()
3019 if len(data) == 4:
3020 p = FourMomentum(data)
3021 momenta.append(p)
3022 elif len(data)>0:
3023 wgt = OneNLOWeight(line, real_type=self.real_type)
3024 wgts.append(wgt)
3025
3026 assert len(wgts) == int(nb_wgt)
3027
3028 get_weights_for_momenta = dict( (i,[]) for i in range(1,nb_event+1) )
3029 size_momenta = 0
3030 for wgt in wgts:
3031 if wgt.momenta_config in get_weights_for_momenta:
3032 get_weights_for_momenta[wgt.momenta_config].append(wgt)
3033 else:
3034 if size_momenta == 0: size_momenta = wgt.nexternal
3035 assert size_momenta == wgt.nexternal
3036 get_weights_for_momenta[wgt.momenta_config] = [wgt]
3037
3038 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)
3039
3040
3041 for key in range(1, nb_event+1):
3042 wgts = get_weights_for_momenta[key]
3043 if not wgts:
3044 continue
3045 if size_momenta == 0: size_momenta = wgts[0].nexternal
3046 p = momenta[size_momenta*(key-1):key*size_momenta]
3047 evt = self.BasicEvent(p, wgts, self.event, self.real_type)
3048 if len(evt) == size_momenta:
3049 for wgt in wgts:
3050 if not wgt.type in self.real_type:
3051 continue
3052 if evt.check_fks_singularity(wgt.to_merge_pdg[0]-1,
3053 wgt.to_merge_pdg[1]-1,
3054 nb_init=sum(1 for p in self.event if p.status==-1),
3055 threshold=self.threshold):
3056 get_weights_for_momenta[wgt.momenta_config].remove(wgt)
3057 get_weights_for_momenta[wgt.born_related].append(wgt)
3058 wgt.momenta_config = wgt.born_related
3059
3060 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)
3061
3062 self.cevents = []
3063 for key in range(1, nb_event+1):
3064 if key in get_weights_for_momenta:
3065 wgt = get_weights_for_momenta[key]
3066 if not wgt:
3067 continue
3068 pdg_to_event = {}
3069 for w in wgt:
3070 pdgs = w.pdgs
3071 if w.momenta_config == w.born_related:
3072 pdgs = list(pdgs)
3073 ind1, ind2 = [ind-1 for ind in w.to_merge_pdg]
3074 if ind1> ind2:
3075 ind1, ind2 = ind2, ind1
3076 pdgs.pop(ind1)
3077 pdgs.insert(ind1, w.merge_new_pdg )
3078 pdgs.pop(ind2)
3079 pdgs = tuple(pdgs)
3080 if pdgs not in pdg_to_event:
3081 p = momenta[size_momenta*(key-1):key*size_momenta]
3082 evt = self.BasicEvent(p, [w], self.event, self.real_type)
3083 self.cevents.append(evt)
3084 pdg_to_event[pdgs] = evt
3085 else:
3086 pdg_to_event[pdgs].wgts.append(w)
3087
3088 if __debug__:
3089 nb_wgt_check = 0
3090 for cevt in self.cevents:
3091 nb_wgt_check += len(cevt.wgts)
3092 assert nb_wgt_check == int(nb_wgt)
3093
3094
3095
3096 if '__main__' == __name__:
3097
3098 if False:
3099 lhe = EventFile('unweighted_events.lhe.gz')
3100
3101 start = time.time()
3102 for event in lhe:
3103 event.parse_lo_weight()
3104 print 'old method -> ', time.time()-start
3105 lhe = EventFile('unweighted_events.lhe.gz')
3106
3107 start = time.time()
3108 for event in lhe:
3109 event.parse_lo_weight_test()
3110 print 'new method -> ', time.time()-start
3111
3112
3113
3114 if False:
3115 start = time
3116 lhe = EventFile('unweighted_events.lhe.gz')
3117 output = open('output_events.lhe', 'w')
3118
3119 output.write(lhe.banner)
3120
3121 for event in lhe:
3122 for particle in event:
3123
3124 particle.mass = 0
3125 particle.vtim = 2
3126
3127
3128 output.write(str(event))
3129 output.write('</LesHouchesEvent>\n')
3130
3131
3132 if False:
3133 lhe = EventFile('unweighted_events.lhe.gz')
3134 import matplotlib.pyplot as plt
3135 import matplotlib.gridspec as gridspec
3136 nbins = 100
3137
3138 nb_pass = 0
3139 data = []
3140 for event in lhe:
3141 etaabs = 0
3142 etafinal = 0
3143 for particle in event:
3144 if particle.status==1:
3145 p = FourMomentum(particle)
3146 eta = p.pseudorapidity
3147 if abs(eta) > etaabs:
3148 etafinal = eta
3149 etaabs = abs(eta)
3150 if etaabs < 4:
3151 data.append(etafinal)
3152 nb_pass +=1
3153
3154
3155 print nb_pass
3156 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1])
3157 gs1.update(wspace=0, hspace=0)
3158 ax = plt.subplot(gs1[0])
3159
3160 n, bins, patches = ax.hist(data, nbins, histtype='step', label='original')
3161 ax_c = ax.twinx()
3162 ax_c.set_ylabel('MadGraph5_aMC@NLO')
3163 ax_c.yaxis.set_label_coords(1.01, 0.25)
3164 ax_c.set_yticks(ax.get_yticks())
3165 ax_c.set_yticklabels([])
3166 ax.set_xlim([-4,4])
3167 print "bin value:", n
3168 print "start/end point of bins", bins
3169 plt.axis('on')
3170 plt.xlabel('weight ratio')
3171 plt.show()
3172
3173
3174
3175 if False:
3176 lhe = EventFile('unweighted_events.lhe')
3177 import matplotlib.pyplot as plt
3178 import matplotlib.gridspec as gridspec
3179 nbins = 100
3180
3181
3182 mtau, wtau = 1.777, 4.027000e-13
3183 nb_pass = 0
3184 data, data2, data3 = [], [], []
3185 for event in lhe:
3186 nb_pass +=1
3187 if nb_pass > 10000:
3188 break
3189 tau1 = FourMomentum()
3190 tau2 = FourMomentum()
3191 for part in event:
3192 if part.pid in [-12,11,16]:
3193 momenta = FourMomentum(part)
3194 tau1 += momenta
3195 elif part.pid == 15:
3196 tau2 += FourMomentum(part)
3197
3198 if abs((mtau-tau2.mass())/wtau)<1e6 and tau2.mass() >1:
3199 data.append((tau1.mass()-mtau)/wtau)
3200 data2.append((tau2.mass()-mtau)/wtau)
3201 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1])
3202 gs1.update(wspace=0, hspace=0)
3203 ax = plt.subplot(gs1[0])
3204
3205 n, bins, patches = ax.hist(data2, nbins, histtype='step', label='original')
3206 n2, bins2, patches2 = ax.hist(data, bins=bins, histtype='step',label='reconstructed')
3207 import cmath
3208
3209 breit = lambda m : math.sqrt(4*math.pi)*1/(((m)**2-mtau**2)**2+(mtau*wtau)**2)*wtau
3210
3211 data3 = [breit(mtau + x*wtau)*wtau*16867622.6624*50 for x in bins]
3212
3213 ax.plot(bins, data3,label='breit-wigner')
3214
3215 ax.legend()
3216
3217 ax_c = ax.twinx()
3218 ax_c.set_ylabel('MadGraph5_aMC@NLO')
3219 ax_c.yaxis.set_label_coords(1.01, 0.25)
3220 ax_c.set_yticks(ax.get_yticks())
3221 ax_c.set_yticklabels([])
3222
3223 plt.title('invariant mass of tau LHE/reconstructed')
3224 plt.axis('on')
3225 ax.set_xticklabels([])
3226
3227 ax = plt.subplot(gs1[1])
3228 data4 = [n[i]/(data3[i]) for i in range(nbins)]
3229 ax.plot(bins, data4 + [0] , 'b')
3230 data4 = [n2[i]/(data3[i]) for i in range(nbins)]
3231 ax.plot(bins, data4 + [0] , 'g')
3232 ax.set_ylim([0,2])
3233
3234 tick = ax.get_yticks()
3235 ax.set_yticks(tick[:-1])
3236
3237
3238 plt.axis('on')
3239 plt.xlabel('(M - Mtau)/Wtau')
3240 plt.show()
3241