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