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