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"] == 0:
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"] == 0:
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 = this_4mom.boost(FourMomentum(new_particle))
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
1731 """check various property of the events"""
1732
1733
1734 E, px, py, pz = 0,0,0,0
1735 absE, abspx, abspy, abspz = 0,0,0,0
1736 for particle in self:
1737 coeff = 1
1738 if particle.status == -1:
1739 coeff = -1
1740 elif particle.status != 1:
1741 continue
1742 E += coeff * particle.E
1743 absE += abs(particle.E)
1744 px += coeff * particle.px
1745 py += coeff * particle.py
1746 pz += coeff * particle.pz
1747 abspx += abs(particle.px)
1748 abspy += abs(particle.py)
1749 abspz += abs(particle.pz)
1750
1751 threshold = 5e-7
1752 if E/absE > threshold:
1753 logger.critical(self)
1754 raise Exception, "Do not conserve Energy %s, %s" % (E/absE, E)
1755 if px/abspx > threshold:
1756 logger.critical(self)
1757 raise Exception, "Do not conserve Px %s, %s" % (px/abspx, px)
1758 if py/abspy > threshold:
1759 logger.critical(self)
1760 raise Exception, "Do not conserve Py %s, %s" % (py/abspy, py)
1761 if pz/abspz > threshold:
1762 logger.critical(self)
1763 raise Exception, "Do not conserve Pz %s, %s" % (pz/abspz, pz)
1764
1765
1766 self.check_color_structure()
1767
1769 """read the line corresponding to global event line
1770 format of the line is:
1771 Nexternal IEVENT WEIGHT SCALE AEW AS
1772 """
1773 inputs = line.split()
1774 assert len(inputs) == 6
1775 self.nexternal=int(inputs[0])
1776 self.ievent=int(inputs[1])
1777 self.wgt=float(inputs[2])
1778 self.scale=float(inputs[3])
1779 self.aqed=float(inputs[4])
1780 self.aqcd=float(inputs[5])
1781
1783 """Return the unique tag identifying the SubProcesses for the generation.
1784 Usefull for program like MadSpin and Reweight module."""
1785
1786 initial, final, order = [], [], [[], []]
1787 for particle in self:
1788 if particle.status == -1:
1789 initial.append(particle.pid)
1790 order[0].append(particle.pid)
1791 elif particle.status == 1:
1792 final.append(particle.pid)
1793 order[1].append(particle.pid)
1794 initial.sort(), final.sort()
1795 tag = (tuple(initial), tuple(final))
1796 return tag, order
1797
1798 @staticmethod
1799 - def mass_shuffle(momenta, sqrts, new_mass, new_sqrts=None):
1800 """use the RAMBO method to shuffle the PS. initial sqrts is preserved."""
1801
1802 if not new_sqrts:
1803 new_sqrts = sqrts
1804
1805 oldm = [p.mass_sqr for p in momenta]
1806 newm = [m**2 for m in new_mass]
1807 tot_mom = sum(momenta, FourMomentum())
1808 if tot_mom.pt2 > 1e-5:
1809 boost_back = FourMomentum(tot_mom.mass,0,0,0).boost_to_restframe(tot_mom)
1810 for i,m in enumerate(momenta):
1811 momenta[i] = m.boost_to_restframe(tot_mom)
1812
1813
1814 f = lambda chi: new_sqrts - sum(math.sqrt(max(0, M + chi**2*(p.E**2-m)))
1815 for M,p,m in zip(newm, momenta,oldm))
1816
1817 df = lambda chi: -1* sum(chi*(p.E**2-m)/math.sqrt(max(0,(p.E**2-m)*chi**2+M))
1818 for M,p,m in zip(newm, momenta,oldm))
1819
1820 if sum(new_mass) > new_sqrts:
1821 return momenta, 0
1822 try:
1823 chi = misc.newtonmethod(f, df, 1.0, error=1e-7,maxiter=1000)
1824 except:
1825 return momenta, 0
1826
1827 new_momenta = []
1828 for i,p in enumerate(momenta):
1829 new_momenta.append(
1830 FourMomentum(math.sqrt(newm[i]+chi**2*(p.E**2-oldm[i])),
1831 chi*p.px, chi*p.py, chi*p.pz))
1832
1833
1834
1835
1836
1837
1838
1839 jac = chi**(3*len(momenta)-3)
1840 jac *= reduce(operator.mul,[p.E/k.E for p,k in zip(momenta, new_momenta)],1)
1841 jac *= sum(p.norm_sq/p.E for p in momenta)
1842 jac /= sum(k.norm_sq/k.E for k in new_momenta)
1843
1844
1845 if tot_mom.pt2 > 1e-5:
1846 for i,m in enumerate(new_momenta):
1847 new_momenta[i] = m.boost_to_restframe(boost_back)
1848 return new_momenta, jac
1849
1850
1851
1852
1854 """routine to rescale the mass via RAMBO method. no internal mass preserve.
1855 sqrts is preserve (RAMBO algo)
1856 """
1857
1858 old_momenta = []
1859 new_masses = []
1860 change_mass = False
1861 for part in self:
1862 if part.status == 1:
1863 old_momenta.append(FourMomentum(part))
1864 new_masses.append(new_param_card.get_value('mass', abs(part.pid)))
1865 if not misc.equal(part.mass, new_masses[-1], 4, zero_limit=10):
1866 change_mass = True
1867
1868 if not change_mass:
1869 return 1
1870
1871 sqrts = self.sqrts
1872
1873
1874 new_mom, jac = self.mass_shuffle(old_momenta, sqrts, new_masses)
1875
1876
1877 ind =0
1878 for part in self:
1879 if part.status==1:
1880 part.E, part.px, part.py, part.pz, part.mass = \
1881 new_mom[ind].E, new_mom[ind].px, new_mom[ind].py, new_mom[ind].pz,new_mom[ind].mass
1882 ind+=1
1883 return jac
1884
1886 """routine to rescale the momenta to change the invariant mass"""
1887
1888 old_momenta = []
1889 incoming = []
1890 masses = []
1891 for part in self:
1892 if part.status == -1:
1893 incoming.append(FourMomentum(part))
1894 if part.status == 1:
1895 old_momenta.append(FourMomentum(part))
1896 masses.append(part.mass)
1897
1898 p_init = FourMomentum()
1899 p_inits = []
1900 n_init = 0
1901 for p in incoming:
1902 n_init +=1
1903 p_init += p
1904 p_inits.append(p)
1905 old_sqrts = p_init.mass
1906
1907 new_mom, jac = self.mass_shuffle(old_momenta, old_sqrts, masses, new_sqrts=new_sqrts)
1908
1909
1910 ind =0
1911 for part in self:
1912 if part.status==1:
1913 part.E, part.px, part.py, part.pz, part.mass = \
1914 new_mom[ind].E, new_mom[ind].px, new_mom[ind].py, new_mom[ind].pz,new_mom[ind].mass
1915 ind+=1
1916
1917
1918 p_init = FourMomentum()
1919 for part in self:
1920 if part.status==1:
1921 p_init += part
1922 if n_init == 1:
1923 for part in self:
1924 if part.status == -1:
1925 part.E, part.px, part.py, part.pz = \
1926 p_init.E, p_init.px, p_init.py, p_init.pz
1927 elif n_init ==2:
1928 if not misc.equal(p_init.px, 0) or not misc.equal(p_init.py, 0):
1929 raise Exception
1930 if not misc.equal(p_inits[0].px, 0) or not misc.equal(p_inits[0].py, 0):
1931 raise Exception
1932
1933
1934
1935
1936 eta = math.log(2*p_inits[0].E/old_sqrts)
1937 new_p = [[new_sqrts/2*math.exp(eta), 0., 0., new_sqrts/2*math.exp(eta)],
1938 [new_sqrts/2*math.exp(-eta), 0., 0., -new_sqrts/2*math.exp(-eta)]]
1939
1940 ind=0
1941 for part in self:
1942 if part.status == -1:
1943 part.E, part.px, part.py, part.pz = new_p[ind]
1944 ind+=1
1945 if ind ==2:
1946 break
1947 else:
1948 raise Exception
1949
1950 return jac
1951
1952
1954 """return a list with the helicities in the order asked for"""
1955
1956
1957 order = [list(get_order[0]), list(get_order[1])]
1958 out = [9] *(len(order[0])+len(order[1]))
1959 for i, part in enumerate(self):
1960 if part.status == 1:
1961 try:
1962 ind = order[1].index(part.pid)
1963 except ValueError, error:
1964 if not allow_reversed:
1965 raise error
1966 else:
1967 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
1968 try:
1969 return self.get_helicity(order, False)
1970 except ValueError:
1971 raise error
1972 position = len(order[0]) + ind
1973 order[1][ind] = 0
1974 elif part.status == -1:
1975 try:
1976 ind = order[0].index(part.pid)
1977 except ValueError, error:
1978 if not allow_reversed:
1979 raise error
1980 else:
1981 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
1982 try:
1983 return self.get_helicity(order, False)
1984 except ValueError:
1985 raise error
1986
1987 position = ind
1988 order[0][ind] = 0
1989 else:
1990 continue
1991 out[position] = int(part.helicity)
1992 return out
1993
1994
1996 """check the validity of the color structure"""
1997
1998
1999 color_index = collections.defaultdict(int)
2000 for particle in self:
2001 if particle.status in [-1,1]:
2002 if particle.color1:
2003 color_index[particle.color1] +=1
2004 if -7 < particle.pdg < 0:
2005 raise Exception, "anti-quark with color tag"
2006 if particle.color2:
2007 color_index[particle.color2] +=1
2008 if 7 > particle.pdg > 0:
2009 raise Exception, "quark with anti-color tag"
2010
2011
2012 for key,value in color_index.items():
2013 if value > 2:
2014 print self
2015 print key, value
2016 raise Exception, 'Wrong color_flow'
2017
2018
2019
2020 check = []
2021 popup_index = []
2022 for particle in self:
2023 mothers = []
2024 childs = []
2025 if particle.mother1:
2026 mothers.append(particle.mother1)
2027 if particle.mother2 and particle.mother2 is not particle.mother1:
2028 mothers.append(particle.mother2)
2029 if not mothers:
2030 continue
2031 if (particle.mother1.event_id, particle.mother2.event_id) in check:
2032 continue
2033 check.append((particle.mother1.event_id, particle.mother2.event_id))
2034
2035 childs = [p for p in self if p.mother1 is particle.mother1 and \
2036 p.mother2 is particle.mother2]
2037
2038 mcolors = []
2039 manticolors = []
2040 for m in mothers:
2041 if m.color1:
2042 if m.color1 in manticolors:
2043 manticolors.remove(m.color1)
2044 else:
2045 mcolors.append(m.color1)
2046 if m.color2:
2047 if m.color2 in mcolors:
2048 mcolors.remove(m.color2)
2049 else:
2050 manticolors.append(m.color2)
2051 ccolors = []
2052 canticolors = []
2053 for m in childs:
2054 if m.color1:
2055 if m.color1 in canticolors:
2056 canticolors.remove(m.color1)
2057 else:
2058 ccolors.append(m.color1)
2059 if m.color2:
2060 if m.color2 in ccolors:
2061 ccolors.remove(m.color2)
2062 else:
2063 canticolors.append(m.color2)
2064 for index in mcolors[:]:
2065 if index in ccolors:
2066 mcolors.remove(index)
2067 ccolors.remove(index)
2068 for index in manticolors[:]:
2069 if index in canticolors:
2070 manticolors.remove(index)
2071 canticolors.remove(index)
2072
2073 if mcolors != []:
2074
2075 if len(canticolors) + len(mcolors) != 3:
2076 logger.critical(str(self))
2077 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs])
2078 else:
2079 popup_index += canticolors
2080 elif manticolors != []:
2081
2082 if len(ccolors) + len(manticolors) != 3:
2083 logger.critical(str(self))
2084 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs])
2085 else:
2086 popup_index += ccolors
2087
2088
2089 if len(popup_index) != len(set(popup_index)):
2090 logger.critical(self)
2091 raise Exception, "Wrong color flow: identical poping-up index, %s" % (popup_index)
2092
2094 """two event are the same if they have the same momentum. other info are ignored"""
2095
2096 if other is None:
2097 return False
2098
2099 for i,p in enumerate(self):
2100 if p.E != other[i].E:
2101 return False
2102 elif p.pz != other[i].pz:
2103 return False
2104 elif p.px != other[i].px:
2105 return False
2106 elif p.py != other[i].py:
2107 return False
2108 return True
2109
2110
2112 """return a correctly formatted LHE event"""
2113
2114 out="""<event%(event_flag)s>
2115 %(scale)s
2116 %(particles)s
2117 %(comments)s
2118 %(tag)s
2119 %(reweight)s
2120 </event>
2121 """
2122 if event_id not in ['', None]:
2123 self.eventflag['event'] = str(event_id)
2124
2125 if self.eventflag:
2126 event_flag = ' %s' % ' '.join('%s="%s"' % (k,v) for (k,v) in self.eventflag.items())
2127 else:
2128 event_flag = ''
2129
2130 scale_str = "%2d %6d %+13.7e %14.8e %14.8e %14.8e" % \
2131 (self.nexternal,self.ievent,self.wgt,self.scale,self.aqed,self.aqcd)
2132
2133
2134 if self.reweight_data:
2135
2136 if set(self.reweight_data.keys()) != set(self.reweight_order):
2137 self.reweight_order += [k for k in self.reweight_data.keys() \
2138 if k not in self.reweight_order]
2139
2140 reweight_str = '<rwgt>\n%s\n</rwgt>' % '\n'.join(
2141 '<wgt id=\'%s\'> %+13.7e </wgt>' % (i, float(self.reweight_data[i]))
2142 for i in self.reweight_order if i in self.reweight_data)
2143 else:
2144 reweight_str = ''
2145
2146 tag_str = self.tag
2147 if hasattr(self, 'nloweight') and self.nloweight.modified:
2148 self.rewrite_nlo_weight()
2149 tag_str = self.tag
2150
2151 if self.matched_scale_data:
2152 tmp_scale = ' '.join(['pt_clust_%i=\"%s\"' % (i+1,v)
2153 for i,v in enumerate(self.matched_scale_data)
2154 if v!=-1])
2155 if tmp_scale:
2156 tag_str = "<scales %s></scales>%s" % (tmp_scale, self.tag)
2157
2158 if self.syscalc_data:
2159 keys= ['rscale', 'asrwt', ('pdfrwt', 'beam', '1'), ('pdfrwt', 'beam', '2'),
2160 'matchscale', 'totfact']
2161 sys_str = "<mgrwt>\n"
2162 template = """<%(key)s%(opts)s>%(values)s</%(key)s>\n"""
2163 for k in keys:
2164 if k not in self.syscalc_data:
2165 continue
2166 replace = {}
2167 replace['values'] = self.syscalc_data[k]
2168 if isinstance(k, str):
2169 replace['key'] = k
2170 replace['opts'] = ''
2171 else:
2172 replace['key'] = k[0]
2173 replace['opts'] = ' %s=\"%s\"' % (k[1],k[2])
2174 sys_str += template % replace
2175 sys_str += "</mgrwt>\n"
2176 reweight_str = sys_str + reweight_str
2177
2178 out = out % {'event_flag': event_flag,
2179 'scale': scale_str,
2180 'particles': '\n'.join([str(p) for p in self]),
2181 'tag': tag_str,
2182 'comments': self.comment,
2183 'reweight': reweight_str}
2184
2185 return re.sub('[\n]+', '\n', out)
2186
2187 - def get_momenta(self, get_order, allow_reversed=True):
2188 """return the momenta vector in the order asked for"""
2189
2190
2191 order = [list(get_order[0]), list(get_order[1])]
2192 out = [''] *(len(order[0])+len(order[1]))
2193 for i, part in enumerate(self):
2194 if part.status == 1:
2195 try:
2196 ind = order[1].index(part.pid)
2197 except ValueError, error:
2198 if not allow_reversed:
2199 raise error
2200 else:
2201 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2202 try:
2203 return self.get_momenta_str(order, False)
2204 except ValueError:
2205 raise error
2206 position = len(order[0]) + ind
2207 order[1][ind] = 0
2208 elif part.status == -1:
2209 try:
2210 ind = order[0].index(part.pid)
2211 except ValueError, error:
2212 if not allow_reversed:
2213 raise error
2214 else:
2215 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2216 try:
2217 return self.get_momenta_str(order, False)
2218 except ValueError:
2219 raise error
2220
2221 position = ind
2222 order[0][ind] = 0
2223 else:
2224 continue
2225
2226 out[position] = (part.E, part.px, part.py, part.pz)
2227
2228 return out
2229
2230
2243
2244
2246
2247 scale = 0
2248 for particle in self:
2249 if particle.status != 1:
2250 continue
2251 p=FourMomentum(particle)
2252 scale += math.sqrt(p.mass_sqr + p.pt**2)
2253
2254 return prefactor * scale
2255
2256
2258
2259 scale = 0
2260 for particle in self:
2261 if particle.status != 1:
2262 continue
2263 p = FourMomentum(particle)
2264 pt = p.pt
2265 if (pt>0):
2266 scale += p.E*pt/math.sqrt(pt**2+p.pz**2)
2267
2268 return prefactor * scale
2269
2270 @property
2273
2275
2276 scale = 0
2277 init = []
2278 for particle in self:
2279 if particle.status == -1:
2280 init.append(FourMomentum(particle))
2281 if len(init) == 1:
2282 return init[0].mass
2283 elif len(init)==2:
2284 return math.sqrt((init[0]+init[1])**2)
2285
2286
2287
2288
2290 """return the momenta str in the order asked for"""
2291
2292 out = self.get_momenta(get_order, allow_reversed)
2293
2294 format = '%.12f'
2295 format_line = ' '.join([format]*4) + ' \n'
2296 out = [format_line % one for one in out]
2297 out = ''.join(out).replace('e','d')
2298 return out
2299
2301 """A class to allow to read both gzip and not gzip file.
2302 containing only weight from pythia --generated by SysCalc"""
2303
2304 - def __new__(self, path, mode='r', *args, **opt):
2305 if path.endswith(".gz"):
2306 try:
2307 return gzip.GzipFile.__new__(WeightFileGzip, path, mode, *args, **opt)
2308 except IOError, error:
2309 raise
2310 except Exception, error:
2311 if mode == 'r':
2312 misc.gunzip(path)
2313 return file.__new__(WeightFileNoGzip, path[:-3], mode, *args, **opt)
2314 else:
2315 return file.__new__(WeightFileNoGzip, path, mode, *args, **opt)
2316
2317
2318 - def __init__(self, path, mode='r', *args, **opt):
2319 """open file and read the banner [if in read mode]"""
2320
2321 super(EventFile, self).__init__(path, mode, *args, **opt)
2322 self.banner = ''
2323 if mode == 'r':
2324 line = ''
2325 while '</header>' not in line.lower():
2326 try:
2327 line = super(EventFile, self).next()
2328 except StopIteration:
2329 self.seek(0)
2330 self.banner = ''
2331 break
2332 if "<event" in line.lower():
2333 self.seek(0)
2334 self.banner = ''
2335 break
2336
2337 self.banner += line
2338
2342
2345
2348 """a convenient object for 4-momenta operation"""
2349
2350 - def __init__(self, obj=0, px=0, py=0, pz=0, E=0):
2351 """initialize the four momenta"""
2352
2353 if obj is 0 and E:
2354 obj = E
2355
2356 if isinstance(obj, (FourMomentum, Particle)):
2357 px = obj.px
2358 py = obj.py
2359 pz = obj.pz
2360 E = obj.E
2361 elif isinstance(obj, (list, tuple)):
2362 assert len(obj) ==4
2363 E = obj[0]
2364 px = obj[1]
2365 py = obj[2]
2366 pz = obj[3]
2367 elif isinstance(obj, str):
2368 obj = [float(i) for i in obj.split()]
2369 assert len(obj) ==4
2370 E = obj[0]
2371 px = obj[1]
2372 py = obj[2]
2373 pz = obj[3]
2374 else:
2375 E =obj
2376
2377
2378 self.E = float(E)
2379 self.px = float(px)
2380 self.py = float(py)
2381 self.pz = float(pz)
2382
2383 @property
2385 """return the mass"""
2386 return math.sqrt(max(self.E**2 - self.px**2 - self.py**2 - self.pz**2,0))
2387
2388 @property
2390 """return the mass square"""
2391 return self.E**2 - self.px**2 - self.py**2 - self.pz**2
2392
2393 @property
2395 return math.sqrt(max(0, self.pt2))
2396
2397 @property
2399 norm = math.sqrt(self.px**2 + self.py**2+self.pz**2)
2400 return 0.5* math.log((norm - self.pz) / (norm + self.pz))
2401
2402 @property
2404 return 0.5* math.log((self.E +self.pz) / (self.E - self.pz))
2405
2406
2407 @property
2409 """ return the pt square """
2410
2411 return self.px**2 + self.py**2
2412
2413 @property
2415 """ return |\vec p| """
2416 return math.sqrt(self.px**2 + self.py**2 + self.pz**2)
2417
2418 @property
2420 """ return |\vec p|^2 """
2421 return self.px**2 + self.py**2 + self.pz**2
2422
2423 @property
2425 """return the mass square"""
2426 import math
2427 return math.atan(math.sqrt((self.px**2+self.py**2)/self.pz**2))
2428
2429
2431
2432 assert isinstance(obj, FourMomentum)
2433 new = FourMomentum(self.E+obj.E,
2434 self.px + obj.px,
2435 self.py + obj.py,
2436 self.pz + obj.pz)
2437 return new
2438
2440 """update the object with the sum"""
2441 self.E += obj.E
2442 self.px += obj.px
2443 self.py += obj.py
2444 self.pz += obj.pz
2445 return self
2446
2448
2449 assert isinstance(obj, FourMomentum)
2450 new = FourMomentum(self.E-obj.E,
2451 self.px - obj.px,
2452 self.py - obj.py,
2453 self.pz - obj.pz)
2454 return new
2455
2457 """update the object with the sum"""
2458 self.E -= obj.E
2459 self.px -= obj.px
2460 self.py -= obj.py
2461 self.pz -= obj.pz
2462 return self
2463
2465 if isinstance(obj, FourMomentum):
2466 return self.E*obj.E - self.px *obj.px - self.py * obj.py - self.pz * obj.pz
2467 elif isinstance(obj, (float, int)):
2468 return FourMomentum(obj*self.E,obj*self.px,obj*self.py,obj*self.pz )
2469 else:
2470 raise NotImplemented
2471 __rmul__ = __mul__
2472
2474 assert power in [1,2]
2475
2476 if power == 1:
2477 return FourMomentum(self)
2478 elif power == 2:
2479 return self.mass_sqr
2480
2482 return 'FourMomentum(%s,%s,%s,%s)' % (self.E, self.px, self.py,self.pz)
2483
2484 - def __str__(self, mode='python'):
2485 if mode == 'python':
2486 return self.__repr__()
2487 elif mode == 'fortran':
2488 return '%.10e %.10e %.10e %.10e' % self.get_tuple()
2489
2491 return (self.E, self.px, self.py,self.pz)
2492
2494 """mom 4-momenta is suppose to be given in the rest frame of this 4-momenta.
2495 the output is the 4-momenta in the frame of this 4-momenta
2496 function copied from HELAS routine."""
2497
2498
2499 pt = self.px**2 + self.py**2 + self.pz**2
2500 if pt:
2501 s3product = self.px * mom.px + self.py * mom.py + self.pz * mom.pz
2502 mass = self.mass
2503 lf = (mom.E + (self.E - mass) * s3product / pt ) / mass
2504 return FourMomentum(E=(self.E*mom.E+s3product)/mass,
2505 px=mom.px + self.px * lf,
2506 py=mom.py + self.py * lf,
2507 pz=mom.pz + self.pz * lf)
2508 else:
2509 return FourMomentum(mom)
2510
2511 - def zboost(self, pboost=None, E=0, pz=0):
2512 """Both momenta should be in the same frame.
2513 The boost perform correspond to the boost required to set pboost at
2514 rest (only z boost applied).
2515 """
2516 if isinstance(pboost, FourMomentum):
2517 E = pboost.E
2518 pz = pboost.pz
2519
2520
2521 gamma = E / math.sqrt(E**2-pz**2)
2522 gammabeta = pz / math.sqrt(E**2-pz**2)
2523
2524 out = FourMomentum([gamma*self.E - gammabeta*self.pz,
2525 self.px,
2526 self.py,
2527 gamma*self.pz - gammabeta*self.E])
2528
2529 if abs(out.pz) < 1e-6 * out.E:
2530 out.pz = 0
2531 return out
2532
2534 """apply the boost transformation such that pboost is at rest in the new frame.
2535 First apply a rotation to allign the pboost to the z axis and then use
2536 zboost routine (see above)
2537 """
2538
2539 if pboost.px == 0 == pboost.py:
2540 out = self.zboost(E=pboost.E,pz=pboost.pz)
2541 return out
2542
2543
2544
2545
2546
2547
2548
2549 p = math.sqrt( pboost.px**2 + pboost.py**2+ pboost.pz**2)
2550 cosF = pboost.pz / p
2551 sinF = math.sqrt(1-cosF**2)
2552 sinT = pboost.py/p/sinF
2553 cosT = pboost.px/p/sinF
2554
2555 out=FourMomentum([self.E,
2556 self.px*cosT*cosF + self.py*sinT*cosF-self.pz*sinF,
2557 -self.px*sinT+ self.py*cosT,
2558 self.px*cosT*sinF + self.py*sinT*sinF + self.pz*cosF
2559 ])
2560 out = out.zboost(E=pboost.E,pz=p)
2561 return out
2562
2567
2568 - def __init__(self, input, real_type=(1,11)):
2569 """ """
2570
2571 self.real_type = real_type
2572 if isinstance(input, str):
2573 self.parse(input)
2574
2575 - def __str__(self, mode='display'):
2576
2577 if mode == 'display':
2578 out = """ pwgt: %(pwgt)s
2579 born, real : %(born)s %(real)s
2580 pdgs : %(pdgs)s
2581 bjks : %(bjks)s
2582 scales**2, gs: %(scales2)s %(gs)s
2583 born/real related : %(born_related)s %(real_related)s
2584 type / nfks : %(type)s %(nfks)s
2585 to merge : %(to_merge_pdg)s in %(merge_new_pdg)s
2586 ref_wgt : %(ref_wgt)s""" % self.__dict__
2587 return out
2588 elif mode == 'formatted':
2589 format_var = []
2590 variable = []
2591
2592 def to_add_full(f, v, format_var, variable):
2593 """ function to add to the formatted output"""
2594 if isinstance(v, list):
2595 format_var += [f]*len(v)
2596 variable += v
2597 else:
2598 format_var.append(f)
2599 variable.append(v)
2600 to_add = lambda x,y: to_add_full(x,y, format_var, variable)
2601
2602 to_add('%.10e', [p*self.bias_wgt for p in self.pwgt])
2603 to_add('%.10e', self.born)
2604 to_add('%.10e', self.real)
2605 to_add('%i', self.nexternal)
2606 to_add('%i', self.pdgs)
2607 to_add('%i', self.qcdpower)
2608 to_add('%.10e', self.bjks)
2609 to_add('%.10e', self.scales2)
2610 to_add('%.10e', self.gs)
2611 to_add('%i', [self.born_related, self.real_related])
2612 to_add('%i' , [self.type, self.nfks])
2613 to_add('%i' , self.to_merge_pdg)
2614 to_add('%i', self.merge_new_pdg)
2615 to_add('%.10e', self.ref_wgt*self.bias_wgt)
2616 to_add('%.10e', self.bias_wgt)
2617 return ' '.join(format_var) % tuple(variable)
2618
2619
2620 - def parse(self, text, keep_bias=False):
2621 """parse the line and create the related object.
2622 keep bias allow to not systematically correct for the bias in the written information"""
2623
2624
2625
2626 data = text.split()
2627
2628
2629
2630
2631 self.pwgt = [float(f) for f in data[:3]]
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645
2646
2647
2648
2649
2650 self.born = float(data[3])
2651 self.real = float(data[4])
2652
2653
2654 self.nexternal = int(data[5])
2655
2656
2657 self.pdgs = [int(i) for i in data[6:6+self.nexternal]]
2658 flag = 6+self.nexternal
2659
2660
2661 self.qcdpower = int(data[flag])
2662
2663
2664 self.bjks = [float(f) for f in data[flag+1:flag+3]]
2665
2666
2667 self.scales2 = [float(f) for f in data[flag+3:flag+6]]
2668
2669
2670 self.gs = float(data[flag+6])
2671
2672
2673
2674
2675
2676
2677
2678 self.born_related = int(data[flag+7])
2679 self.real_related = int(data[flag+8])
2680
2681
2682
2683
2684
2685
2686
2687
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
2698
2699
2700
2701
2702
2703 self.type = int(data[flag+9])
2704
2705
2706
2707
2708 self.nfks = int(data[flag+10])
2709
2710
2711
2712
2713
2714
2715
2716 self.to_merge_pdg = [int (f) for f in data[flag+11:flag+13]]
2717
2718
2719 self.merge_new_pdg = int(data[flag+13])
2720
2721
2722
2723
2724
2725 self.ref_wgt = float(data[flag+14])
2726
2727
2728
2729 try:
2730 self.bias_wgt = float(data[flag+15])
2731 except IndexError:
2732 self.bias_wgt = 1.0
2733
2734 if not keep_bias:
2735 self.ref_wgt /= self.bias_wgt
2736 self.pwgt = [p/self.bias_wgt for p in self.pwgt]
2737
2738
2739 if self.type in self.real_type:
2740 self.momenta_config = self.real_related
2741 else:
2742 self.momenta_config = self.born_related
2743
2746
2748
2749
2750 - def __init__(self, momenta, wgts, event, real_type=(1,11)):
2751
2752 list.__init__(self, momenta)
2753 assert self
2754 self.soft = False
2755 self.wgts = wgts
2756 self.pdgs = list(wgts[0].pdgs)
2757 self.event = event
2758 self.real_type = real_type
2759
2760 if wgts[0].momenta_config == wgts[0].born_related:
2761
2762 ind1, ind2 = [ind-1 for ind in wgts[0].to_merge_pdg]
2763 if ind1> ind2:
2764 ind1, ind2 = ind2, ind1
2765 if ind1 >= sum(1 for p in event if p.status==-1):
2766 new_p = self[ind1] + self[ind2]
2767 else:
2768 new_p = self[ind1] - self[ind2]
2769 self.pop(ind1)
2770 self.insert(ind1, new_p)
2771 self.pop(ind2)
2772 self.pdgs.pop(ind1)
2773 self.pdgs.insert(ind1, wgts[0].merge_new_pdg )
2774 self.pdgs.pop(ind2)
2775
2776
2777 elif any(w.type in self.real_type for w in wgts):
2778 if any(w.type not in self.real_type for w in wgts):
2779 raise Exception
2780
2781
2782
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792 else:
2793 raise Exception
2794
2796 """check that the propagator associated to ij is not too light
2797 [related to soft-collinear singularity]"""
2798
2799 if threshold is None:
2800 threshold = 1e-8
2801
2802 if ind1> ind2:
2803 ind1, ind2 = ind2, ind1
2804 if ind1 >= nb_init:
2805 new_p = self[ind1] + self[ind2]
2806 else:
2807 new_p = self[ind1] - self[ind2]
2808
2809 inv_mass = new_p.mass_sqr
2810 if nb_init == 2:
2811 shat = (self[0]+self[1]).mass_sqr
2812 else:
2813 shat = self[0].mass_sqr
2814
2815
2816 if (abs(inv_mass)/shat < threshold):
2817 return True
2818 else:
2819 return False
2820
2821
2824
2826 """ return the tag and order for this basic event"""
2827 (initial, _), _ = self.event.get_tag_and_order()
2828 order = self.get_pdg_code()
2829
2830
2831 initial, out = order[:len(initial)], order[len(initial):]
2832 initial.sort()
2833 out.sort()
2834 return (tuple(initial), tuple(out)), order
2835
2836 - def get_momenta(self, get_order, allow_reversed=True):
2837 """return the momenta vector in the order asked for"""
2838
2839
2840 order = [list(get_order[0]), list(get_order[1])]
2841 out = [''] *(len(order[0])+len(order[1]))
2842 pdgs = self.get_pdg_code()
2843 for pos, part in enumerate(self):
2844 if pos < len(get_order[0]):
2845 try:
2846 ind = order[0].index(pdgs[pos])
2847 except ValueError, error:
2848 if not allow_reversed:
2849 raise error
2850 else:
2851 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2852 try:
2853 return self.get_momenta(order, False)
2854 except ValueError:
2855 raise error
2856
2857
2858 position = ind
2859 order[0][ind] = 0
2860 else:
2861 try:
2862 ind = order[1].index(pdgs[pos])
2863 except ValueError, error:
2864 if not allow_reversed:
2865 raise error
2866 else:
2867 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]]
2868 try:
2869 return self.get_momenta(order, False)
2870 except ValueError:
2871 raise error
2872 position = len(order[0]) + ind
2873 order[1][ind] = 0
2874
2875 out[position] = (part.E, part.px, part.py, part.pz)
2876
2877 return out
2878
2879
2881 return [9] * len(self)
2882
2883 @property
2885 return self.event.aqcd
2886
2888
2889 scale = 0
2890 for particle in self:
2891 p = particle
2892 scale += math.sqrt(max(0, p.mass_sqr + p.pt**2))
2893
2894 return prefactor * scale
2895
2897
2898 scale = 0
2899 for particle in self:
2900 p = particle
2901 pt = p.pt
2902 if (pt>0):
2903 scale += p.E*pt/math.sqrt(pt**2+p.pz**2)
2904
2905 return prefactor * scale
2906
2907
2909
2910 scale = 0
2911 nb_init = 0
2912 for particle in event:
2913 if particle.status == -1:
2914 nb_init+=1
2915 if nb_init == 1:
2916 return self[0].mass
2917 elif nb_init==2:
2918 return math.sqrt((self[0]+self[1])**2)
2919
2920
2921
2922
2923 - def __init__(self, input, event, real_type=(1,11), threshold=None):
2924
2925 self.real_type = real_type
2926 self.event = event
2927 self.total_wgt = 0.
2928 self.nb_event = 0
2929 self.nb_wgts = 0
2930 self.threshold = threshold
2931 self.modified = False
2932
2933
2934 if isinstance(input, str):
2935 self.parse(input)
2936
2937
2938
2940 """create the object from the string information (see example below)"""
2941
2942
2943
2944
2945
2946
2947
2948
2949
2950
2951
2952
2953
2954
2955
2956
2957
2958
2959
2960
2961
2962 text = text.lower().replace('d','e')
2963 all_line = text.split('\n')
2964
2965 first_line =''
2966 while not first_line.strip():
2967 first_line = all_line.pop(0)
2968
2969 wgt, nb_wgt, nb_event, _ = first_line.split()
2970 self.total_wgt = float(wgt.replace('d','e'))
2971 nb_wgt, nb_event = int(nb_wgt), int(nb_event)
2972 self.nb_wgt, self.nb_event = nb_wgt, nb_event
2973
2974 momenta = []
2975 self.momenta = momenta
2976 wgts = []
2977 for line in all_line:
2978 data = line.split()
2979 if len(data) == 4:
2980 p = FourMomentum(data)
2981 momenta.append(p)
2982 elif len(data)>0:
2983 wgt = OneNLOWeight(line, real_type=self.real_type)
2984 wgts.append(wgt)
2985
2986 assert len(wgts) == int(nb_wgt)
2987
2988 get_weights_for_momenta = dict( (i,[]) for i in range(1,nb_event+1) )
2989 size_momenta = 0
2990 for wgt in wgts:
2991 if wgt.momenta_config in get_weights_for_momenta:
2992 get_weights_for_momenta[wgt.momenta_config].append(wgt)
2993 else:
2994 if size_momenta == 0: size_momenta = wgt.nexternal
2995 assert size_momenta == wgt.nexternal
2996 get_weights_for_momenta[wgt.momenta_config] = [wgt]
2997
2998 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)
2999
3000
3001 for key in range(1, nb_event+1):
3002 wgts = get_weights_for_momenta[key]
3003 if not wgts:
3004 continue
3005 if size_momenta == 0: size_momenta = wgts[0].nexternal
3006 p = momenta[size_momenta*(key-1):key*size_momenta]
3007 evt = self.BasicEvent(p, wgts, self.event, self.real_type)
3008 if len(evt) == size_momenta:
3009 for wgt in wgts:
3010 if not wgt.type in self.real_type:
3011 continue
3012 if evt.check_fks_singularity(wgt.to_merge_pdg[0]-1,
3013 wgt.to_merge_pdg[1]-1,
3014 nb_init=sum(1 for p in self.event if p.status==-1),
3015 threshold=self.threshold):
3016 get_weights_for_momenta[wgt.momenta_config].remove(wgt)
3017 get_weights_for_momenta[wgt.born_related].append(wgt)
3018 wgt.momenta_config = wgt.born_related
3019
3020 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)
3021
3022 self.cevents = []
3023 for key in range(1, nb_event+1):
3024 if key in get_weights_for_momenta:
3025 wgt = get_weights_for_momenta[key]
3026 if not wgt:
3027 continue
3028 pdg_to_event = {}
3029 for w in wgt:
3030 pdgs = w.pdgs
3031 if w.momenta_config == w.born_related:
3032 pdgs = list(pdgs)
3033 ind1, ind2 = [ind-1 for ind in w.to_merge_pdg]
3034 if ind1> ind2:
3035 ind1, ind2 = ind2, ind1
3036 pdgs.pop(ind1)
3037 pdgs.insert(ind1, w.merge_new_pdg )
3038 pdgs.pop(ind2)
3039 pdgs = tuple(pdgs)
3040 if pdgs not in pdg_to_event:
3041 p = momenta[size_momenta*(key-1):key*size_momenta]
3042 evt = self.BasicEvent(p, [w], self.event, self.real_type)
3043 self.cevents.append(evt)
3044 pdg_to_event[pdgs] = evt
3045 else:
3046 pdg_to_event[pdgs].wgts.append(w)
3047
3048 if __debug__:
3049 nb_wgt_check = 0
3050 for cevt in self.cevents:
3051 nb_wgt_check += len(cevt.wgts)
3052 assert nb_wgt_check == int(nb_wgt)
3053
3054
3055
3056 if '__main__' == __name__:
3057
3058 if False:
3059 lhe = EventFile('unweighted_events.lhe.gz')
3060
3061 start = time.time()
3062 for event in lhe:
3063 event.parse_lo_weight()
3064 print 'old method -> ', time.time()-start
3065 lhe = EventFile('unweighted_events.lhe.gz')
3066
3067 start = time.time()
3068 for event in lhe:
3069 event.parse_lo_weight_test()
3070 print 'new method -> ', time.time()-start
3071
3072
3073
3074 if False:
3075 start = time
3076 lhe = EventFile('unweighted_events.lhe.gz')
3077 output = open('output_events.lhe', 'w')
3078
3079 output.write(lhe.banner)
3080
3081 for event in lhe:
3082 for particle in event:
3083
3084 particle.mass = 0
3085 particle.vtim = 2
3086
3087
3088 output.write(str(event))
3089 output.write('</LesHouchesEvent>\n')
3090
3091
3092 if False:
3093 lhe = EventFile('unweighted_events.lhe.gz')
3094 import matplotlib.pyplot as plt
3095 import matplotlib.gridspec as gridspec
3096 nbins = 100
3097
3098 nb_pass = 0
3099 data = []
3100 for event in lhe:
3101 etaabs = 0
3102 etafinal = 0
3103 for particle in event:
3104 if particle.status==1:
3105 p = FourMomentum(particle)
3106 eta = p.pseudorapidity
3107 if abs(eta) > etaabs:
3108 etafinal = eta
3109 etaabs = abs(eta)
3110 if etaabs < 4:
3111 data.append(etafinal)
3112 nb_pass +=1
3113
3114
3115 print nb_pass
3116 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1])
3117 gs1.update(wspace=0, hspace=0)
3118 ax = plt.subplot(gs1[0])
3119
3120 n, bins, patches = ax.hist(data, nbins, histtype='step', label='original')
3121 ax_c = ax.twinx()
3122 ax_c.set_ylabel('MadGraph5_aMC@NLO')
3123 ax_c.yaxis.set_label_coords(1.01, 0.25)
3124 ax_c.set_yticks(ax.get_yticks())
3125 ax_c.set_yticklabels([])
3126 ax.set_xlim([-4,4])
3127 print "bin value:", n
3128 print "start/end point of bins", bins
3129 plt.axis('on')
3130 plt.xlabel('weight ratio')
3131 plt.show()
3132
3133
3134
3135 if False:
3136 lhe = EventFile('unweighted_events.lhe')
3137 import matplotlib.pyplot as plt
3138 import matplotlib.gridspec as gridspec
3139 nbins = 100
3140
3141
3142 mtau, wtau = 1.777, 4.027000e-13
3143 nb_pass = 0
3144 data, data2, data3 = [], [], []
3145 for event in lhe:
3146 nb_pass +=1
3147 if nb_pass > 10000:
3148 break
3149 tau1 = FourMomentum()
3150 tau2 = FourMomentum()
3151 for part in event:
3152 if part.pid in [-12,11,16]:
3153 momenta = FourMomentum(part)
3154 tau1 += momenta
3155 elif part.pid == 15:
3156 tau2 += FourMomentum(part)
3157
3158 if abs((mtau-tau2.mass())/wtau)<1e6 and tau2.mass() >1:
3159 data.append((tau1.mass()-mtau)/wtau)
3160 data2.append((tau2.mass()-mtau)/wtau)
3161 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1])
3162 gs1.update(wspace=0, hspace=0)
3163 ax = plt.subplot(gs1[0])
3164
3165 n, bins, patches = ax.hist(data2, nbins, histtype='step', label='original')
3166 n2, bins2, patches2 = ax.hist(data, bins=bins, histtype='step',label='reconstructed')
3167 import cmath
3168
3169 breit = lambda m : math.sqrt(4*math.pi)*1/(((m)**2-mtau**2)**2+(mtau*wtau)**2)*wtau
3170
3171 data3 = [breit(mtau + x*wtau)*wtau*16867622.6624*50 for x in bins]
3172
3173 ax.plot(bins, data3,label='breit-wigner')
3174
3175 ax.legend()
3176
3177 ax_c = ax.twinx()
3178 ax_c.set_ylabel('MadGraph5_aMC@NLO')
3179 ax_c.yaxis.set_label_coords(1.01, 0.25)
3180 ax_c.set_yticks(ax.get_yticks())
3181 ax_c.set_yticklabels([])
3182
3183 plt.title('invariant mass of tau LHE/reconstructed')
3184 plt.axis('on')
3185 ax.set_xticklabels([])
3186
3187 ax = plt.subplot(gs1[1])
3188 data4 = [n[i]/(data3[i]) for i in range(nbins)]
3189 ax.plot(bins, data4 + [0] , 'b')
3190 data4 = [n2[i]/(data3[i]) for i in range(nbins)]
3191 ax.plot(bins, data4 + [0] , 'g')
3192 ax.set_ylim([0,2])
3193
3194 tick = ax.get_yticks()
3195 ax.set_yticks(tick[:-1])
3196
3197
3198 plt.axis('on')
3199 plt.xlabel('(M - Mtau)/Wtau')
3200 plt.show()
3201