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