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
812
814 currpos = super(EventFileGzip, self).tell()
815 if not currpos:
816 currpos = self.size
817 return currpos
818
820 """A way to read a standard event file"""
821
822 - def close(self,*args, **opts):
827
829 """a class to read simultaneously multiple file and read them in mixing them.
830 Unweighting can be done at the same time.
831 The number of events in each file need to be provide in advance
832 (if not provide the file is first read to find that number"""
833
834 - def __new__(cls, start_list=[],parse=True):
836
837 - def __init__(self, start_list=[], parse=True):
838 """if trunc_error is define here then this allow
839 to only read all the files twice and not three times."""
840 self.eventgroup = False
841 self.files = []
842 self.parsefile = parse
843 self.banner = ''
844 self.initial_nb_events = []
845 self.total_event_in_files = 0
846 self.curr_nb_events = []
847 self.allcross = []
848 self.error = []
849 self.across = []
850 self.scales = []
851 if start_list:
852 if parse:
853 for p in start_list:
854 self.add(p)
855 else:
856 self.files = start_list
857 self._configure = False
858
859 - def close(self,*args,**opts):
862
863 - def add(self, path, cross, error, across, nb_event=0, scale=1):
864 """ add a file to the pool, across allow to reweight the sum of weight
865 in the file to the given cross-section
866 """
867
868 if across == 0:
869
870 return
871
872 obj = EventFile(path)
873 obj.eventgroup = self.eventgroup
874 if len(self.files) == 0 and not self.banner:
875 self.banner = obj.banner
876 self.curr_nb_events.append(0)
877 self.initial_nb_events.append(0)
878 self.allcross.append(cross)
879 self.across.append(across)
880 self.error.append(error)
881 self.scales.append(scale)
882 self.files.append(obj)
883 if nb_event:
884 obj.len = nb_event
885 self._configure = False
886 return obj
887
890
892
893 if not self._configure:
894 self.configure()
895
896 remaining_event = self.total_event_in_files - sum(self.curr_nb_events)
897 if remaining_event == 0:
898 raise StopIteration
899
900 nb_event = random.randint(1, remaining_event)
901 sum_nb=0
902 for i, obj in enumerate(self.files):
903 sum_nb += self.initial_nb_events[i] - self.curr_nb_events[i]
904 if nb_event <= sum_nb:
905 self.curr_nb_events[i] += 1
906 event = obj.next()
907 if not self.eventgroup:
908 event.sample_scale = self.scales[i]
909 else:
910 for evt in event:
911 evt.sample_scale = self.scales[i]
912 return event
913 else:
914 raise Exception
915
916
918 """define the part of the init_banner"""
919
920 if not self.banner:
921 return
922
923
924 grouped_cross = {}
925 grouped_error = {}
926 for i,ff in enumerate(self.files):
927 filename = ff.name
928 from_init = False
929 Pdir = [P for P in filename.split(os.path.sep) if P.startswith('P')]
930 if Pdir:
931 Pdir = Pdir[-1]
932 group = Pdir.split("_")[0][1:]
933 if not group.isdigit():
934 from_init = True
935 else:
936 from_init = True
937
938 if not from_init:
939 if group in grouped_cross:
940 grouped_cross[group] += self.allcross[i]
941 grouped_error[group] += self.error[i]**2
942 else:
943 grouped_cross[group] = self.allcross[i]
944 grouped_error[group] = self.error[i]**2
945 else:
946 ban = banner_mod.Banner(ff.banner)
947 for line in ban['init'].split('\n'):
948 splitline = line.split()
949 if len(splitline)==4:
950 cross, error, _, group = splitline
951 if int(group) in grouped_cross:
952 grouped_cross[group] += float(cross)
953 grouped_error[group] += float(error)**2
954 else:
955 grouped_cross[group] = float(cross)
956 grouped_error[group] = float(error)**2
957 nb_group = len(grouped_cross)
958
959
960 try:
961 run_card = self.banner.run_card
962 except:
963 run_card = self.banner.charge_card("run_card")
964
965 init_information = run_card.get_banner_init_information()
966
967 if proc_charac and proc_charac['ninitial'] == 1:
968
969 init_information = run_card.get_banner_init_information()
970 event = self.next()
971 init_information["idbmup1"] = event[0].pdg
972 init_information["ebmup1"] = event[0].mass
973 init_information["idbmup2"] = 0
974 init_information["ebmup2"] = 0
975 self.seek(0)
976 else:
977
978 if init_information["idbmup1"] == 0:
979 event = self.next()
980 init_information["idbmup1"]= event[0].pdg
981 if init_information["idbmup2"] == 0:
982 init_information["idbmup2"]= event[1].pdg
983 self.seek(0)
984 if init_information["idbmup2"] == 0:
985 event = self.next()
986 init_information["idbmup2"] = event[1].pdg
987 self.seek(0)
988
989 init_information["nprup"] = nb_group
990
991 if run_card["lhe_version"] < 3:
992 init_information["generator_info"] = ""
993 else:
994 init_information["generator_info"] = "<generator name='MadGraph5_aMC@NLO' version='%s'>please cite 1405.0301 </generator>\n" \
995 % misc.get_pkg_info()['version']
996
997
998 cross_info = "%(cross)e %(error)e %(wgt)e %(id)i"
999 init_information["cross_info"] = []
1000 for id in grouped_cross:
1001 conv = {"id": int(id), "cross": grouped_cross[id], "error": math.sqrt(grouped_error[id]),
1002 "wgt": wgt}
1003 init_information["cross_info"].append( cross_info % conv)
1004 init_information["cross_info"] = '\n'.join(init_information["cross_info"])
1005 init_information['lha_stra'] = -1 * abs(lha_strategy)
1006
1007 template_init =\
1008 """ %(idbmup1)i %(idbmup2)i %(ebmup1)e %(ebmup2)e %(pdfgup1)i %(pdfgup2)i %(pdfsup1)i %(pdfsup2)i %(lha_stra)i %(nprup)i
1009 %(cross_info)s
1010 %(generator_info)s
1011 """
1012
1013 self.banner["init"] = template_init % init_information
1014
1015
1016
1018 """ scan once the file to return
1019 - the list of the hightest weight (of size trunc_error*NB_EVENT
1020 - the cross-section by type of process
1021 - the total number of events in the files
1022 In top of that it initialise the information for the next routine
1023 to determine how to choose which file to read
1024 """
1025 self.seek(0)
1026 all_wgt = []
1027 total_event = 0
1028 sum_cross = collections.defaultdict(int)
1029 for i,f in enumerate(self.files):
1030 nb_event = 0
1031
1032
1033 cross = collections.defaultdict(int)
1034 new_wgt =[]
1035 for event in f:
1036 nb_event += 1
1037 total_event += 1
1038 event.sample_scale = 1
1039 wgt = getwgt(event)
1040 cross['all'] += wgt
1041 cross['abs'] += abs(wgt)
1042 cross[event.ievent] += wgt
1043 new_wgt.append(abs(wgt))
1044
1045 if nb_event % 20000 == 0:
1046 new_wgt.sort()
1047
1048 nb_keep = max(20, int(nb_event*trunc_error*15))
1049 new_wgt = new_wgt[-nb_keep:]
1050 if nb_event == 0:
1051 raise Exception
1052
1053 self.initial_nb_events[i] = nb_event
1054 self.scales[i] = self.across[i]/cross['abs'] if self.across[i] else 1
1055
1056
1057 for key in cross:
1058 sum_cross[key] += cross[key]* self.scales[i]
1059 all_wgt +=[self.scales[i] * w for w in new_wgt]
1060 all_wgt.sort()
1061 nb_keep = max(20, int(total_event*trunc_error*10))
1062 all_wgt = all_wgt[-nb_keep:]
1063
1064 self.total_event_in_files = total_event
1065
1066 all_wgt.sort()
1067
1068 nb_keep = max(20, int(total_event*trunc_error*10))
1069 all_wgt = all_wgt[-nb_keep:]
1070 self.seek(0)
1071 self._configure = True
1072 return all_wgt, sum_cross, total_event
1073
1080
1082
1083 return len(self.files)
1084
1085 - def seek(self, pos):
1086 """ """
1087
1088 if pos !=0:
1089 raise Exception
1090 for i in range(len(self)):
1091 self.curr_nb_events[i] = 0
1092 for f in self.files:
1093 f.seek(pos)
1094
1095 - def unweight(self, outputpath, get_wgt, **opts):
1096 """unweight the current file according to wgt information wgt.
1097 which can either be a fct of the event or a tag in the rwgt list.
1098 max_wgt allow to do partial unweighting.
1099 trunc_error allow for dynamical partial unweighting
1100 event_target reweight for that many event with maximal trunc_error.
1101 (stop to write event when target is reached)
1102 """
1103
1104 if isinstance(get_wgt, str):
1105 unwgt_name =get_wgt
1106 def get_wgt_multi(event):
1107 event.parse_reweight()
1108 return event.reweight_data[unwgt_name] * event.sample_scale
1109 else:
1110 unwgt_name = get_wgt.func_name
1111 get_wgt_multi = lambda event: get_wgt(event) * event.sample_scale
1112
1113
1114 if 'proc_charac' in opts:
1115 if opts['proc_charac']:
1116 proc_charac = opts['proc_charac']
1117 else:
1118 proc_charac=None
1119 del opts['proc_charac']
1120 else:
1121 proc_charac = None
1122
1123 if 'event_target' in opts and opts['event_target']:
1124 if 'normalization' in opts:
1125 if opts['normalization'] == 'sum':
1126 new_wgt = sum(self.across)/opts['event_target']
1127 strategy = 3
1128 elif opts['normalization'] == 'average':
1129 strategy = 4
1130 new_wgt = sum(self.across)
1131 elif opts['normalization'] == 'unit':
1132 strategy =3
1133 new_wgt = 1.
1134 else:
1135 strategy = 4
1136 new_wgt = sum(self.across)
1137 self.define_init_banner(new_wgt, strategy, proc_charac=proc_charac)
1138 self.written_weight = new_wgt
1139 elif 'write_init' in opts and opts['write_init']:
1140 self.define_init_banner(0,0, proc_charac=proc_charac)
1141 del opts['write_init']
1142 return super(MultiEventFile, self).unweight(outputpath, get_wgt_multi, **opts)
1143
1144 - def write(self, path, random=False, banner=None, get_info=False):
1193
1202
1203
1204
1205 -class Event(list):
1206 """Class storing a single event information (list of particles + global information)"""
1207
1208 warning_order = True
1209
1211 """The initialization of an empty Event (or one associate to a text file)"""
1212 list.__init__(self)
1213
1214
1215 self.nexternal = 0
1216 self.ievent = 0
1217 self.wgt = 0
1218 self.aqcd = 0
1219 self.scale = 0
1220 self.aqed = 0
1221 self.aqcd = 0
1222
1223 self.tag = ''
1224 self.eventflag = {}
1225 self.comment = ''
1226 self.reweight_data = {}
1227 self.matched_scale_data = None
1228 self.syscalc_data = {}
1229 if text:
1230 self.parse(text)
1231
1232
1233
1235 """Take the input file and create the structured information"""
1236
1237 status = 'first'
1238 for line in text.split('\n'):
1239 line = line.strip()
1240 if not line:
1241 continue
1242 elif line[0] == '#':
1243 self.comment += '%s\n' % line
1244 continue
1245 elif line.startswith('<event'):
1246 if '=' in line:
1247 found = re.findall(r"""(\w*)=(?:(?:['"])([^'"]*)(?=['"])|(\S*))""",line)
1248
1249
1250 self.eventflag = dict((n, a1) if a1 else (n,a2) for n,a1,a2 in found)
1251
1252 continue
1253
1254 elif 'first' == status:
1255 if '<rwgt>' in line:
1256 status = 'tag'
1257 else:
1258 self.assign_scale_line(line)
1259 status = 'part'
1260 continue
1261 if '<' in line:
1262 status = 'tag'
1263
1264 if 'part' == status:
1265 part = Particle(line, event=self)
1266 if part.E != 0:
1267 self.append(part)
1268 elif self.nexternal:
1269 self.nexternal-=1
1270 else:
1271 if '</event>' in line:
1272 line = line.replace('</event>','',1)
1273 self.tag += '%s\n' % line
1274
1275 self.assign_mother()
1276
1278 """convert the number in actual particle"""
1279
1280
1281 if all(p.status != -1 for p in self):
1282 if self.warning_order:
1283 Event.warning_order = False
1284 logger.warning("Weird format for lhe format: no incoming particle... adding a fake one")
1285
1286 mother = Particle(event=self)
1287 mother.status = -1
1288 mother.pid = 0
1289 self.insert(0,mother)
1290 mother.color2 = 0
1291 mother.event_id = 0
1292 self.nexternal += 1
1293 for p in self[1:]:
1294 p.mother1 = 1
1295 p.mother2 = 1
1296 p.event_id += 1
1297
1298
1299
1300 for i,particle in enumerate(self):
1301 if i < particle.mother1 or i < particle.mother2:
1302 if self.warning_order:
1303 logger.warning("Order of particle in the event did not agree with parent/child order. This might be problematic for some code.")
1304 Event.warning_order = False
1305 self.reorder_mother_child()
1306 return self.assign_mother()
1307
1308 if particle.mother1:
1309 try:
1310 particle.mother1 = self[int(particle.mother1) -1]
1311 except Exception:
1312 logger.warning("WRONG MOTHER INFO %s", self)
1313 particle.mother1 = 0
1314 if particle.mother2:
1315 try:
1316 particle.mother2 = self[int(particle.mother2) -1]
1317 except Exception:
1318 logger.warning("WRONG MOTHER INFO %s", self)
1319 particle.mother2 = 0
1320
1322 """change all the weights by a given ratio"""
1323
1324 self.wgt *= ratio
1325 self.parse_reweight()
1326 for key in self.reweight_data:
1327 self.reweight_data[key] *= ratio
1328 return self.wgt
1329
1331 """check and correct the mother/child position.
1332 only correct one order by call (but this is a recursive call)"""
1333
1334 tomove, position = None, None
1335 for i,particle in enumerate(self):
1336 if i < particle.mother1:
1337
1338 tomove, position = i, particle.mother1-1
1339 break
1340 if i < particle.mother2:
1341 tomove, position = i, particle.mother2-1
1342
1343
1344 if not tomove:
1345 return
1346
1347
1348 particle = self.pop(tomove)
1349 self.insert(int(position), particle)
1350
1351
1352 for i, particle in enumerate(self):
1353 particle.event_id = i
1354
1355 m1, m2 = particle.mother1, particle.mother2
1356 if m1 == tomove +1:
1357 particle.mother1 = position+1
1358 elif tomove < m1 <= position +1:
1359 particle.mother1 -= 1
1360 if m2 == tomove +1:
1361 particle.mother2 = position+1
1362 elif tomove < m2 <= position +1:
1363 particle.mother2 -= 1
1364
1365 return self.reorder_mother_child()
1366
1367
1368
1369
1370
1371
1373 """Parse the re-weight information in order to return a dictionary
1374 {key: value}. If no group is define group should be '' """
1375 if self.reweight_data:
1376 return self.reweight_data
1377 self.reweight_data = {}
1378 self.reweight_order = []
1379 start, stop = self.tag.find('<rwgt>'), self.tag.find('</rwgt>')
1380 if start != -1 != stop :
1381 pattern = re.compile(r'''<\s*wgt id=(?:\'|\")(?P<id>[^\'\"]+)(?:\'|\")\s*>\s*(?P<val>[\ded+-.]*)\s*</wgt>''',re.I)
1382 data = pattern.findall(self.tag[start:stop])
1383 try:
1384 self.reweight_data = dict([(pid, float(value)) for (pid, value) in data
1385 if not self.reweight_order.append(pid)])
1386
1387 except ValueError, error:
1388 raise Exception, 'Event File has unvalid weight. %s' % error
1389 self.tag = self.tag[:start] + self.tag[stop+7:]
1390 return self.reweight_data
1391
1393 """ """
1394 if hasattr(self, 'nloweight'):
1395 return self.nloweight
1396
1397 start, stop = self.tag.find('<mgrwgt>'), self.tag.find('</mgrwgt>')
1398 if start != -1 != stop :
1399
1400 text = self.tag[start+8:stop]
1401 self.nloweight = NLO_PARTIALWEIGHT(text, self, real_type=real_type,
1402 threshold=threshold)
1403 return self.nloweight
1404
1406 """get the string associate to the weight"""
1407
1408 text="""<mgrwgt>
1409 %(total_wgt).10e %(nb_wgt)i %(nb_event)i 0
1410 %(event)s
1411 %(wgt)s
1412 </mgrwgt>"""
1413
1414
1415 if not wgt:
1416 if not hasattr(self, 'nloweight'):
1417 return
1418 wgt = self.nloweight
1419
1420 data = {'total_wgt': wgt.total_wgt,
1421 'nb_wgt': wgt.nb_wgt,
1422 'nb_event': wgt.nb_event,
1423 'event': '\n'.join(p.__str__(mode='fortran') for p in wgt.momenta),
1424 'wgt':'\n'.join(w.__str__(mode='formatted')
1425 for e in wgt.cevents for w in e.wgts)}
1426
1427 data['total_wgt'] = sum([w.ref_wgt for e in wgt.cevents for w in e.wgts])
1428 start, stop = self.tag.find('<mgrwgt>'), self.tag.find('</mgrwgt>')
1429
1430 self.tag = self.tag[:start] + text % data + self.tag[stop+9:]
1431
1432
1434 """ """
1435
1436
1437 if hasattr(self, 'loweight'):
1438 return self.loweight
1439
1440 if not hasattr(Event, 'loweight_pattern'):
1441 Event.loweight_pattern = re.compile('''<rscale>\s*(?P<nqcd>\d+)\s+(?P<ren_scale>[\d.e+-]+)\s*</rscale>\s*\n\s*
1442 <asrwt>\s*(?P<asrwt>[\s\d.+-e]+)\s*</asrwt>\s*\n\s*
1443 <pdfrwt\s+beam=["']?1["']?\>\s*(?P<beam1>[\s\d.e+-]*)\s*</pdfrwt>\s*\n\s*
1444 <pdfrwt\s+beam=["']?2["']?\>\s*(?P<beam2>[\s\d.e+-]*)\s*</pdfrwt>\s*\n\s*
1445 <totfact>\s*(?P<totfact>[\d.e+-]*)\s*</totfact>
1446 ''',re.X+re.I+re.M)
1447
1448 start, stop = self.tag.find('<mgrwt>'), self.tag.find('</mgrwt>')
1449
1450 if start != -1 != stop :
1451 text = self.tag[start+8:stop]
1452
1453 info = Event.loweight_pattern.search(text)
1454 if not info:
1455 raise Exception, '%s not parsed'% text
1456 self.loweight={}
1457 self.loweight['n_qcd'] = int(info.group('nqcd'))
1458 self.loweight['ren_scale'] = float(info.group('ren_scale'))
1459 self.loweight['asrwt'] =[float(x) for x in info.group('asrwt').split()[1:]]
1460 self.loweight['tot_fact'] = float(info.group('totfact'))
1461
1462 args = info.group('beam1').split()
1463 npdf = int(args[0])
1464 self.loweight['n_pdfrw1'] = npdf
1465 self.loweight['pdf_pdg_code1'] = [int(i) for i in args[1:1+npdf]]
1466 self.loweight['pdf_x1'] = [float(i) for i in args[1+npdf:1+2*npdf]]
1467 self.loweight['pdf_q1'] = [float(i) for i in args[1+2*npdf:1+3*npdf]]
1468 args = info.group('beam2').split()
1469 npdf = int(args[0])
1470 self.loweight['n_pdfrw2'] = npdf
1471 self.loweight['pdf_pdg_code2'] = [int(i) for i in args[1:1+npdf]]
1472 self.loweight['pdf_x2'] = [float(i) for i in args[1+npdf:1+2*npdf]]
1473 self.loweight['pdf_q2'] = [float(i) for i in args[1+2*npdf:1+3*npdf]]
1474
1475 else:
1476 return None
1477 return self.loweight
1478
1479
1481 """Parse the line containing the starting scale for the shower"""
1482
1483 if self.matched_scale_data is not None:
1484 return self.matched_scale_data
1485
1486 self.matched_scale_data = []
1487
1488
1489 pattern = re.compile("<scales\s|</scales>")
1490 data = re.split(pattern,self.tag)
1491 if len(data) == 1:
1492 return []
1493 else:
1494 tmp = {}
1495 start,content, end = data
1496 self.tag = "%s%s" % (start, end)
1497 pattern = re.compile("pt_clust_(\d*)=\"([\de+-.]*)\"")
1498 for id,value in pattern.findall(content):
1499 tmp[int(id)] = float(value)
1500 for i in range(1, len(self)+1):
1501 if i in tmp:
1502 self.matched_scale_data.append(tmp[i])
1503 else:
1504 self.matched_scale_data.append(-1)
1505 return self.matched_scale_data
1506
1508 """ parse the flag for syscalc between <mgrwt></mgrwt>
1509 <mgrwt>
1510 <rscale> 3 0.26552898E+03</rscale>
1511 <asrwt>0</asrwt>
1512 <pdfrwt beam="1"> 1 21 0.14527945E+00 0.26552898E+03</pdfrwt>
1513 <pdfrwt beam="2"> 1 21 0.15249110E-01 0.26552898E+03</pdfrwt>
1514 <totfact> 0.10344054E+04</totfact>
1515 </mgrwt>
1516 """
1517 if self.syscalc_data:
1518 return self.syscalc_data
1519
1520 pattern = re.compile("<mgrwt>|</mgrwt>")
1521 pattern2 = re.compile("<(?P<tag>[\w]*)(?:\s*(\w*)=[\"'](.*)[\"']\s*|\s*)>(.*)</(?P=tag)>")
1522 data = re.split(pattern,self.tag)
1523 if len(data) == 1:
1524 return []
1525 else:
1526 tmp = {}
1527 start,content, end = data
1528 self.tag = "%s%s" % (start, end)
1529 for tag, key, keyval, tagval in pattern2.findall(content):
1530 if key:
1531 self.syscalc_data[(tag, key, keyval)] = tagval
1532 else:
1533 self.syscalc_data[tag] = tagval
1534 return self.syscalc_data
1535
1536
1537 - def add_decay_to_particle(self, position, decay_event):
1538 """define the decay of the particle id by the event pass in argument"""
1539
1540 this_particle = self[position]
1541
1542 this_particle.status = 2
1543 this_particle.helicity = 0
1544
1545
1546 decay_particle = decay_event[0]
1547 this_4mom = FourMomentum(this_particle)
1548 nb_part = len(self)
1549
1550 thres = decay_particle.E*1e-10
1551 assert max(decay_particle.px, decay_particle.py, decay_particle.pz) < thres,\
1552 "not on rest particle %s %s %s %s" % (decay_particle.E, decay_particle.px,decay_particle.py,decay_particle.pz)
1553
1554 self.nexternal += decay_event.nexternal -1
1555 old_scales = list(self.parse_matching_scale())
1556 if old_scales:
1557 jet_position = sum(1 for i in range(position) if self[i].status==1)
1558 initial_pos = sum(1 for i in range(position) if self[i].status==-1)
1559 self.matched_scale_data.pop(initial_pos+jet_position)
1560
1561
1562 for particle in decay_event[1:]:
1563
1564 new_particle = Particle(particle, self)
1565 new_particle.event_id = len(self)
1566 self.append(new_particle)
1567 if old_scales:
1568 self.matched_scale_data.append(old_scales[initial_pos+jet_position])
1569
1570 new_momentum = this_4mom.boost(FourMomentum(new_particle))
1571 new_particle.set_momentum(new_momentum)
1572
1573 for tag in ['mother1', 'mother2']:
1574 mother = getattr(particle, tag)
1575 if isinstance(mother, Particle):
1576 mother_id = getattr(particle, tag).event_id
1577 if mother_id == 0:
1578 setattr(new_particle, tag, this_particle)
1579 else:
1580 try:
1581 setattr(new_particle, tag, self[nb_part + mother_id -1])
1582 except Exception, error:
1583 print error
1584 misc.sprint( self)
1585 misc.sprint(nb_part + mother_id -1)
1586 misc.sprint(tag)
1587 misc.sprint(position, decay_event)
1588 misc.sprint(particle)
1589 misc.sprint(len(self), nb_part + mother_id -1)
1590 raise
1591 elif tag == "mother2" and isinstance(particle.mother1, Particle):
1592 new_particle.mother2 = this_particle
1593 else:
1594 raise Exception, "Something weird happens. Please report it for investigation"
1595
1596
1597 max_color=501
1598 for particle in self[:nb_part]:
1599 max_color=max(max_color, particle.color1, particle.color2)
1600
1601
1602 color_mapping = {}
1603 color_mapping[decay_particle.color1] = this_particle.color1
1604 color_mapping[decay_particle.color2] = this_particle.color2
1605 for particle in self[nb_part:]:
1606 if particle.color1:
1607 if particle.color1 not in color_mapping:
1608 max_color +=1
1609 color_mapping[particle.color1] = max_color
1610 particle.color1 = max_color
1611 else:
1612 particle.color1 = color_mapping[particle.color1]
1613 if particle.color2:
1614 if particle.color2 not in color_mapping:
1615 max_color +=1
1616 color_mapping[particle.color2] = max_color
1617 particle.color2 = max_color
1618 else:
1619 particle.color2 = color_mapping[particle.color2]
1620
1622 """use auto-recursion"""
1623
1624 pdg_to_decay = dict(pdg_to_decay)
1625
1626 for i,particle in enumerate(self):
1627 if particle.status != 1:
1628 continue
1629 if particle.pdg in pdg_to_decay and pdg_to_decay[particle.pdg]:
1630 one_decay = pdg_to_decay[particle.pdg].pop()
1631 self.add_decay_to_particle(i, one_decay)
1632 return self.add_decays(pdg_to_decay)
1633 return self
1634
1635
1636
1638
1639 to_remove = []
1640 if event_id is not None:
1641 to_remove.append(self[event_id])
1642
1643 if pdg_code:
1644 for particle in self:
1645 if particle.pid == pdg_code:
1646 to_remove.append(particle)
1647
1648 new_event = Event()
1649
1650 for tag in ['nexternal', 'ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']:
1651 setattr(new_event, tag, getattr(self, tag))
1652
1653 for particle in self:
1654 if isinstance(particle.mother1, Particle) and particle.mother1 in to_remove:
1655 to_remove.append(particle)
1656 if particle.status == 1:
1657 new_event.nexternal -= 1
1658 continue
1659 elif isinstance(particle.mother2, Particle) and particle.mother2 in to_remove:
1660 to_remove.append(particle)
1661 if particle.status == 1:
1662 new_event.nexternal -= 1
1663 continue
1664 else:
1665 new_event.append(Particle(particle))
1666
1667
1668
1669 for pos, particle in enumerate(new_event):
1670 particle.event_id = pos
1671 if particle in to_remove:
1672 particle.status = 1
1673 return new_event
1674
1675 - def get_decay(self, pdg_code=0, event_id=None):
1676
1677 to_start = []
1678 if event_id is not None:
1679 to_start.append(self[event_id])
1680
1681 elif pdg_code:
1682 for particle in self:
1683 if particle.pid == pdg_code:
1684 to_start.append(particle)
1685 break
1686
1687 new_event = Event()
1688
1689 for tag in ['ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']:
1690 setattr(new_event, tag, getattr(self, tag))
1691
1692
1693 old2new = {}
1694 new_decay_part = Particle(to_start[0])
1695 new_decay_part.mother1 = None
1696 new_decay_part.mother2 = None
1697 new_decay_part.status = -1
1698 old2new[new_decay_part.event_id] = len(old2new)
1699 new_event.append(new_decay_part)
1700
1701
1702
1703 for particle in self:
1704 if isinstance(particle.mother1, Particle) and particle.mother1.event_id in old2new\
1705 or isinstance(particle.mother2, Particle) and particle.mother2.event_id in old2new:
1706 old2new[particle.event_id] = len(old2new)
1707 new_event.append(Particle(particle))
1708
1709
1710
1711 nexternal = 0
1712 for pos, particle in enumerate(new_event):
1713 particle.event_id = pos
1714 if particle.mother1:
1715 particle.mother1 = new_event[old2new[particle.mother1.event_id]]
1716 if particle.mother2:
1717 particle.mother2 = new_event[old2new[particle.mother2.event_id]]
1718 if particle.status in [-1,1]:
1719 nexternal +=1
1720 new_event.nexternal = nexternal
1721
1722 return new_event
1723
1724
1726 """check various property of the events"""
1727
1728
1729 E, px, py, pz = 0,0,0,0
1730 absE, abspx, abspy, abspz = 0,0,0,0
1731 for particle in self:
1732 coeff = 1
1733 if particle.status == -1:
1734 coeff = -1
1735 elif particle.status != 1:
1736 continue
1737 E += coeff * particle.E
1738 absE += abs(particle.E)
1739 px += coeff * particle.px
1740 py += coeff * particle.py
1741 pz += coeff * particle.pz
1742 abspx += abs(particle.px)
1743 abspy += abs(particle.py)
1744 abspz += abs(particle.pz)
1745
1746 threshold = 5e-7
1747 if E/absE > threshold:
1748 logger.critical(self)
1749 raise Exception, "Do not conserve Energy %s, %s" % (E/absE, E)
1750 if px/abspx > threshold:
1751 logger.critical(self)
1752 raise Exception, "Do not conserve Px %s, %s" % (px/abspx, px)
1753 if py/abspy > threshold:
1754 logger.critical(self)
1755 raise Exception, "Do not conserve Py %s, %s" % (py/abspy, py)
1756 if pz/abspz > threshold:
1757 logger.critical(self)
1758 raise Exception, "Do not conserve Pz %s, %s" % (pz/abspz, pz)
1759
1760
1761 self.check_color_structure()
1762
1764 """read the line corresponding to global event line
1765 format of the line is:
1766 Nexternal IEVENT WEIGHT SCALE AEW AS
1767 """
1768 inputs = line.split()
1769 assert len(inputs) == 6
1770 self.nexternal=int(inputs[0])
1771 self.ievent=int(inputs[1])
1772 self.wgt=float(inputs[2])
1773 self.scale=float(inputs[3])
1774 self.aqed=float(inputs[4])
1775 self.aqcd=float(inputs[5])
1776
1778 """Return the unique tag identifying the SubProcesses for the generation.
1779 Usefull for program like MadSpin and Reweight module."""
1780
1781 initial, final, order = [], [], [[], []]
1782 for particle in self:
1783 if particle.status == -1:
1784 initial.append(particle.pid)
1785 order[0].append(particle.pid)
1786 elif particle.status == 1:
1787 final.append(particle.pid)
1788 order[1].append(particle.pid)
1789 initial.sort(), final.sort()
1790 tag = (tuple(initial), tuple(final))
1791 return tag, order
1792
1793 @staticmethod
1794 - def mass_shuffle(momenta, sqrts, new_mass, new_sqrts=None):
1795 """use the RAMBO method to shuffle the PS. initial sqrts is preserved."""
1796
1797 if not new_sqrts:
1798 new_sqrts = sqrts
1799
1800 oldm = [p.mass_sqr for p in momenta]
1801 newm = [m**2 for m in new_mass]
1802 tot_mom = sum(momenta, FourMomentum())
1803 if tot_mom.pt2 > 1e-5:
1804 boost_back = FourMomentum(tot_mom.mass,0,0,0).boost_to_restframe(tot_mom)
1805 for i,m in enumerate(momenta):
1806 momenta[i] = m.boost_to_restframe(tot_mom)
1807
1808
1809 f = lambda chi: new_sqrts - sum(math.sqrt(max(0, M + chi**2*(p.E**2-m)))
1810 for M,p,m in zip(newm, momenta,oldm))
1811
1812 df = lambda chi: -1* sum(chi*(p.E**2-m)/math.sqrt(max(0,(p.E**2-m)*chi**2+M))
1813 for M,p,m in zip(newm, momenta,oldm))
1814
1815 if sum(new_mass) > new_sqrts:
1816 return momenta, 0
1817 try:
1818 chi = misc.newtonmethod(f, df, 1.0, error=1e-7,maxiter=1000)
1819 except:
1820 return momenta, 0
1821
1822 new_momenta = []
1823 for i,p in enumerate(momenta):
1824 new_momenta.append(
1825 FourMomentum(math.sqrt(newm[i]+chi**2*(p.E**2-oldm[i])),
1826 chi*p.px, chi*p.py, chi*p.pz))
1827
1828
1829
1830
1831
1832
1833
1834 jac = chi**(3*len(momenta)-3)
1835 jac *= reduce(operator.mul,[p.E/k.E for p,k in zip(momenta, new_momenta)],1)
1836 jac *= sum(p.norm_sq/p.E for p in momenta)
1837 jac /= sum(k.norm_sq/k.E for k in new_momenta)
1838
1839
1840 if tot_mom.pt2 > 1e-5:
1841 for i,m in enumerate(new_momenta):
1842 new_momenta[i] = m.boost_to_restframe(boost_back)
1843 return new_momenta, jac
1844
1845
1846
1847
1849 """routine to rescale the mass via RAMBO method. no internal mass preserve.
1850 sqrts is preserve (RAMBO algo)
1851 """
1852
1853 old_momenta = []
1854 new_masses = []
1855 change_mass = False
1856 for part in self:
1857 if part.status == 1:
1858 old_momenta.append(FourMomentum(part))
1859 new_masses.append(new_param_card.get_value('mass', abs(part.pid)))
1860 if not misc.equal(part.mass, new_masses[-1], 4, zero_limit=10):
1861 change_mass = True
1862
1863 if not change_mass:
1864 return 1
1865
1866 sqrts = self.sqrts
1867
1868
1869 new_mom, jac = self.mass_shuffle(old_momenta, sqrts, new_masses)
1870
1871
1872 ind =0
1873 for part in self:
1874 if part.status==1:
1875 part.E, part.px, part.py, part.pz, part.mass = \
1876 new_mom[ind].E, new_mom[ind].px, new_mom[ind].py, new_mom[ind].pz,new_mom[ind].mass
1877 ind+=1
1878 return jac
1879
1881 """routine to rescale the momenta to change the invariant mass"""
1882
1883 old_momenta = []
1884 incoming = []
1885 masses = []
1886 for part in self:
1887 if part.status == -1:
1888 incoming.append(FourMomentum(part))
1889 if part.status == 1:
1890 old_momenta.append(FourMomentum(part))
1891 masses.append(part.mass)
1892
1893 p_init = FourMomentum()
1894 p_inits = []
1895 n_init = 0
1896 for p in incoming:
1897 n_init +=1
1898 p_init += p
1899 p_inits.append(p)
1900 old_sqrts = p_init.mass
1901
1902 new_mom, jac = self.mass_shuffle(old_momenta, old_sqrts, masses, new_sqrts=new_sqrts)
1903
1904
1905 ind =0
1906 for part in self:
1907 if part.status==1:
1908 part.E, part.px, part.py, part.pz, part.mass = \
1909 new_mom[ind].E, new_mom[ind].px, new_mom[ind].py, new_mom[ind].pz,new_mom[ind].mass
1910 ind+=1
1911
1912
1913 p_init = FourMomentum()
1914 for part in self:
1915 if part.status==1:
1916 p_init += part
1917 if n_init == 1:
1918 for part in self:
1919 if part.status == -1:
1920 part.E, part.px, part.py, part.pz = \
1921 p_init.E, p_init.px, p_init.py, p_init.pz
1922 elif n_init ==2:
1923 if not misc.equal(p_init.px, 0) or not misc.equal(p_init.py, 0):
1924 raise Exception
1925 if not misc.equal(p_inits[0].px, 0) or not misc.equal(p_inits[0].py, 0):
1926 raise Exception
1927
1928
1929
1930
1931 eta = math.log(2*p_inits[0].E/old_sqrts)
1932 new_p = [[new_sqrts/2*math.exp(eta), 0., 0., new_sqrts/2*math.exp(eta)],
1933 [new_sqrts/2*math.exp(-eta), 0., 0., -new_sqrts/2*math.exp(-eta)]]
1934
1935 ind=0
1936 for part in self:
1937 if part.status == -1:
1938 part.E, part.px, part.py, part.pz = new_p[ind]
1939 ind+=1
1940 if ind ==2:
1941 break
1942 else:
1943 raise Exception
1944
1945 return jac
1946
1947
1949 """return a list with the helicities in the order asked for"""
1950
1951
1952 order = [list(get_order[0]), list(get_order[1])]
1953 out = [9] *(len(order[0])+len(order[1]))
1954 for i, part in enumerate(self):
1955 if part.status == 1:
1956 try:
1957 ind = order[1].index(part.pid)
1958 except ValueError, error:
1959 if not allow_reversed:
1960 raise error
1961 else:
1962 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
1963 try:
1964 return self.get_helicity(order, False)
1965 except ValueError:
1966 raise error
1967 position = len(order[0]) + ind
1968 order[1][ind] = 0
1969 elif part.status == -1:
1970 try:
1971 ind = order[0].index(part.pid)
1972 except ValueError, error:
1973 if not allow_reversed:
1974 raise error
1975 else:
1976 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
1977 try:
1978 return self.get_helicity(order, False)
1979 except ValueError:
1980 raise error
1981
1982 position = ind
1983 order[0][ind] = 0
1984 else:
1985 continue
1986 out[position] = int(part.helicity)
1987 return out
1988
1989
1991 """check the validity of the color structure"""
1992
1993
1994 color_index = collections.defaultdict(int)
1995 for particle in self:
1996 if particle.status in [-1,1]:
1997 if particle.color1:
1998 color_index[particle.color1] +=1
1999 if -7 < particle.pdg < 0:
2000 raise Exception, "anti-quark with color tag"
2001 if particle.color2:
2002 color_index[particle.color2] +=1
2003 if 7 > particle.pdg > 0:
2004 raise Exception, "quark with anti-color tag"
2005
2006
2007 for key,value in color_index.items():
2008 if value > 2:
2009 print self
2010 print key, value
2011 raise Exception, 'Wrong color_flow'
2012
2013
2014
2015 check = []
2016 popup_index = []
2017 for particle in self:
2018 mothers = []
2019 childs = []
2020 if particle.mother1:
2021 mothers.append(particle.mother1)
2022 if particle.mother2 and particle.mother2 is not particle.mother1:
2023 mothers.append(particle.mother2)
2024 if not mothers:
2025 continue
2026 if (particle.mother1.event_id, particle.mother2.event_id) in check:
2027 continue
2028 check.append((particle.mother1.event_id, particle.mother2.event_id))
2029
2030 childs = [p for p in self if p.mother1 is particle.mother1 and \
2031 p.mother2 is particle.mother2]
2032
2033 mcolors = []
2034 manticolors = []
2035 for m in mothers:
2036 if m.color1:
2037 if m.color1 in manticolors:
2038 manticolors.remove(m.color1)
2039 else:
2040 mcolors.append(m.color1)
2041 if m.color2:
2042 if m.color2 in mcolors:
2043 mcolors.remove(m.color2)
2044 else:
2045 manticolors.append(m.color2)
2046 ccolors = []
2047 canticolors = []
2048 for m in childs:
2049 if m.color1:
2050 if m.color1 in canticolors:
2051 canticolors.remove(m.color1)
2052 else:
2053 ccolors.append(m.color1)
2054 if m.color2:
2055 if m.color2 in ccolors:
2056 ccolors.remove(m.color2)
2057 else:
2058 canticolors.append(m.color2)
2059 for index in mcolors[:]:
2060 if index in ccolors:
2061 mcolors.remove(index)
2062 ccolors.remove(index)
2063 for index in manticolors[:]:
2064 if index in canticolors:
2065 manticolors.remove(index)
2066 canticolors.remove(index)
2067
2068 if mcolors != []:
2069
2070 if len(canticolors) + len(mcolors) != 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 += canticolors
2075 elif manticolors != []:
2076
2077 if len(ccolors) + len(manticolors) != 3:
2078 logger.critical(str(self))
2079 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs])
2080 else:
2081 popup_index += ccolors
2082
2083
2084 if len(popup_index) != len(set(popup_index)):
2085 logger.critical(self)
2086 raise Exception, "Wrong color flow: identical poping-up index, %s" % (popup_index)
2087
2089 """two event are the same if they have the same momentum. other info are ignored"""
2090
2091 if other is None:
2092 return False
2093
2094 for i,p in enumerate(self):
2095 if p.E != other[i].E:
2096 return False
2097 elif p.pz != other[i].pz:
2098 return False
2099 elif p.px != other[i].px:
2100 return False
2101 elif p.py != other[i].py:
2102 return False
2103 return True
2104
2105
2107 """return a correctly formatted LHE event"""
2108
2109 out="""<event%(event_flag)s>
2110 %(scale)s
2111 %(particles)s
2112 %(comments)s
2113 %(tag)s
2114 %(reweight)s
2115 </event>
2116 """
2117 if event_id not in ['', None]:
2118 self.eventflag['event'] = str(event_id)
2119
2120 if self.eventflag:
2121 event_flag = ' %s' % ' '.join('%s="%s"' % (k,v) for (k,v) in self.eventflag.items())
2122 else:
2123 event_flag = ''
2124
2125 scale_str = "%2d %6d %+13.7e %14.8e %14.8e %14.8e" % \
2126 (self.nexternal,self.ievent,self.wgt,self.scale,self.aqed,self.aqcd)
2127
2128
2129 if self.reweight_data:
2130
2131 if set(self.reweight_data.keys()) != set(self.reweight_order):
2132 self.reweight_order += [k for k in self.reweight_data.keys() \
2133 if k not in self.reweight_order]
2134
2135 reweight_str = '<rwgt>\n%s\n</rwgt>' % '\n'.join(
2136 '<wgt id=\'%s\'> %+13.7e </wgt>' % (i, float(self.reweight_data[i]))
2137 for i in self.reweight_order if i in self.reweight_data)
2138 else:
2139 reweight_str = ''
2140
2141 tag_str = self.tag
2142 if hasattr(self, 'nloweight') and self.nloweight.modified:
2143 self.rewrite_nlo_weight()
2144 tag_str = self.tag
2145
2146 if self.matched_scale_data:
2147 tmp_scale = ' '.join(['pt_clust_%i=\"%s\"' % (i+1,v)
2148 for i,v in enumerate(self.matched_scale_data)
2149 if v!=-1])
2150 if tmp_scale:
2151 tag_str = "<scales %s></scales>%s" % (tmp_scale, self.tag)
2152
2153 if self.syscalc_data:
2154 keys= ['rscale', 'asrwt', ('pdfrwt', 'beam', '1'), ('pdfrwt', 'beam', '2'),
2155 'matchscale', 'totfact']
2156 sys_str = "<mgrwt>\n"
2157 template = """<%(key)s%(opts)s>%(values)s</%(key)s>\n"""
2158 for k in keys:
2159 if k not in self.syscalc_data:
2160 continue
2161 replace = {}
2162 replace['values'] = self.syscalc_data[k]
2163 if isinstance(k, str):
2164 replace['key'] = k
2165 replace['opts'] = ''
2166 else:
2167 replace['key'] = k[0]
2168 replace['opts'] = ' %s=\"%s\"' % (k[1],k[2])
2169 sys_str += template % replace
2170 sys_str += "</mgrwt>\n"
2171 reweight_str = sys_str + reweight_str
2172
2173 out = out % {'event_flag': event_flag,
2174 'scale': scale_str,
2175 'particles': '\n'.join([str(p) for p in self]),
2176 'tag': tag_str,
2177 'comments': self.comment,
2178 'reweight': reweight_str}
2179
2180 return re.sub('[\n]+', '\n', out)
2181
2182 - def get_momenta(self, get_order, allow_reversed=True):
2183 """return the momenta vector in the order asked for"""
2184
2185
2186 order = [list(get_order[0]), list(get_order[1])]
2187 out = [''] *(len(order[0])+len(order[1]))
2188 for i, part in enumerate(self):
2189 if part.status == 1:
2190 try:
2191 ind = order[1].index(part.pid)
2192 except ValueError, error:
2193 if not allow_reversed:
2194 raise error
2195 else:
2196 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2197 try:
2198 return self.get_momenta_str(order, False)
2199 except ValueError:
2200 raise error
2201 position = len(order[0]) + ind
2202 order[1][ind] = 0
2203 elif part.status == -1:
2204 try:
2205 ind = order[0].index(part.pid)
2206 except ValueError, error:
2207 if not allow_reversed:
2208 raise error
2209 else:
2210 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2211 try:
2212 return self.get_momenta_str(order, False)
2213 except ValueError:
2214 raise error
2215
2216 position = ind
2217 order[0][ind] = 0
2218 else:
2219 continue
2220
2221 out[position] = (part.E, part.px, part.py, part.pz)
2222
2223 return out
2224
2225
2238
2239
2241
2242 scale = 0
2243 for particle in self:
2244 if particle.status != 1:
2245 continue
2246 p=FourMomentum(particle)
2247 scale += math.sqrt(p.mass_sqr + p.pt**2)
2248
2249 return prefactor * scale
2250
2251
2253
2254 scale = 0
2255 for particle in self:
2256 if particle.status != 1:
2257 continue
2258 p = FourMomentum(particle)
2259 pt = p.pt
2260 if (pt>0):
2261 scale += p.E*pt/math.sqrt(pt**2+p.pz**2)
2262
2263 return prefactor * scale
2264
2265 @property
2268
2270
2271 scale = 0
2272 init = []
2273 for particle in self:
2274 if particle.status == -1:
2275 init.append(FourMomentum(particle))
2276 if len(init) == 1:
2277 return init[0].mass
2278 elif len(init)==2:
2279 return math.sqrt((init[0]+init[1])**2)
2280
2281
2282
2283
2285 """return the momenta str in the order asked for"""
2286
2287 out = self.get_momenta(get_order, allow_reversed)
2288
2289 format = '%.12f'
2290 format_line = ' '.join([format]*4) + ' \n'
2291 out = [format_line % one for one in out]
2292 out = ''.join(out).replace('e','d')
2293 return out
2294
2296 """A class to allow to read both gzip and not gzip file.
2297 containing only weight from pythia --generated by SysCalc"""
2298
2299 - def __new__(self, path, mode='r', *args, **opt):
2300 if path.endswith(".gz"):
2301 try:
2302 return gzip.GzipFile.__new__(WeightFileGzip, path, mode, *args, **opt)
2303 except IOError, error:
2304 raise
2305 except Exception, error:
2306 if mode == 'r':
2307 misc.gunzip(path)
2308 return file.__new__(WeightFileNoGzip, path[:-3], mode, *args, **opt)
2309 else:
2310 return file.__new__(WeightFileNoGzip, path, mode, *args, **opt)
2311
2312
2313 - def __init__(self, path, mode='r', *args, **opt):
2314 """open file and read the banner [if in read mode]"""
2315
2316 super(EventFile, self).__init__(path, mode, *args, **opt)
2317 self.banner = ''
2318 if mode == 'r':
2319 line = ''
2320 while '</header>' not in line.lower():
2321 try:
2322 line = super(EventFile, self).next()
2323 except StopIteration:
2324 self.seek(0)
2325 self.banner = ''
2326 break
2327 if "<event" in line.lower():
2328 self.seek(0)
2329 self.banner = ''
2330 break
2331
2332 self.banner += line
2333
2337
2340
2343 """a convenient object for 4-momenta operation"""
2344
2345 - def __init__(self, obj=0, px=0, py=0, pz=0, E=0):
2346 """initialize the four momenta"""
2347
2348 if obj is 0 and E:
2349 obj = E
2350
2351 if isinstance(obj, (FourMomentum, Particle)):
2352 px = obj.px
2353 py = obj.py
2354 pz = obj.pz
2355 E = obj.E
2356 elif isinstance(obj, (list, tuple)):
2357 assert len(obj) ==4
2358 E = obj[0]
2359 px = obj[1]
2360 py = obj[2]
2361 pz = obj[3]
2362 elif isinstance(obj, str):
2363 obj = [float(i) for i in obj.split()]
2364 assert len(obj) ==4
2365 E = obj[0]
2366 px = obj[1]
2367 py = obj[2]
2368 pz = obj[3]
2369 else:
2370 E =obj
2371
2372
2373 self.E = float(E)
2374 self.px = float(px)
2375 self.py = float(py)
2376 self.pz = float(pz)
2377
2378 @property
2380 """return the mass"""
2381 return math.sqrt(max(self.E**2 - self.px**2 - self.py**2 - self.pz**2,0))
2382
2383 @property
2385 """return the mass square"""
2386 return self.E**2 - self.px**2 - self.py**2 - self.pz**2
2387
2388 @property
2390 return math.sqrt(max(0, self.pt2))
2391
2392 @property
2394 norm = math.sqrt(self.px**2 + self.py**2+self.pz**2)
2395 return 0.5* math.log((norm - self.pz) / (norm + self.pz))
2396
2397 @property
2399 return 0.5* math.log((self.E +self.pz) / (self.E - self.pz))
2400
2401
2402 @property
2404 """ return the pt square """
2405
2406 return self.px**2 + self.py**2
2407
2408 @property
2410 """ return |\vec p| """
2411 return math.sqrt(self.px**2 + self.py**2 + self.pz**2)
2412
2413 @property
2415 """ return |\vec p|^2 """
2416 return self.px**2 + self.py**2 + self.pz**2
2417
2418 @property
2420 """return the mass square"""
2421 import math
2422 return math.atan(math.sqrt((self.px**2+self.py**2)/self.pz**2))
2423
2424
2426
2427 assert isinstance(obj, FourMomentum)
2428 new = FourMomentum(self.E+obj.E,
2429 self.px + obj.px,
2430 self.py + obj.py,
2431 self.pz + obj.pz)
2432 return new
2433
2435 """update the object with the sum"""
2436 self.E += obj.E
2437 self.px += obj.px
2438 self.py += obj.py
2439 self.pz += obj.pz
2440 return self
2441
2443
2444 assert isinstance(obj, FourMomentum)
2445 new = FourMomentum(self.E-obj.E,
2446 self.px - obj.px,
2447 self.py - obj.py,
2448 self.pz - obj.pz)
2449 return new
2450
2452 """update the object with the sum"""
2453 self.E -= obj.E
2454 self.px -= obj.px
2455 self.py -= obj.py
2456 self.pz -= obj.pz
2457 return self
2458
2460 if isinstance(obj, FourMomentum):
2461 return self.E*obj.E - self.px *obj.px - self.py * obj.py - self.pz * obj.pz
2462 elif isinstance(obj, (float, int)):
2463 return FourMomentum(obj*self.E,obj*self.px,obj*self.py,obj*self.pz )
2464 else:
2465 raise NotImplemented
2466 __rmul__ = __mul__
2467
2469 assert power in [1,2]
2470
2471 if power == 1:
2472 return FourMomentum(self)
2473 elif power == 2:
2474 return self.mass_sqr
2475
2477 return 'FourMomentum(%s,%s,%s,%s)' % (self.E, self.px, self.py,self.pz)
2478
2479 - def __str__(self, mode='python'):
2480 if mode == 'python':
2481 return self.__repr__()
2482 elif mode == 'fortran':
2483 return '%.10e %.10e %.10e %.10e' % self.get_tuple()
2484
2486 return (self.E, self.px, self.py,self.pz)
2487
2489 """mom 4-momenta is suppose to be given in the rest frame of this 4-momenta.
2490 the output is the 4-momenta in the frame of this 4-momenta
2491 function copied from HELAS routine."""
2492
2493
2494 pt = self.px**2 + self.py**2 + self.pz**2
2495 if pt:
2496 s3product = self.px * mom.px + self.py * mom.py + self.pz * mom.pz
2497 mass = self.mass
2498 lf = (mom.E + (self.E - mass) * s3product / pt ) / mass
2499 return FourMomentum(E=(self.E*mom.E+s3product)/mass,
2500 px=mom.px + self.px * lf,
2501 py=mom.py + self.py * lf,
2502 pz=mom.pz + self.pz * lf)
2503 else:
2504 return FourMomentum(mom)
2505
2506 - def zboost(self, pboost=None, E=0, pz=0):
2507 """Both momenta should be in the same frame.
2508 The boost perform correspond to the boost required to set pboost at
2509 rest (only z boost applied).
2510 """
2511 if isinstance(pboost, FourMomentum):
2512 E = pboost.E
2513 pz = pboost.pz
2514
2515
2516 gamma = E / math.sqrt(E**2-pz**2)
2517 gammabeta = pz / math.sqrt(E**2-pz**2)
2518
2519 out = FourMomentum([gamma*self.E - gammabeta*self.pz,
2520 self.px,
2521 self.py,
2522 gamma*self.pz - gammabeta*self.E])
2523
2524 if abs(out.pz) < 1e-6 * out.E:
2525 out.pz = 0
2526 return out
2527
2529 """apply the boost transformation such that pboost is at rest in the new frame.
2530 First apply a rotation to allign the pboost to the z axis and then use
2531 zboost routine (see above)
2532 """
2533
2534 if pboost.px == 0 == pboost.py:
2535 out = self.zboost(E=pboost.E,pz=pboost.pz)
2536 return out
2537
2538
2539
2540
2541
2542
2543
2544 p = math.sqrt( pboost.px**2 + pboost.py**2+ pboost.pz**2)
2545 cosF = pboost.pz / p
2546 sinF = math.sqrt(1-cosF**2)
2547 sinT = pboost.py/p/sinF
2548 cosT = pboost.px/p/sinF
2549
2550 out=FourMomentum([self.E,
2551 self.px*cosT*cosF + self.py*sinT*cosF-self.pz*sinF,
2552 -self.px*sinT+ self.py*cosT,
2553 self.px*cosT*sinF + self.py*sinT*sinF + self.pz*cosF
2554 ])
2555 out = out.zboost(E=pboost.E,pz=p)
2556 return out
2557
2562
2563 - def __init__(self, input, real_type=(1,11)):
2564 """ """
2565
2566 self.real_type = real_type
2567 if isinstance(input, str):
2568 self.parse(input)
2569
2570 - def __str__(self, mode='display'):
2571
2572 if mode == 'display':
2573 out = """ pwgt: %(pwgt)s
2574 born, real : %(born)s %(real)s
2575 pdgs : %(pdgs)s
2576 bjks : %(bjks)s
2577 scales**2, gs: %(scales2)s %(gs)s
2578 born/real related : %(born_related)s %(real_related)s
2579 type / nfks : %(type)s %(nfks)s
2580 to merge : %(to_merge_pdg)s in %(merge_new_pdg)s
2581 ref_wgt : %(ref_wgt)s""" % self.__dict__
2582 return out
2583 elif mode == 'formatted':
2584 format_var = []
2585 variable = []
2586
2587 def to_add_full(f, v, format_var, variable):
2588 """ function to add to the formatted output"""
2589 if isinstance(v, list):
2590 format_var += [f]*len(v)
2591 variable += v
2592 else:
2593 format_var.append(f)
2594 variable.append(v)
2595 to_add = lambda x,y: to_add_full(x,y, format_var, variable)
2596
2597 to_add('%.10e', [p*self.bias_wgt for p in self.pwgt])
2598 to_add('%.10e', self.born)
2599 to_add('%.10e', self.real)
2600 to_add('%i', self.nexternal)
2601 to_add('%i', self.pdgs)
2602 to_add('%i', self.qcdpower)
2603 to_add('%.10e', self.bjks)
2604 to_add('%.10e', self.scales2)
2605 to_add('%.10e', self.gs)
2606 to_add('%i', [self.born_related, self.real_related])
2607 to_add('%i' , [self.type, self.nfks])
2608 to_add('%i' , self.to_merge_pdg)
2609 to_add('%i', self.merge_new_pdg)
2610 to_add('%.10e', self.ref_wgt*self.bias_wgt)
2611 to_add('%.10e', self.bias_wgt)
2612 return ' '.join(format_var) % tuple(variable)
2613
2614
2615 - def parse(self, text, keep_bias=False):
2616 """parse the line and create the related object.
2617 keep bias allow to not systematically correct for the bias in the written information"""
2618
2619
2620
2621 data = text.split()
2622
2623
2624
2625
2626 self.pwgt = [float(f) for f in data[:3]]
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645 self.born = float(data[3])
2646 self.real = float(data[4])
2647
2648
2649 self.nexternal = int(data[5])
2650
2651
2652 self.pdgs = [int(i) for i in data[6:6+self.nexternal]]
2653 flag = 6+self.nexternal
2654
2655
2656 self.qcdpower = int(data[flag])
2657
2658
2659 self.bjks = [float(f) for f in data[flag+1:flag+3]]
2660
2661
2662 self.scales2 = [float(f) for f in data[flag+3:flag+6]]
2663
2664
2665 self.gs = float(data[flag+6])
2666
2667
2668
2669
2670
2671
2672
2673 self.born_related = int(data[flag+7])
2674 self.real_related = int(data[flag+8])
2675
2676
2677
2678
2679
2680
2681
2682
2683
2684
2685
2686
2687
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
2698 self.type = int(data[flag+9])
2699
2700
2701
2702
2703 self.nfks = int(data[flag+10])
2704
2705
2706
2707
2708
2709
2710
2711 self.to_merge_pdg = [int (f) for f in data[flag+11:flag+13]]
2712
2713
2714 self.merge_new_pdg = int(data[flag+13])
2715
2716
2717
2718
2719
2720 self.ref_wgt = float(data[flag+14])
2721
2722
2723
2724 try:
2725 self.bias_wgt = float(data[flag+15])
2726 except IndexError:
2727 self.bias_wgt = 1.0
2728
2729 if not keep_bias:
2730 self.ref_wgt /= self.bias_wgt
2731 self.pwgt = [p/self.bias_wgt for p in self.pwgt]
2732
2733
2734 if self.type in self.real_type:
2735 self.momenta_config = self.real_related
2736 else:
2737 self.momenta_config = self.born_related
2738
2741
2743
2744
2745 - def __init__(self, momenta, wgts, event, real_type=(1,11)):
2746
2747 list.__init__(self, momenta)
2748 assert self
2749 self.soft = False
2750 self.wgts = wgts
2751 self.pdgs = list(wgts[0].pdgs)
2752 self.event = event
2753 self.real_type = real_type
2754
2755 if wgts[0].momenta_config == wgts[0].born_related:
2756
2757 ind1, ind2 = [ind-1 for ind in wgts[0].to_merge_pdg]
2758 if ind1> ind2:
2759 ind1, ind2 = ind2, ind1
2760 if ind1 >= sum(1 for p in event if p.status==-1):
2761 new_p = self[ind1] + self[ind2]
2762 else:
2763 new_p = self[ind1] - self[ind2]
2764 self.pop(ind1)
2765 self.insert(ind1, new_p)
2766 self.pop(ind2)
2767 self.pdgs.pop(ind1)
2768 self.pdgs.insert(ind1, wgts[0].merge_new_pdg )
2769 self.pdgs.pop(ind2)
2770
2771
2772 elif any(w.type in self.real_type for w in wgts):
2773 if any(w.type not in self.real_type for w in wgts):
2774 raise Exception
2775
2776
2777
2778
2779
2780
2781
2782
2783
2784
2785
2786
2787 else:
2788 raise Exception
2789
2791 """check that the propagator associated to ij is not too light
2792 [related to soft-collinear singularity]"""
2793
2794 if threshold is None:
2795 threshold = 1e-8
2796
2797 if ind1> ind2:
2798 ind1, ind2 = ind2, ind1
2799 if ind1 >= nb_init:
2800 new_p = self[ind1] + self[ind2]
2801 else:
2802 new_p = self[ind1] - self[ind2]
2803
2804 inv_mass = new_p.mass_sqr
2805 if nb_init == 2:
2806 shat = (self[0]+self[1]).mass_sqr
2807 else:
2808 shat = self[0].mass_sqr
2809
2810
2811 if (abs(inv_mass)/shat < threshold):
2812 return True
2813 else:
2814 return False
2815
2816
2819
2821 """ return the tag and order for this basic event"""
2822 (initial, _), _ = self.event.get_tag_and_order()
2823 order = self.get_pdg_code()
2824
2825
2826 initial, out = order[:len(initial)], order[len(initial):]
2827 initial.sort()
2828 out.sort()
2829 return (tuple(initial), tuple(out)), order
2830
2831 - def get_momenta(self, get_order, allow_reversed=True):
2832 """return the momenta vector in the order asked for"""
2833
2834
2835 order = [list(get_order[0]), list(get_order[1])]
2836 out = [''] *(len(order[0])+len(order[1]))
2837 pdgs = self.get_pdg_code()
2838 for pos, part in enumerate(self):
2839 if pos < len(get_order[0]):
2840 try:
2841 ind = order[0].index(pdgs[pos])
2842 except ValueError, error:
2843 if not allow_reversed:
2844 raise error
2845 else:
2846 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2847 try:
2848 return self.get_momenta(order, False)
2849 except ValueError:
2850 raise error
2851
2852
2853 position = ind
2854 order[0][ind] = 0
2855 else:
2856 try:
2857 ind = order[1].index(pdgs[pos])
2858 except ValueError, error:
2859 if not allow_reversed:
2860 raise error
2861 else:
2862 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2863 try:
2864 return self.get_momenta(order, False)
2865 except ValueError:
2866 raise error
2867 position = len(order[0]) + ind
2868 order[1][ind] = 0
2869
2870 out[position] = (part.E, part.px, part.py, part.pz)
2871
2872 return out
2873
2874
2876 return [9] * len(self)
2877
2878 @property
2880 return self.event.aqcd
2881
2883
2884 scale = 0
2885 for particle in self:
2886 p = particle
2887 scale += math.sqrt(max(0, p.mass_sqr + p.pt**2))
2888
2889 return prefactor * scale
2890
2892
2893 scale = 0
2894 for particle in self:
2895 p = particle
2896 pt = p.pt
2897 if (pt>0):
2898 scale += p.E*pt/math.sqrt(pt**2+p.pz**2)
2899
2900 return prefactor * scale
2901
2902
2904
2905 scale = 0
2906 nb_init = 0
2907 for particle in event:
2908 if particle.status == -1:
2909 nb_init+=1
2910 if nb_init == 1:
2911 return self[0].mass
2912 elif nb_init==2:
2913 return math.sqrt((self[0]+self[1])**2)
2914
2915
2916
2917
2918 - def __init__(self, input, event, real_type=(1,11), threshold=None):
2919
2920 self.real_type = real_type
2921 self.event = event
2922 self.total_wgt = 0.
2923 self.nb_event = 0
2924 self.nb_wgts = 0
2925 self.threshold = threshold
2926 self.modified = False
2927
2928
2929 if isinstance(input, str):
2930 self.parse(input)
2931
2932
2933
2935 """create the object from the string information (see example below)"""
2936
2937
2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948
2949
2950
2951
2952
2953
2954
2955
2956
2957 text = text.lower().replace('d','e')
2958 all_line = text.split('\n')
2959
2960 first_line =''
2961 while not first_line.strip():
2962 first_line = all_line.pop(0)
2963
2964 wgt, nb_wgt, nb_event, _ = first_line.split()
2965 self.total_wgt = float(wgt.replace('d','e'))
2966 nb_wgt, nb_event = int(nb_wgt), int(nb_event)
2967 self.nb_wgt, self.nb_event = nb_wgt, nb_event
2968
2969 momenta = []
2970 self.momenta = momenta
2971 wgts = []
2972 for line in all_line:
2973 data = line.split()
2974 if len(data) == 4:
2975 p = FourMomentum(data)
2976 momenta.append(p)
2977 elif len(data)>0:
2978 wgt = OneNLOWeight(line, real_type=self.real_type)
2979 wgts.append(wgt)
2980
2981 assert len(wgts) == int(nb_wgt)
2982
2983 get_weights_for_momenta = dict( (i,[]) for i in range(1,nb_event+1) )
2984 size_momenta = 0
2985 for wgt in wgts:
2986 if wgt.momenta_config in get_weights_for_momenta:
2987 get_weights_for_momenta[wgt.momenta_config].append(wgt)
2988 else:
2989 if size_momenta == 0: size_momenta = wgt.nexternal
2990 assert size_momenta == wgt.nexternal
2991 get_weights_for_momenta[wgt.momenta_config] = [wgt]
2992
2993 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)
2994
2995
2996 for key in range(1, nb_event+1):
2997 wgts = get_weights_for_momenta[key]
2998 if not wgts:
2999 continue
3000 if size_momenta == 0: size_momenta = wgts[0].nexternal
3001 p = momenta[size_momenta*(key-1):key*size_momenta]
3002 evt = self.BasicEvent(p, wgts, self.event, self.real_type)
3003 if len(evt) == size_momenta:
3004 for wgt in wgts:
3005 if not wgt.type in self.real_type:
3006 continue
3007 if evt.check_fks_singularity(wgt.to_merge_pdg[0]-1,
3008 wgt.to_merge_pdg[1]-1,
3009 nb_init=sum(1 for p in self.event if p.status==-1),
3010 threshold=self.threshold):
3011 get_weights_for_momenta[wgt.momenta_config].remove(wgt)
3012 get_weights_for_momenta[wgt.born_related].append(wgt)
3013 wgt.momenta_config = wgt.born_related
3014
3015 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)
3016
3017 self.cevents = []
3018 for key in range(1, nb_event+1):
3019 if key in get_weights_for_momenta:
3020 wgt = get_weights_for_momenta[key]
3021 if not wgt:
3022 continue
3023 pdg_to_event = {}
3024 for w in wgt:
3025 pdgs = w.pdgs
3026 if w.momenta_config == w.born_related:
3027 pdgs = list(pdgs)
3028 ind1, ind2 = [ind-1 for ind in w.to_merge_pdg]
3029 if ind1> ind2:
3030 ind1, ind2 = ind2, ind1
3031 pdgs.pop(ind1)
3032 pdgs.insert(ind1, w.merge_new_pdg )
3033 pdgs.pop(ind2)
3034 pdgs = tuple(pdgs)
3035 if pdgs not in pdg_to_event:
3036 p = momenta[size_momenta*(key-1):key*size_momenta]
3037 evt = self.BasicEvent(p, [w], self.event, self.real_type)
3038 self.cevents.append(evt)
3039 pdg_to_event[pdgs] = evt
3040 else:
3041 pdg_to_event[pdgs].wgts.append(w)
3042
3043 if __debug__:
3044 nb_wgt_check = 0
3045 for cevt in self.cevents:
3046 nb_wgt_check += len(cevt.wgts)
3047 assert nb_wgt_check == int(nb_wgt)
3048
3049
3050
3051 if '__main__' == __name__:
3052
3053 if False:
3054 lhe = EventFile('unweighted_events.lhe.gz')
3055
3056 start = time.time()
3057 for event in lhe:
3058 event.parse_lo_weight()
3059 print 'old method -> ', time.time()-start
3060 lhe = EventFile('unweighted_events.lhe.gz')
3061
3062 start = time.time()
3063 for event in lhe:
3064 event.parse_lo_weight_test()
3065 print 'new method -> ', time.time()-start
3066
3067
3068
3069 if False:
3070 start = time
3071 lhe = EventFile('unweighted_events.lhe.gz')
3072 output = open('output_events.lhe', 'w')
3073
3074 output.write(lhe.banner)
3075
3076 for event in lhe:
3077 for particle in event:
3078
3079 particle.mass = 0
3080 particle.vtim = 2
3081
3082
3083 output.write(str(event))
3084 output.write('</LesHouchesEvent>\n')
3085
3086
3087 if False:
3088 lhe = EventFile('unweighted_events.lhe.gz')
3089 import matplotlib.pyplot as plt
3090 import matplotlib.gridspec as gridspec
3091 nbins = 100
3092
3093 nb_pass = 0
3094 data = []
3095 for event in lhe:
3096 etaabs = 0
3097 etafinal = 0
3098 for particle in event:
3099 if particle.status==1:
3100 p = FourMomentum(particle)
3101 eta = p.pseudorapidity
3102 if abs(eta) > etaabs:
3103 etafinal = eta
3104 etaabs = abs(eta)
3105 if etaabs < 4:
3106 data.append(etafinal)
3107 nb_pass +=1
3108
3109
3110 print nb_pass
3111 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1])
3112 gs1.update(wspace=0, hspace=0)
3113 ax = plt.subplot(gs1[0])
3114
3115 n, bins, patches = ax.hist(data, nbins, histtype='step', label='original')
3116 ax_c = ax.twinx()
3117 ax_c.set_ylabel('MadGraph5_aMC@NLO')
3118 ax_c.yaxis.set_label_coords(1.01, 0.25)
3119 ax_c.set_yticks(ax.get_yticks())
3120 ax_c.set_yticklabels([])
3121 ax.set_xlim([-4,4])
3122 print "bin value:", n
3123 print "start/end point of bins", bins
3124 plt.axis('on')
3125 plt.xlabel('weight ratio')
3126 plt.show()
3127
3128
3129
3130 if False:
3131 lhe = EventFile('unweighted_events.lhe')
3132 import matplotlib.pyplot as plt
3133 import matplotlib.gridspec as gridspec
3134 nbins = 100
3135
3136
3137 mtau, wtau = 1.777, 4.027000e-13
3138 nb_pass = 0
3139 data, data2, data3 = [], [], []
3140 for event in lhe:
3141 nb_pass +=1
3142 if nb_pass > 10000:
3143 break
3144 tau1 = FourMomentum()
3145 tau2 = FourMomentum()
3146 for part in event:
3147 if part.pid in [-12,11,16]:
3148 momenta = FourMomentum(part)
3149 tau1 += momenta
3150 elif part.pid == 15:
3151 tau2 += FourMomentum(part)
3152
3153 if abs((mtau-tau2.mass())/wtau)<1e6 and tau2.mass() >1:
3154 data.append((tau1.mass()-mtau)/wtau)
3155 data2.append((tau2.mass()-mtau)/wtau)
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(data2, nbins, histtype='step', label='original')
3161 n2, bins2, patches2 = ax.hist(data, bins=bins, histtype='step',label='reconstructed')
3162 import cmath
3163
3164 breit = lambda m : math.sqrt(4*math.pi)*1/(((m)**2-mtau**2)**2+(mtau*wtau)**2)*wtau
3165
3166 data3 = [breit(mtau + x*wtau)*wtau*16867622.6624*50 for x in bins]
3167
3168 ax.plot(bins, data3,label='breit-wigner')
3169
3170 ax.legend()
3171
3172 ax_c = ax.twinx()
3173 ax_c.set_ylabel('MadGraph5_aMC@NLO')
3174 ax_c.yaxis.set_label_coords(1.01, 0.25)
3175 ax_c.set_yticks(ax.get_yticks())
3176 ax_c.set_yticklabels([])
3177
3178 plt.title('invariant mass of tau LHE/reconstructed')
3179 plt.axis('on')
3180 ax.set_xticklabels([])
3181
3182 ax = plt.subplot(gs1[1])
3183 data4 = [n[i]/(data3[i]) for i in range(nbins)]
3184 ax.plot(bins, data4 + [0] , 'b')
3185 data4 = [n2[i]/(data3[i]) for i in range(nbins)]
3186 ax.plot(bins, data4 + [0] , 'g')
3187 ax.set_ylim([0,2])
3188
3189 tick = ax.get_yticks()
3190 ax.set_yticks(tick[:-1])
3191
3192
3193 plt.axis('on')
3194 plt.xlabel('(M - Mtau)/Wtau')
3195 plt.show()
3196