Package madgraph :: Package various :: Module lhe_parser
[hide private]
[frames] | no frames]

Source Code for Module madgraph.various.lhe_parser

   1  from __future__ import division 
   2  import collections 
   3  import random 
   4  import re 
   5  import numbers 
   6  import math 
   7  import time 
   8  import os 
   9  import shutil 
  10  import sys 
  11   
  12  pjoin = os.path.join 
  13   
  14  if '__main__' == __name__: 
  15      import sys 
  16      sys.path.append('../../') 
  17  import misc 
  18  import logging 
  19  import gzip 
  20  import banner as banner_mod 
  21  logger = logging.getLogger("madgraph.lhe_parser") 
22 23 -class Particle(object):
24 """ """ 25 # regular expression not use anymore to speed up the computation 26 #pattern=re.compile(r'''^\s* 27 # (?P<pid>-?\d+)\s+ #PID 28 # (?P<status>-?\d+)\s+ #status (1 for output particle) 29 # (?P<mother1>-?\d+)\s+ #mother 30 # (?P<mother2>-?\d+)\s+ #mother 31 # (?P<color1>[+-e.\d]*)\s+ #color1 32 # (?P<color2>[+-e.\d]*)\s+ #color2 33 # (?P<px>[+-e.\d]*)\s+ #px 34 # (?P<py>[+-e.\d]*)\s+ #py 35 # (?P<pz>[+-e.\d]*)\s+ #pz 36 # (?P<E>[+-e.\d]*)\s+ #E 37 # (?P<mass>[+-e.\d]*)\s+ #mass 38 # (?P<vtim>[+-e.\d]*)\s+ #displace vertex 39 # (?P<helicity>[+-e.\d]*)\s* #helicity 40 # ($|(?P<comment>\#[\d|D]*)) #comment/end of string 41 # ''',66) #verbose+ignore case 42 43 44
45 - def __init__(self, line=None, event=None):
46 """ """ 47 48 if isinstance(line, Particle): 49 for key in line.__dict__: 50 setattr(self, key, getattr(line, key)) 51 if event: 52 self.event = event 53 return 54 55 self.event = event 56 self.event_id = len(event) #not yet in the event 57 # LHE information 58 self.pid = 0 59 self.status = 0 # -1:initial. 1:final. 2: propagator 60 self.mother1 = None 61 self.mother2 = None 62 self.color1 = 0 63 self.color2 = None 64 self.px = 0 65 self.py = 0 66 self.pz = 0 67 self.E = 0 68 self.mass = 0 69 self.vtim = 0 70 self.helicity = 9 71 self.rwgt = 0 72 self.comment = '' 73 74 if line: 75 self.parse(line)
76 77 @property
78 - def pdg(self):
79 "convenient alias" 80 return self.pid
81
82 - def parse(self, line):
83 """parse the line""" 84 85 args = line.split() 86 keys = ['pid', 'status','mother1','mother2','color1', 'color2', 'px','py','pz','E', 87 'mass','vtim','helicity'] 88 89 for key,value in zip(keys,args): 90 setattr(self, key, float(value)) 91 self.pid = int(self.pid) 92 93 self.comment = ' '.join(args[len(keys):]) 94 if self.comment.startswith(('|','#')): 95 self.comment = self.comment[1:]
96 97 # Note that mother1/mother2 will be modified by the Event parse function to replace the 98 # integer by a pointer to the actual particle object. 99
100 - def __str__(self):
101 """string representing the particles""" 102 return " %8d %2d %4d %4d %4d %4d %+13.10e %+13.10e %+13.10e %14.10e %14.10e %10.4e %10.4e" \ 103 % (self.pid, 104 self.status, 105 (self.mother1 if isinstance(self.mother1, numbers.Number) else self.mother1.event_id+1) if self.mother1 else 0, 106 (self.mother2 if isinstance(self.mother2, numbers.Number) else self.mother2.event_id+1) if self.mother2 else 0, 107 self.color1, 108 self.color2, 109 self.px, 110 self.py, 111 self.pz, 112 self.E, 113 self.mass, 114 self.vtim, 115 self.helicity)
116
117 - def __eq__(self, other):
118 119 if not isinstance(other, Particle): 120 return False 121 if self.pid == other.pid and \ 122 self.status == other.status and \ 123 self.mother1 == other.mother1 and \ 124 self.mother2 == other.mother2 and \ 125 self.color1 == other.color1 and \ 126 self.color2 == other.color2 and \ 127 self.px == other.px and \ 128 self.py == other.py and \ 129 self.pz == other.pz and \ 130 self.E == other.E and \ 131 self.mass == other.mass and \ 132 self.vtim == other.vtim and \ 133 self.helicity == other.helicity: 134 return True 135 return False
136
137 - def set_momentum(self, momentum):
138 139 self.E = momentum.E 140 self.px = momentum.px 141 self.py = momentum.py 142 self.pz = momentum.pz
143
144 - def add_decay(self, decay_event):
145 """associate to this particle the decay in the associate event""" 146 147 return self.event.add_decay_to_particle(self.event_id, decay_event)
148
149 - def __repr__(self):
150 return 'Particle("%s", event=%s)' % (str(self), self.event)
151
152 153 -class EventFile(object):
154 """A class to allow to read both gzip and not gzip file""" 155 156
157 - def __new__(self, path, mode='r', *args, **opt):
158 159 if not path.endswith(".gz"): 160 return file.__new__(EventFileNoGzip, path, mode, *args, **opt) 161 elif mode == 'r' and not os.path.exists(path) and os.path.exists(path[:-3]): 162 return EventFile.__new__(EventFileNoGzip, path[:-3], mode, *args, **opt) 163 else: 164 try: 165 return gzip.GzipFile.__new__(EventFileGzip, path, mode, *args, **opt) 166 except IOError, error: 167 raise 168 except Exception, error: 169 if mode == 'r': 170 misc.gunzip(path) 171 return file.__new__(EventFileNoGzip, path[:-3], mode, *args, **opt)
172 173
174 - def __init__(self, path, mode='r', *args, **opt):
175 """open file and read the banner [if in read mode]""" 176 177 self.parsing = True # check if/when we need to parse the event. 178 self.eventgroup = False 179 try: 180 super(EventFile, self).__init__(path, mode, *args, **opt) 181 except IOError: 182 if '.gz' in path and isinstance(self, EventFileNoGzip) and\ 183 mode == 'r' and os.path.exists(path[:-3]): 184 super(EventFile, self).__init__(path[:-3], mode, *args, **opt) 185 else: 186 raise 187 188 self.banner = '' 189 if mode == 'r': 190 line = '' 191 while '</init>' not in line.lower(): 192 try: 193 line = super(EventFile, self).next() 194 except StopIteration: 195 self.seek(0) 196 self.banner = '' 197 break 198 if "<event" in line.lower(): 199 self.seek(0) 200 self.banner = '' 201 break 202 203 self.banner += line
204
205 - def get_banner(self):
206 """return a banner object""" 207 import madgraph.various.banner as banner 208 if isinstance(self.banner, banner.Banner): 209 return self.banner 210 211 output = banner.Banner() 212 output.read_banner(self.banner) 213 return output
214 215 @property
216 - def cross(self):
217 """return the cross-section of the file #from the banner""" 218 try: 219 return self._cross 220 except Exception: 221 pass 222 223 onebanner = self.get_banner() 224 self._cross = onebanner.get_cross() 225 return self._cross
226
227 - def __len__(self):
228 if self.closed: 229 return 0 230 if hasattr(self,"len"): 231 return self.len 232 233 init_pos = self.tell() 234 self.seek(0) 235 nb_event=0 236 with misc.TMP_variable(self, 'parsing', False): 237 for _ in self: 238 nb_event +=1 239 self.len = nb_event 240 self.seek(init_pos) 241 return self.len
242
243 - def next(self):
244 """get next event""" 245 246 if not self.eventgroup: 247 text = '' 248 line = '' 249 mode = 0 250 while '</event>' not in line: 251 line = super(EventFile, self).next() 252 if '<event' in line: 253 mode = 1 254 text = '' 255 if mode: 256 text += line 257 258 if self.parsing: 259 return Event(text) 260 else: 261 return text 262 else: 263 events = [] 264 text = '' 265 line = '' 266 mode = 0 267 while '</eventgroup>' not in line: 268 line = super(EventFile, self).next() 269 if '<eventgroup' in line: 270 events=[] 271 text = '' 272 elif '<event' in line: 273 text='' 274 mode=1 275 elif '</event>' in line: 276 if self.parsing: 277 events.append(Event(text)) 278 else: 279 events.append(text) 280 text = '' 281 mode = 0 282 if mode: 283 text += line 284 if len(events) == 0: 285 return self.next() 286 return events
287 288
289 - def initialize_unweighting(self, get_wgt, trunc_error):
290 """ scan once the file to return 291 - the list of the hightest weight (of size trunc_error*NB_EVENT 292 - the cross-section by type of process 293 - the total number of events in the file 294 """ 295 296 # We need to loop over the event file to get some information about the 297 # new cross-section/ wgt of event. 298 self.seek(0) 299 all_wgt = [] 300 cross = collections.defaultdict(int) 301 nb_event = 0 302 for event in self: 303 nb_event +=1 304 wgt = get_wgt(event) 305 cross['all'] += wgt 306 cross['abs'] += abs(wgt) 307 cross[event.ievent] += wgt 308 all_wgt.append(abs(wgt)) 309 # avoid all_wgt to be too large 310 if nb_event % 20000 == 0: 311 all_wgt.sort() 312 # drop the lowest weight 313 nb_keep = max(20, int(nb_event*trunc_error*15)) 314 all_wgt = all_wgt[-nb_keep:] 315 316 #final selection of the interesting weight to keep 317 all_wgt.sort() 318 # drop the lowest weight 319 nb_keep = max(20, int(nb_event*trunc_error*10)) 320 all_wgt = all_wgt[-nb_keep:] 321 self.seek(0) 322 return all_wgt, cross, nb_event
323
324 - def write_events(self, event):
325 """ write a single events or a list of event 326 if self.eventgroup is ON, then add <eventgroup> around the lists of events 327 """ 328 if isinstance(event, Event): 329 if self.eventgroup: 330 self.write('<eventgroup>\n%s\n</eventgroup>\n' % event) 331 else: 332 self.write(str(event)) 333 elif isinstance(event, list): 334 if self.eventgroup: 335 self.write('<eventgroup>\n') 336 for evt in event: 337 self.write(str(evt)) 338 if self.eventgroup: 339 self.write('</eventgroup>\n')
340
341 - def unweight(self, outputpath, get_wgt=None, max_wgt=0, trunc_error=0, 342 event_target=0, log_level=logging.INFO, normalization='average'):
343 """unweight the current file according to wgt information wgt. 344 which can either be a fct of the event or a tag in the rwgt list. 345 max_wgt allow to do partial unweighting. 346 trunc_error allow for dynamical partial unweighting 347 event_target reweight for that many event with maximal trunc_error. 348 (stop to write event when target is reached) 349 """ 350 if not get_wgt: 351 def weight(event): 352 return event.wgt
353 get_wgt = weight 354 unwgt_name = "central weight" 355 elif isinstance(get_wgt, str): 356 unwgt_name =get_wgt 357 def get_wgt(event): 358 event.parse_reweight() 359 return event.reweight_data[unwgt_name]
360 else: 361 unwgt_name = get_wgt.func_name 362 363 # check which weight to write 364 if hasattr(self, "written_weight"): 365 written_weight = lambda x: math.copysign(self.written_weight,float(x)) 366 else: 367 written_weight = lambda x: x 368 369 all_wgt, cross, nb_event = self.initialize_unweighting(get_wgt, trunc_error) 370 371 # function that need to be define on the flight 372 def max_wgt_for_trunc(trunc): 373 """find the weight with the maximal truncation.""" 374 375 xsum = 0 376 i=1 377 while (xsum - all_wgt[-i] * (i-1) <= cross['abs'] * trunc): 378 max_wgt = all_wgt[-i] 379 xsum += all_wgt[-i] 380 i +=1 381 if i == len(all_wgt): 382 break 383 384 return max_wgt 385 # end of the function 386 387 # choose the max_weight 388 if not max_wgt: 389 if trunc_error == 0 or len(all_wgt)<2 or event_target: 390 max_wgt = all_wgt[-1] 391 else: 392 max_wgt = max_wgt_for_trunc(trunc_error) 393 394 # need to modify the banner so load it to an object 395 if self.banner: 396 try: 397 import internal 398 except: 399 import madgraph.various.banner as banner_module 400 else: 401 import internal.banner as banner_module 402 if not isinstance(self.banner, banner_module.Banner): 403 banner = self.get_banner() 404 # 1. modify the cross-section 405 banner.modify_init_cross(cross) 406 # 3. add information about change in weight 407 banner["unweight"] = "unweighted by %s" % unwgt_name 408 else: 409 banner = self.banner 410 # modify the lha strategy 411 curr_strategy = banner.get_lha_strategy() 412 if normalization in ['unit', 'sum']: 413 strategy = 3 414 else: 415 strategy = 4 416 if curr_strategy >0: 417 banner.set_lha_strategy(abs(strategy)) 418 else: 419 banner.set_lha_strategy(-1*abs(strategy)) 420 421 # Do the reweighting (up to 20 times if we have target_event) 422 nb_try = 20 423 nb_keep = 0 424 for i in range(nb_try): 425 self.seek(0) 426 if event_target: 427 if i==0: 428 max_wgt = max_wgt_for_trunc(0) 429 else: 430 #guess the correct max_wgt based on last iteration 431 efficiency = nb_keep/nb_event 432 needed_efficiency = event_target/nb_event 433 last_max_wgt = max_wgt 434 needed_max_wgt = last_max_wgt * efficiency / needed_efficiency 435 436 min_max_wgt = max_wgt_for_trunc(trunc_error) 437 max_wgt = max(min_max_wgt, needed_max_wgt) 438 max_wgt = min(max_wgt, all_wgt[-1]) 439 if max_wgt == last_max_wgt: 440 if nb_keep <= event_target and log_level>=10: 441 logger.log(log_level+10,"fail to reach target %s", event_target) 442 break 443 else: 444 break 445 446 #create output file (here since we are sure that we have to rewrite it) 447 if outputpath: 448 outfile = EventFile(outputpath, "w") 449 # need to write banner information 450 # need to see what to do with rwgt information! 451 if self.banner and outputpath: 452 banner.write(outfile, close_tag=False) 453 454 # scan the file 455 nb_keep = 0 456 trunc_cross = 0 457 for event in self: 458 r = random.random() 459 wgt = get_wgt(event) 460 if abs(wgt) < r * max_wgt: 461 continue 462 elif wgt > 0: 463 nb_keep += 1 464 event.wgt = written_weight(max(wgt, max_wgt)) 465 if abs(wgt) > max_wgt: 466 trunc_cross += abs(wgt) - max_wgt 467 if event_target ==0 or nb_keep <= event_target: 468 if outputpath: 469 outfile.write(str(event)) 470 471 elif wgt < 0: 472 nb_keep += 1 473 event.wgt = -1* written_weight(max(abs(wgt), max_wgt)) 474 if abs(wgt) > max_wgt: 475 trunc_cross += abs(wgt) - max_wgt 476 if outputpath and (event_target ==0 or nb_keep <= event_target): 477 outfile.write(str(event)) 478 479 if event_target and nb_keep > event_target: 480 if not outputpath: 481 #no outputpath define -> wants only the nb of unweighted events 482 continue 483 elif event_target and i != nb_try-1 and nb_keep >= event_target *1.05: 484 outfile.write("</LesHouchesEvents>\n") 485 outfile.close() 486 #logger.log(log_level, "Found Too much event %s. Try to reduce truncation" % nb_keep) 487 continue 488 else: 489 outfile.write("</LesHouchesEvents>\n") 490 outfile.close() 491 break 492 elif event_target == 0: 493 if outputpath: 494 outfile.write("</LesHouchesEvents>\n") 495 outfile.close() 496 break 497 elif outputpath: 498 outfile.write("</LesHouchesEvents>\n") 499 outfile.close() 500 # logger.log(log_level, "Found only %s event. Reduce max_wgt" % nb_keep) 501 502 else: 503 # pass here if event_target > 0 and all the attempt fail. 504 logger.log(log_level+10,"fail to reach target event %s (iteration=%s)", event_target,i) 505 506 # logger.log(log_level, "Final maximum weight used for final "+\ 507 # "unweighting is %s yielding %s events." % (max_wgt,nb_keep)) 508 509 if event_target: 510 nb_events_unweighted = nb_keep 511 nb_keep = min( event_target, nb_keep) 512 else: 513 nb_events_unweighted = nb_keep 514 515 logger.log(log_level, "write %i event (efficiency %.2g %%, truncation %.2g %%) after %i iteration(s)", 516 nb_keep, nb_events_unweighted/nb_event*100, trunc_cross/cross['abs']*100, i) 517 518 #correct the weight in the file if not the correct number of event 519 if nb_keep != event_target and hasattr(self, "written_weight") and strategy !=4: 520 written_weight = lambda x: math.copysign(self.written_weight*event_target/nb_keep, float(x)) 521 startfile = EventFile(outputpath) 522 tmpname = pjoin(os.path.dirname(outputpath), "wgtcorrected_"+ os.path.basename(outputpath)) 523 outfile = EventFile(tmpname, "w") 524 outfile.write(startfile.banner) 525 for event in startfile: 526 event.wgt = written_weight(event.wgt) 527 outfile.write(str(event)) 528 outfile.write("</LesHouchesEvents>\n") 529 startfile.close() 530 outfile.close() 531 shutil.move(tmpname, outputpath) 532 533 534 535 536 self.max_wgt = max_wgt 537 return nb_keep 538
539 - def apply_fct_on_event(self, *fcts, **opts):
540 """ apply one or more fct on all event. """ 541 542 opt= {"print_step": 5000, "maxevent":float("inf"),'no_output':False} 543 opt.update(opts) 544 start = time.time() 545 nb_fct = len(fcts) 546 out = [] 547 for i in range(nb_fct): 548 out.append([]) 549 self.seek(0) 550 nb_event = 0 551 for event in self: 552 nb_event += 1 553 if opt["print_step"] and (nb_event % opt["print_step"]) == 0: 554 if hasattr(self,"len"): 555 print("currently at %s/%s event [%is]" % (nb_event, self.len, time.time()-start)) 556 else: 557 print("currently at %s event [%is]" % (nb_event, time.time()-start)) 558 for i in range(nb_fct): 559 value = fcts[i](event) 560 if not opt['no_output']: 561 out[i].append(value) 562 if nb_event > opt['maxevent']: 563 break 564 if nb_fct == 1: 565 return out[0] 566 else: 567 return out
568
569 - def split(self, nb_event=0, partition=None, cwd=os.path.curdir, zip=False):
570 """split the file in multiple file. Do not change the weight!""" 571 572 nb_file = -1 573 for i, event in enumerate(self): 574 if (not (partition is None) and i==sum(partition[:nb_file+1])) or \ 575 (partition is None and i % nb_event == 0): 576 if i: 577 #close previous file 578 current.write('</LesHouchesEvent>\n') 579 current.close() 580 # create the new file 581 nb_file +=1 582 # If end of partition then finish writing events here. 583 if not partition is None and (nb_file+1>len(partition)): 584 return nb_file 585 if zip: 586 current = EventFile(pjoin(cwd,'%s_%s.lhe.gz' % (self.name, nb_file)),'w') 587 else: 588 current = open(pjoin(cwd,'%s_%s.lhe' % (self.name, nb_file)),'w') 589 current.write(self.banner) 590 current.write(str(event)) 591 if i!=0: 592 current.write('</LesHouchesEvent>\n') 593 current.close() 594 return nb_file +1
595
596 - def update_HwU(self, hwu, fct, name='lhe', keep_wgt=False, maxevents=sys.maxint):
597 """take a HwU and add this event file for the function fct""" 598 599 if not isinstance(hwu, list): 600 hwu = [hwu] 601 602 class HwUUpdater(object): 603 604 def __init__(self, fct, keep_wgt): 605 606 self.fct = fct 607 self.first = True 608 self.keep_wgt = keep_wgt
609 610 def add(self, event): 611 612 value = self.fct(event) 613 # initialise the curve for the first call 614 if self.first: 615 for h in hwu: 616 # register the variables 617 if isinstance(value, dict): 618 h.add_line(value.keys()) 619 else: 620 621 h.add_line(name) 622 if self.keep_wgt is True: 623 event.parse_reweight() 624 h.add_line(['%s_%s' % (name, key) 625 for key in event.reweight_data]) 626 elif self.keep_wgt: 627 h.add_line(self.keep_wgt.values()) 628 self.first = False 629 # Fill the histograms 630 for h in hwu: 631 if isinstance(value, tuple): 632 h.addEvent(value[0], value[1]) 633 else: 634 h.addEvent(value,{name:event.wgt}) 635 if self.keep_wgt: 636 event.parse_reweight() 637 if self.keep_wgt is True: 638 data = dict(('%s_%s' % (name, key),event.reweight_data[key]) 639 for key in event.reweight_data) 640 h.addEvent(value, data) 641 else: 642 data = dict(( value,event.reweight_data[key]) 643 for key,value in self.keep_wgt.items()) 644 h.addEvent(value, data) 645 646 647 648 self.apply_fct_on_event(HwUUpdater(fct,keep_wgt).add, no_output=True,maxevent=maxevents) 649 return hwu 650
651 - def create_syscalc_data(self, out_path, pythia_input=None):
652 """take the lhe file and add the matchscale from the pythia_input file""" 653 654 if pythia_input: 655 def next_data(): 656 for line in open(pythia_input): 657 if line.startswith('#'): 658 continue 659 data = line.split() 660 print (int(data[0]), data[-3], data[-2], data[-1]) 661 yield (int(data[0]), data[-3], data[-2], data[-1])
662 else: 663 def next_data(): 664 i=0 665 while 1: 666 yield [i,0,0,0] 667 i+=1 668 sys_iterator = next_data() 669 #ensure that we are at the beginning of the file 670 self.seek(0) 671 out = open(out_path,'w') 672 673 pdf_pattern = re.compile(r'''<init>(.*)</init>''', re.M+re.S) 674 init = pdf_pattern.findall(self.banner)[0].split('\n',2)[1] 675 id1, id2, _, _, _, _, pdf1,pdf2,_,_ = init.split() 676 id = [int(id1), int(id2)] 677 type = [] 678 for i in range(2): 679 if abs(id[i]) == 2212: 680 if i > 0: 681 type.append(1) 682 else: 683 type.append(-1) 684 else: 685 type.append(0) 686 pdf = max(int(pdf1),int(pdf2)) 687 688 out.write("<header>\n" + \ 689 "<orgpdf>%i</orgpdf>\n" % pdf + \ 690 "<beams> %s %s</beams>\n" % tuple(type) + \ 691 "</header>\n") 692 693 694 nevt, smin, smax, scomp = sys_iterator.next() 695 for i, orig_event in enumerate(self): 696 if i < nevt: 697 continue 698 new_event = Event() 699 sys = orig_event.parse_syscalc_info() 700 new_event.syscalc_data = sys 701 if smin: 702 new_event.syscalc_data['matchscale'] = "%s %s %s" % (smin, scomp, smax) 703 out.write(str(new_event), nevt) 704 try: 705 nevt, smin, smax, scomp = sys_iterator.next() 706 except StopIteration: 707 break 708
709 710 711 712 713 714 715 716 -class EventFileGzip(EventFile, gzip.GzipFile):
717 """A way to read/write a gzipped lhef event"""
718
719 -class EventFileNoGzip(EventFile, file):
720 """A way to read a standard event file"""
721
722 -class MultiEventFile(EventFile):
723 """a class to read simultaneously multiple file and read them in mixing them. 724 Unweighting can be done at the same time. 725 The number of events in each file need to be provide in advance 726 (if not provide the file is first read to find that number""" 727
728 - def __new__(cls, start_list=[],parse=True):
729 return object.__new__(MultiEventFile)
730
731 - def __init__(self, start_list=[], parse=True):
732 """if trunc_error is define here then this allow 733 to only read all the files twice and not three times.""" 734 self.eventgroup = False 735 self.files = [] 736 self.parsefile = parse #if self.files is formatted or just the path 737 self.banner = '' 738 self.initial_nb_events = [] 739 self.total_event_in_files = 0 740 self.curr_nb_events = [] 741 self.allcross = [] 742 self.error = [] 743 self.across = [] 744 self.scales = [] 745 if start_list: 746 if parse: 747 for p in start_list: 748 self.add(p) 749 else: 750 self.files = start_list 751 self._configure = False
752
753 - def close(self,*args,**opts):
754 for f in self.files: 755 f.close(*args, **opts)
756
757 - def add(self, path, cross, error, across, nb_event=0, scale=1):
758 """ add a file to the pool, across allow to reweight the sum of weight 759 in the file to the given cross-section 760 """ 761 762 if across == 0: 763 # No event linked to this channel -> so no need to include it 764 return 765 766 obj = EventFile(path) 767 obj.eventgroup = self.eventgroup 768 if len(self.files) == 0 and not self.banner: 769 self.banner = obj.banner 770 self.curr_nb_events.append(0) 771 self.initial_nb_events.append(0) 772 self.allcross.append(cross) 773 self.across.append(across) 774 self.error.append(error) 775 self.scales.append(scale) 776 self.files.append(obj) 777 if nb_event: 778 obj.len = nb_event 779 self._configure = False 780 return obj
781
782 - def __iter__(self):
783 return self
784
785 - def next(self):
786 787 if not self._configure: 788 self.configure() 789 790 remaining_event = self.total_event_in_files - sum(self.curr_nb_events) 791 if remaining_event == 0: 792 raise StopIteration 793 # determine which file need to be read 794 nb_event = random.randint(1, remaining_event) 795 sum_nb=0 796 for i, obj in enumerate(self.files): 797 sum_nb += self.initial_nb_events[i] - self.curr_nb_events[i] 798 if nb_event <= sum_nb: 799 self.curr_nb_events[i] += 1 800 event = obj.next() 801 if not self.eventgroup: 802 event.sample_scale = self.scales[i] # for file reweighting 803 else: 804 for evt in event: 805 evt.sample_scale = self.scales[i] 806 return event 807 else: 808 raise Exception
809 810
811 - def define_init_banner(self, wgt, lha_strategy, proc_charac=None):
812 """define the part of the init_banner""" 813 814 if not self.banner: 815 return 816 817 # compute the cross-section of each splitted channel 818 grouped_cross = {} 819 grouped_error = {} 820 for i,ff in enumerate(self.files): 821 filename = ff.name 822 from_init = False 823 Pdir = [P for P in filename.split(os.path.sep) if P.startswith('P')] 824 if Pdir: 825 Pdir = Pdir[-1] 826 group = Pdir.split("_")[0][1:] 827 if not group.isdigit(): 828 from_init = True 829 else: 830 from_init = True 831 832 if not from_init: 833 if group in grouped_cross: 834 grouped_cross[group] += self.allcross[i] 835 grouped_error[group] += self.error[i]**2 836 else: 837 grouped_cross[group] = self.allcross[i] 838 grouped_error[group] = self.error[i]**2 839 else: 840 ban = banner_mod.Banner(ff.banner) 841 for line in ban['init'].split('\n'): 842 splitline = line.split() 843 if len(splitline)==4: 844 cross, error, _, group = splitline 845 if int(group) in grouped_cross: 846 grouped_cross[group] += float(cross) 847 grouped_error[group] += float(error)**2 848 else: 849 grouped_cross[group] = float(cross) 850 grouped_error[group] = float(error)**2 851 nb_group = len(grouped_cross) 852 853 # compute the information for the first line 854 try: 855 run_card = self.banner.run_card 856 except: 857 run_card = self.banner.charge_card("run_card") 858 859 init_information = run_card.get_banner_init_information() 860 #correct for special case 861 if proc_charac and proc_charac['ninitial'] == 1: 862 #special case for 1>N 863 init_information = run_card.get_banner_init_information() 864 event = self.next() 865 init_information["idbmup1"] = event[0].pdg 866 init_information["ebmup1"] = event[0].mass 867 init_information["idbmup2"] = 0 868 init_information["ebmup2"] = 0 869 self.seek(0) 870 else: 871 # check special case without PDF for one (or both) beam 872 if init_information["idbmup1"] == 0: 873 event = self.next() 874 init_information["idbmup1"]= event[0].pdg 875 if init_information["idbmup2"] == 0: 876 init_information["idbmup2"]= event[1].pdg 877 self.seek(0) 878 if init_information["idbmup2"] == 0: 879 event = self.next() 880 init_information["idbmup2"] = event[1].pdg 881 self.seek(0) 882 883 init_information["nprup"] = nb_group 884 885 if run_card["lhe_version"] < 3: 886 init_information["generator_info"] = "" 887 else: 888 init_information["generator_info"] = "<generator name='MadGraph5_aMC@NLO' version='%s'>please cite 1405.0301 </generator>\n" \ 889 % misc.get_pkg_info()['version'] 890 891 # cross_information: 892 cross_info = "%(cross)e %(error)e %(wgt)e %(id)i" 893 init_information["cross_info"] = [] 894 for id in grouped_cross: 895 conv = {"id": int(id), "cross": grouped_cross[id], "error": math.sqrt(grouped_error[id]), 896 "wgt": wgt} 897 init_information["cross_info"].append( cross_info % conv) 898 init_information["cross_info"] = '\n'.join(init_information["cross_info"]) 899 init_information['lha_stra'] = -1 * abs(lha_strategy) 900 901 template_init =\ 902 """ %(idbmup1)i %(idbmup2)i %(ebmup1)e %(ebmup2)e %(pdfgup1)i %(pdfgup2)i %(pdfsup1)i %(pdfsup2)i %(lha_stra)i %(nprup)i 903 %(cross_info)s 904 %(generator_info)s 905 """ 906 907 self.banner["init"] = template_init % init_information
908 909 910
911 - def initialize_unweighting(self, getwgt, trunc_error):
912 """ scan once the file to return 913 - the list of the hightest weight (of size trunc_error*NB_EVENT 914 - the cross-section by type of process 915 - the total number of events in the files 916 In top of that it initialise the information for the next routine 917 to determine how to choose which file to read 918 """ 919 self.seek(0) 920 all_wgt = [] 921 total_event = 0 922 sum_cross = collections.defaultdict(int) 923 for i,f in enumerate(self.files): 924 nb_event = 0 925 # We need to loop over the event file to get some information about the 926 # new cross-section/ wgt of event. 927 cross = collections.defaultdict(int) 928 new_wgt =[] 929 for event in f: 930 nb_event += 1 931 total_event += 1 932 event.sample_scale = 1 933 wgt = getwgt(event) 934 cross['all'] += wgt 935 cross['abs'] += abs(wgt) 936 cross[event.ievent] += wgt 937 new_wgt.append(abs(wgt)) 938 # avoid all_wgt to be too large 939 if nb_event % 20000 == 0: 940 new_wgt.sort() 941 # drop the lowest weight 942 nb_keep = max(20, int(nb_event*trunc_error*15)) 943 new_wgt = new_wgt[-nb_keep:] 944 if nb_event == 0: 945 raise Exception 946 # store the information 947 self.initial_nb_events[i] = nb_event 948 self.scales[i] = self.across[i]/cross['abs'] if self.across[i] else 1 949 #misc.sprint("sum of wgt in event %s is %s. Should be %s => scale %s (nb_event: %s)" 950 # % (i, cross['all'], self.allcross[i], self.scales[i], nb_event)) 951 for key in cross: 952 sum_cross[key] += cross[key]* self.scales[i] 953 all_wgt +=[self.scales[i] * w for w in new_wgt] 954 all_wgt.sort() 955 nb_keep = max(20, int(total_event*trunc_error*10)) 956 all_wgt = all_wgt[-nb_keep:] 957 958 self.total_event_in_files = total_event 959 #final selection of the interesting weight to keep 960 all_wgt.sort() 961 # drop the lowest weight 962 nb_keep = max(20, int(total_event*trunc_error*10)) 963 all_wgt = all_wgt[-nb_keep:] 964 self.seek(0) 965 self._configure = True 966 return all_wgt, sum_cross, total_event
967
968 - def configure(self):
969 970 self._configure = True 971 for i,f in enumerate(self.files): 972 self.initial_nb_events[i] = len(f) 973 self.total_event_in_files = sum(self.initial_nb_events)
974
975 - def __len__(self):
976 977 return len(self.files)
978
979 - def seek(self, pos):
980 """ """ 981 982 if pos !=0: 983 raise Exception 984 for i in range(len(self)): 985 self.curr_nb_events[i] = 0 986 for f in self.files: 987 f.seek(pos)
988
989 - def unweight(self, outputpath, get_wgt, **opts):
990 """unweight the current file according to wgt information wgt. 991 which can either be a fct of the event or a tag in the rwgt list. 992 max_wgt allow to do partial unweighting. 993 trunc_error allow for dynamical partial unweighting 994 event_target reweight for that many event with maximal trunc_error. 995 (stop to write event when target is reached) 996 """ 997 998 if isinstance(get_wgt, str): 999 unwgt_name =get_wgt 1000 def get_wgt_multi(event): 1001 event.parse_reweight() 1002 return event.reweight_data[unwgt_name] * event.sample_scale
1003 else: 1004 unwgt_name = get_wgt.func_name 1005 get_wgt_multi = lambda event: get_wgt(event) * event.sample_scale 1006 #define the weighting such that we have built-in the scaling 1007 1008 if 'proc_charac' in opts: 1009 if opts['proc_charac']: 1010 proc_charac = opts['proc_charac'] 1011 else: 1012 proc_charac=None 1013 del opts['proc_charac'] 1014 else: 1015 proc_charac = None 1016 1017 if 'event_target' in opts and opts['event_target']: 1018 if 'normalization' in opts: 1019 if opts['normalization'] == 'sum': 1020 new_wgt = sum(self.across)/opts['event_target'] 1021 strategy = 3 1022 elif opts['normalization'] == 'average': 1023 strategy = 4 1024 new_wgt = sum(self.across) 1025 elif opts['normalization'] == 'unit': 1026 strategy =3 1027 new_wgt = 1. 1028 else: 1029 strategy = 4 1030 new_wgt = sum(self.across) 1031 self.define_init_banner(new_wgt, strategy, proc_charac=proc_charac) 1032 self.written_weight = new_wgt 1033 elif 'write_init' in opts and opts['write_init']: 1034 self.define_init_banner(0,0, proc_charac=proc_charac) 1035 del opts['write_init'] 1036 return super(MultiEventFile, self).unweight(outputpath, get_wgt_multi, **opts)
1037
1038 - def write(self, path, random=False, banner=None, get_info=False):
1039 """ """ 1040 1041 if isinstance(path, str): 1042 out = EventFile(path, 'w') 1043 if self.parsefile and not banner: 1044 banner = self.files[0].banner 1045 elif not banner: 1046 firstlhe = EventFile(self.files[0]) 1047 banner = firstlhe.banner 1048 else: 1049 out = path 1050 if banner: 1051 out.write(banner) 1052 nb_event = 0 1053 info = collections.defaultdict(float) 1054 if random and self.open: 1055 for event in self: 1056 nb_event +=1 1057 out.write(event) 1058 if get_info: 1059 event.parse_reweight() 1060 for key, value in event.reweight_data.items(): 1061 info[key] += value 1062 info['central'] += event.wgt 1063 elif not random: 1064 for i,f in enumerate(self.files): 1065 #check if we need to parse the file or not 1066 if not self.parsefile: 1067 if i==0: 1068 try: 1069 lhe = firstlhe 1070 except: 1071 lhe = EventFile(f) 1072 else: 1073 lhe = EventFile(f) 1074 else: 1075 lhe = f 1076 for event in lhe: 1077 nb_event +=1 1078 if get_info: 1079 event.parse_reweight() 1080 for key, value in event.reweight_data.items(): 1081 info[key] += value 1082 info['central'] += event.wgt 1083 out.write(str(event)) 1084 lhe.close() 1085 out.write("</LesHouchesEvents>\n") 1086 return nb_event, info
1087
1088 - def remove(self):
1089 """ """ 1090 if self.parsefile: 1091 for f in self.files: 1092 os.remove(f.name) 1093 else: 1094 for f in self.files: 1095 os.remove(f)
1096
1097 1098 1099 -class Event(list):
1100 """Class storing a single event information (list of particles + global information)""" 1101 1102 warning_order = True # raise a warning if the order of the particle are not in accordance of child/mother 1103
1104 - def __init__(self, text=None):
1105 """The initialization of an empty Event (or one associate to a text file)""" 1106 list.__init__(self) 1107 1108 # First line information 1109 self.nexternal = 0 1110 self.ievent = 0 1111 self.wgt = 0 1112 self.aqcd = 0 1113 self.scale = 0 1114 self.aqed = 0 1115 self.aqcd = 0 1116 # Weight information 1117 self.tag = '' 1118 self.eventflag = {} # for information in <event > 1119 self.comment = '' 1120 self.reweight_data = {} 1121 self.matched_scale_data = None 1122 self.syscalc_data = {} 1123 if text: 1124 self.parse(text)
1125 1126 1127
1128 - def parse(self, text):
1129 """Take the input file and create the structured information""" 1130 #text = re.sub(r'</?event>', '', text) # remove pointless tag 1131 status = 'first' 1132 for line in text.split('\n'): 1133 line = line.strip() 1134 if not line: 1135 continue 1136 elif line[0] == '#': 1137 self.comment += '%s\n' % line 1138 continue 1139 elif line.startswith('<event'): 1140 if '=' in line: 1141 found = re.findall(r"""(\w*)=(?:(?:['"])([^'"]*)(?=['"])|(\S*))""",line) 1142 #for '<event line=4 value=\'3\' error="5" test=" 1 and 2">\n' 1143 #return [('line', '', '4'), ('value', '3', ''), ('error', '5', ''), ('test', ' 1 and 2', '')] 1144 self.eventflag = dict((n, a1) if a1 else (n,a2) for n,a1,a2 in found) 1145 # return {'test': ' 1 and 2', 'line': '4', 'value': '3', 'error': '5'} 1146 continue 1147 1148 elif 'first' == status: 1149 if '<rwgt>' in line: 1150 status = 'tag' 1151 else: 1152 self.assign_scale_line(line) 1153 status = 'part' 1154 continue 1155 if '<' in line: 1156 status = 'tag' 1157 1158 if 'part' == status: 1159 part = Particle(line, event=self) 1160 if part.E != 0: 1161 self.append(part) 1162 elif self.nexternal: 1163 self.nexternal-=1 1164 else: 1165 if '</event>' in line: 1166 line = line.replace('</event>','',1) 1167 self.tag += '%s\n' % line 1168 1169 self.assign_mother()
1170
1171 - def assign_mother(self):
1172 # assign the mother: 1173 for i,particle in enumerate(self): 1174 if i < particle.mother1 or i < particle.mother2: 1175 if self.warning_order: 1176 logger.warning("Order of particle in the event did not agree with parent/child order. This might be problematic for some code.") 1177 Event.warning_order = False 1178 self.reorder_mother_child() 1179 return self.assign_mother() 1180 1181 if particle.mother1: 1182 try: 1183 particle.mother1 = self[int(particle.mother1) -1] 1184 except Exception: 1185 logger.warning("WRONG MOTHER INFO %s", self) 1186 particle.mother1 = 0 1187 if particle.mother2: 1188 try: 1189 particle.mother2 = self[int(particle.mother2) -1] 1190 except Exception: 1191 logger.warning("WRONG MOTHER INFO %s", self) 1192 particle.mother2 = 0
1193
1194 - def rescale_weights(self, ratio):
1195 """change all the weights by a given ratio""" 1196 1197 self.wgt *= ratio 1198 self.parse_reweight() 1199 for key in self.reweight_data: 1200 self.reweight_data[key] *= ratio 1201 return self.wgt
1202
1203 - def reorder_mother_child(self):
1204 """check and correct the mother/child position. 1205 only correct one order by call (but this is a recursive call)""" 1206 1207 tomove, position = None, None 1208 for i,particle in enumerate(self): 1209 if i < particle.mother1: 1210 # move i after particle.mother1 1211 tomove, position = i, particle.mother1-1 1212 break 1213 if i < particle.mother2: 1214 tomove, position = i, particle.mother2-1 1215 1216 # nothing to change -> we are done 1217 if not tomove: 1218 return 1219 1220 # move the particles: 1221 particle = self.pop(tomove) 1222 self.insert(int(position), particle) 1223 1224 #change the mother id/ event_id in the event. 1225 for i, particle in enumerate(self): 1226 particle.event_id = i 1227 #misc.sprint( i, particle.event_id) 1228 m1, m2 = particle.mother1, particle.mother2 1229 if m1 == tomove +1: 1230 particle.mother1 = position+1 1231 elif tomove < m1 <= position +1: 1232 particle.mother1 -= 1 1233 if m2 == tomove +1: 1234 particle.mother2 = position+1 1235 elif tomove < m2 <= position +1: 1236 particle.mother2 -= 1 1237 # re-call the function for the next potential change 1238 return self.reorder_mother_child()
1239 1240 1241 1242 1243 1244
1245 - def parse_reweight(self):
1246 """Parse the re-weight information in order to return a dictionary 1247 {key: value}. If no group is define group should be '' """ 1248 if self.reweight_data: 1249 return self.reweight_data 1250 self.reweight_data = {} 1251 self.reweight_order = [] 1252 start, stop = self.tag.find('<rwgt>'), self.tag.find('</rwgt>') 1253 if start != -1 != stop : 1254 pattern = re.compile(r'''<\s*wgt id=(?:\'|\")(?P<id>[^\'\"]+)(?:\'|\")\s*>\s*(?P<val>[\ded+-.]*)\s*</wgt>''',re.I) 1255 data = pattern.findall(self.tag[start:stop]) 1256 try: 1257 self.reweight_data = dict([(pid, float(value)) for (pid, value) in data 1258 if not self.reweight_order.append(pid)]) 1259 # the if is to create the order file on the flight 1260 except ValueError, error: 1261 raise Exception, 'Event File has unvalid weight. %s' % error 1262 self.tag = self.tag[:start] + self.tag[stop+7:] 1263 return self.reweight_data
1264
1265 - def parse_nlo_weight(self, real_type=(1,11)):
1266 """ """ 1267 if hasattr(self, 'nloweight'): 1268 return self.nloweight 1269 1270 start, stop = self.tag.find('<mgrwgt>'), self.tag.find('</mgrwgt>') 1271 if start != -1 != stop : 1272 1273 text = self.tag[start+8:stop] 1274 self.nloweight = NLO_PARTIALWEIGHT(text, self, real_type=real_type) 1275 return self.nloweight
1276
1277 - def parse_lo_weight(self):
1278 """ """ 1279 1280 1281 if hasattr(self, 'loweight'): 1282 return self.loweight 1283 1284 if not hasattr(Event, 'loweight_pattern'): 1285 Event.loweight_pattern = re.compile('''<rscale>\s*(?P<nqcd>\d+)\s+(?P<ren_scale>[\d.e+-]+)\s*</rscale>\s*\n\s* 1286 <asrwt>\s*(?P<asrwt>[\s\d.+-e]+)\s*</asrwt>\s*\n\s* 1287 <pdfrwt\s+beam=["']?1["']?\>\s*(?P<beam1>[\s\d.e+-]*)\s*</pdfrwt>\s*\n\s* 1288 <pdfrwt\s+beam=["']?2["']?\>\s*(?P<beam2>[\s\d.e+-]*)\s*</pdfrwt>\s*\n\s* 1289 <totfact>\s*(?P<totfact>[\d.e+-]*)\s*</totfact> 1290 ''',re.X+re.I+re.M) 1291 1292 start, stop = self.tag.find('<mgrwt>'), self.tag.find('</mgrwt>') 1293 1294 if start != -1 != stop : 1295 text = self.tag[start+8:stop] 1296 1297 info = Event.loweight_pattern.search(text) 1298 if not info: 1299 raise Exception, '%s not parsed'% text 1300 self.loweight={} 1301 self.loweight['n_qcd'] = int(info.group('nqcd')) 1302 self.loweight['ren_scale'] = float(info.group('ren_scale')) 1303 self.loweight['asrwt'] =[float(x) for x in info.group('asrwt').split()[1:]] 1304 self.loweight['tot_fact'] = float(info.group('totfact')) 1305 1306 args = info.group('beam1').split() 1307 npdf = int(args[0]) 1308 self.loweight['n_pdfrw1'] = npdf 1309 self.loweight['pdf_pdg_code1'] = [int(i) for i in args[1:1+npdf]] 1310 self.loweight['pdf_x1'] = [float(i) for i in args[1+npdf:1+2*npdf]] 1311 self.loweight['pdf_q1'] = [float(i) for i in args[1+2*npdf:1+3*npdf]] 1312 args = info.group('beam2').split() 1313 npdf = int(args[0]) 1314 self.loweight['n_pdfrw2'] = npdf 1315 self.loweight['pdf_pdg_code2'] = [int(i) for i in args[1:1+npdf]] 1316 self.loweight['pdf_x2'] = [float(i) for i in args[1+npdf:1+2*npdf]] 1317 self.loweight['pdf_q2'] = [float(i) for i in args[1+2*npdf:1+3*npdf]] 1318 1319 else: 1320 return None 1321 return self.loweight
1322 1323
1324 - def parse_matching_scale(self):
1325 """Parse the line containing the starting scale for the shower""" 1326 1327 if self.matched_scale_data is not None: 1328 return self.matched_scale_data 1329 1330 self.matched_scale_data = [] 1331 1332 1333 pattern = re.compile("<scales\s|</scales>") 1334 data = re.split(pattern,self.tag) 1335 if len(data) == 1: 1336 return [] 1337 else: 1338 tmp = {} 1339 start,content, end = data 1340 self.tag = "%s%s" % (start, end) 1341 pattern = re.compile("pt_clust_(\d*)=\"([\de+-.]*)\"") 1342 for id,value in pattern.findall(content): 1343 tmp[int(id)] = float(value) 1344 for i in range(1, len(self)+1): 1345 if i in tmp: 1346 self.matched_scale_data.append(tmp[i]) 1347 else: 1348 self.matched_scale_data.append(-1) 1349 return self.matched_scale_data
1350
1351 - def parse_syscalc_info(self):
1352 """ parse the flag for syscalc between <mgrwt></mgrwt> 1353 <mgrwt> 1354 <rscale> 3 0.26552898E+03</rscale> 1355 <asrwt>0</asrwt> 1356 <pdfrwt beam="1"> 1 21 0.14527945E+00 0.26552898E+03</pdfrwt> 1357 <pdfrwt beam="2"> 1 21 0.15249110E-01 0.26552898E+03</pdfrwt> 1358 <totfact> 0.10344054E+04</totfact> 1359 </mgrwt> 1360 """ 1361 if self.syscalc_data: 1362 return self.syscalc_data 1363 1364 pattern = re.compile("<mgrwt>|</mgrwt>") 1365 pattern2 = re.compile("<(?P<tag>[\w]*)(?:\s*(\w*)=[\"'](.*)[\"']\s*|\s*)>(.*)</(?P=tag)>") 1366 data = re.split(pattern,self.tag) 1367 if len(data) == 1: 1368 return [] 1369 else: 1370 tmp = {} 1371 start,content, end = data 1372 self.tag = "%s%s" % (start, end) 1373 for tag, key, keyval, tagval in pattern2.findall(content): 1374 if key: 1375 self.syscalc_data[(tag, key, keyval)] = tagval 1376 else: 1377 self.syscalc_data[tag] = tagval 1378 return self.syscalc_data
1379 1380
1381 - def add_decay_to_particle(self, position, decay_event):
1382 """define the decay of the particle id by the event pass in argument""" 1383 1384 this_particle = self[position] 1385 #change the status to internal particle 1386 this_particle.status = 2 1387 this_particle.helicity = 0 1388 1389 # some usefull information 1390 decay_particle = decay_event[0] 1391 this_4mom = FourMomentum(this_particle) 1392 nb_part = len(self) #original number of particle 1393 1394 thres = decay_particle.E*1e-10 1395 assert max(decay_particle.px, decay_particle.py, decay_particle.pz) < thres,\ 1396 "not on rest particle %s %s %s %s" % (decay_particle.E, decay_particle.px,decay_particle.py,decay_particle.pz) 1397 1398 self.nexternal += decay_event.nexternal -1 1399 old_scales = list(self.parse_matching_scale()) 1400 if old_scales: 1401 jet_position = sum(1 for i in range(position) if self[i].status==1) 1402 initial_pos = sum(1 for i in range(position) if self[i].status==-1) 1403 self.matched_scale_data.pop(initial_pos+jet_position) 1404 # add the particle with only handling the 4-momenta/mother 1405 # color information will be corrected later. 1406 for particle in decay_event[1:]: 1407 # duplicate particle to avoid border effect 1408 new_particle = Particle(particle, self) 1409 new_particle.event_id = len(self) 1410 self.append(new_particle) 1411 if old_scales: 1412 self.matched_scale_data.append(old_scales[initial_pos+jet_position]) 1413 # compute and assign the new four_momenta 1414 new_momentum = this_4mom.boost(FourMomentum(new_particle)) 1415 new_particle.set_momentum(new_momentum) 1416 # compute the new mother 1417 for tag in ['mother1', 'mother2']: 1418 mother = getattr(particle, tag) 1419 if isinstance(mother, Particle): 1420 mother_id = getattr(particle, tag).event_id 1421 if mother_id == 0: 1422 setattr(new_particle, tag, this_particle) 1423 else: 1424 try: 1425 setattr(new_particle, tag, self[nb_part + mother_id -1]) 1426 except Exception, error: 1427 print error 1428 misc.sprint( self) 1429 misc.sprint(nb_part + mother_id -1) 1430 misc.sprint(tag) 1431 misc.sprint(position, decay_event) 1432 misc.sprint(particle) 1433 misc.sprint(len(self), nb_part + mother_id -1) 1434 raise 1435 elif tag == "mother2" and isinstance(particle.mother1, Particle): 1436 new_particle.mother2 = this_particle 1437 else: 1438 raise Exception, "Something weird happens. Please report it for investigation" 1439 # Need to correct the color information of the particle 1440 # first find the first available color index 1441 max_color=501 1442 for particle in self[:nb_part]: 1443 max_color=max(max_color, particle.color1, particle.color2) 1444 1445 # define a color mapping and assign it: 1446 color_mapping = {} 1447 color_mapping[decay_particle.color1] = this_particle.color1 1448 color_mapping[decay_particle.color2] = this_particle.color2 1449 for particle in self[nb_part:]: 1450 if particle.color1: 1451 if particle.color1 not in color_mapping: 1452 max_color +=1 1453 color_mapping[particle.color1] = max_color 1454 particle.color1 = max_color 1455 else: 1456 particle.color1 = color_mapping[particle.color1] 1457 if particle.color2: 1458 if particle.color2 not in color_mapping: 1459 max_color +=1 1460 color_mapping[particle.color2] = max_color 1461 particle.color2 = max_color 1462 else: 1463 particle.color2 = color_mapping[particle.color2]
1464
1465 - def add_decays(self, pdg_to_decay):
1466 """use auto-recursion""" 1467 1468 pdg_to_decay = dict(pdg_to_decay) 1469 1470 for i,particle in enumerate(self): 1471 if particle.status != 1: 1472 continue 1473 if particle.pdg in pdg_to_decay and pdg_to_decay[particle.pdg]: 1474 one_decay = pdg_to_decay[particle.pdg].pop() 1475 self.add_decay_to_particle(i, one_decay) 1476 return self.add_decays(pdg_to_decay) 1477 return self
1478 1479 1480
1481 - def remove_decay(self, pdg_code=0, event_id=None):
1482 1483 to_remove = [] 1484 if event_id is not None: 1485 to_remove.append(self[event_id]) 1486 1487 if pdg_code: 1488 for particle in self: 1489 if particle.pid == pdg_code: 1490 to_remove.append(particle) 1491 1492 new_event = Event() 1493 # copy first line information + ... 1494 for tag in ['nexternal', 'ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']: 1495 setattr(new_event, tag, getattr(self, tag)) 1496 1497 for particle in self: 1498 if isinstance(particle.mother1, Particle) and particle.mother1 in to_remove: 1499 to_remove.append(particle) 1500 if particle.status == 1: 1501 new_event.nexternal -= 1 1502 continue 1503 elif isinstance(particle.mother2, Particle) and particle.mother2 in to_remove: 1504 to_remove.append(particle) 1505 if particle.status == 1: 1506 new_event.nexternal -= 1 1507 continue 1508 else: 1509 new_event.append(Particle(particle)) 1510 1511 #ensure that the event_id is correct for all_particle 1512 # and put the status to 1 for removed particle 1513 for pos, particle in enumerate(new_event): 1514 particle.event_id = pos 1515 if particle in to_remove: 1516 particle.status = 1 1517 return new_event
1518
1519 - def get_decay(self, pdg_code=0, event_id=None):
1520 1521 to_start = [] 1522 if event_id is not None: 1523 to_start.append(self[event_id]) 1524 1525 elif pdg_code: 1526 for particle in self: 1527 if particle.pid == pdg_code: 1528 to_start.append(particle) 1529 break 1530 1531 new_event = Event() 1532 # copy first line information + ... 1533 for tag in ['ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']: 1534 setattr(new_event, tag, getattr(self, tag)) 1535 1536 # Add the decaying particle 1537 old2new = {} 1538 new_decay_part = Particle(to_start[0]) 1539 new_decay_part.mother1 = None 1540 new_decay_part.mother2 = None 1541 new_decay_part.status = -1 1542 old2new[new_decay_part.event_id] = len(old2new) 1543 new_event.append(new_decay_part) 1544 1545 1546 # add the other particle 1547 for particle in self: 1548 if isinstance(particle.mother1, Particle) and particle.mother1.event_id in old2new\ 1549 or isinstance(particle.mother2, Particle) and particle.mother2.event_id in old2new: 1550 old2new[particle.event_id] = len(old2new) 1551 new_event.append(Particle(particle)) 1552 1553 #ensure that the event_id is correct for all_particle 1554 # and correct the mother1/mother2 by the new reference 1555 nexternal = 0 1556 for pos, particle in enumerate(new_event): 1557 particle.event_id = pos 1558 if particle.mother1: 1559 particle.mother1 = new_event[old2new[particle.mother1.event_id]] 1560 if particle.mother2: 1561 particle.mother2 = new_event[old2new[particle.mother2.event_id]] 1562 if particle.status in [-1,1]: 1563 nexternal +=1 1564 new_event.nexternal = nexternal 1565 1566 return new_event
1567 1568
1569 - def check(self):
1570 """check various property of the events""" 1571 1572 #1. Check that the 4-momenta are conserved 1573 E, px, py, pz = 0,0,0,0 1574 absE, abspx, abspy, abspz = 0,0,0,0 1575 for particle in self: 1576 coeff = 1 1577 if particle.status == -1: 1578 coeff = -1 1579 elif particle.status != 1: 1580 continue 1581 E += coeff * particle.E 1582 absE += abs(particle.E) 1583 px += coeff * particle.px 1584 py += coeff * particle.py 1585 pz += coeff * particle.pz 1586 abspx += abs(particle.px) 1587 abspy += abs(particle.py) 1588 abspz += abs(particle.pz) 1589 # check that relative error is under control 1590 threshold = 5e-7 1591 if E/absE > threshold: 1592 logger.critical(self) 1593 raise Exception, "Do not conserve Energy %s, %s" % (E/absE, E) 1594 if px/abspx > threshold: 1595 logger.critical(self) 1596 raise Exception, "Do not conserve Px %s, %s" % (px/abspx, px) 1597 if py/abspy > threshold: 1598 logger.critical(self) 1599 raise Exception, "Do not conserve Py %s, %s" % (py/abspy, py) 1600 if pz/abspz > threshold: 1601 logger.critical(self) 1602 raise Exception, "Do not conserve Pz %s, %s" % (pz/abspz, pz) 1603 1604 #2. check the color of the event 1605 self.check_color_structure()
1606
1607 - def assign_scale_line(self, line):
1608 """read the line corresponding to global event line 1609 format of the line is: 1610 Nexternal IEVENT WEIGHT SCALE AEW AS 1611 """ 1612 inputs = line.split() 1613 assert len(inputs) == 6 1614 self.nexternal=int(inputs[0]) 1615 self.ievent=int(inputs[1]) 1616 self.wgt=float(inputs[2]) 1617 self.scale=float(inputs[3]) 1618 self.aqed=float(inputs[4]) 1619 self.aqcd=float(inputs[5])
1620
1621 - def get_tag_and_order(self):
1622 """Return the unique tag identifying the SubProcesses for the generation. 1623 Usefull for program like MadSpin and Reweight module.""" 1624 1625 initial, final, order = [], [], [[], []] 1626 for particle in self: 1627 if particle.status == -1: 1628 initial.append(particle.pid) 1629 order[0].append(particle.pid) 1630 elif particle.status == 1: 1631 final.append(particle.pid) 1632 order[1].append(particle.pid) 1633 initial.sort(), final.sort() 1634 tag = (tuple(initial), tuple(final)) 1635 return tag, order
1636
1637 - def get_helicity(self, get_order, allow_reversed=True):
1638 """return a list with the helicities in the order asked for""" 1639 1640 #avoid to modify the input 1641 order = [list(get_order[0]), list(get_order[1])] 1642 out = [9] *(len(order[0])+len(order[1])) 1643 for i, part in enumerate(self): 1644 if part.status == 1: #final 1645 try: 1646 ind = order[1].index(part.pid) 1647 except ValueError, error: 1648 if not allow_reversed: 1649 raise error 1650 else: 1651 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 1652 try: 1653 return self.get_helicity(order, False) 1654 except ValueError: 1655 raise error 1656 position = len(order[0]) + ind 1657 order[1][ind] = 0 1658 elif part.status == -1: 1659 try: 1660 ind = order[0].index(part.pid) 1661 except ValueError, error: 1662 if not allow_reversed: 1663 raise error 1664 else: 1665 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 1666 try: 1667 return self.get_helicity(order, False) 1668 except ValueError: 1669 raise error 1670 1671 position = ind 1672 order[0][ind] = 0 1673 else: #intermediate 1674 continue 1675 out[position] = int(part.helicity) 1676 return out
1677 1678
1679 - def check_color_structure(self):
1680 """check the validity of the color structure""" 1681 1682 #1. check that each color is raised only once. 1683 color_index = collections.defaultdict(int) 1684 for particle in self: 1685 if particle.status in [-1,1]: 1686 if particle.color1: 1687 color_index[particle.color1] +=1 1688 if -7 < particle.pdg < 0: 1689 raise Exception, "anti-quark with color tag" 1690 if particle.color2: 1691 color_index[particle.color2] +=1 1692 if 7 > particle.pdg > 0: 1693 raise Exception, "quark with anti-color tag" 1694 1695 1696 for key,value in color_index.items(): 1697 if value > 2: 1698 print self 1699 print key, value 1700 raise Exception, 'Wrong color_flow' 1701 1702 1703 #2. check that each parent present have coherent color-structure 1704 check = [] 1705 popup_index = [] #check that the popup index are created in a unique way 1706 for particle in self: 1707 mothers = [] 1708 childs = [] 1709 if particle.mother1: 1710 mothers.append(particle.mother1) 1711 if particle.mother2 and particle.mother2 is not particle.mother1: 1712 mothers.append(particle.mother2) 1713 if not mothers: 1714 continue 1715 if (particle.mother1.event_id, particle.mother2.event_id) in check: 1716 continue 1717 check.append((particle.mother1.event_id, particle.mother2.event_id)) 1718 1719 childs = [p for p in self if p.mother1 is particle.mother1 and \ 1720 p.mother2 is particle.mother2] 1721 1722 mcolors = [] 1723 manticolors = [] 1724 for m in mothers: 1725 if m.color1: 1726 if m.color1 in manticolors: 1727 manticolors.remove(m.color1) 1728 else: 1729 mcolors.append(m.color1) 1730 if m.color2: 1731 if m.color2 in mcolors: 1732 mcolors.remove(m.color2) 1733 else: 1734 manticolors.append(m.color2) 1735 ccolors = [] 1736 canticolors = [] 1737 for m in childs: 1738 if m.color1: 1739 if m.color1 in canticolors: 1740 canticolors.remove(m.color1) 1741 else: 1742 ccolors.append(m.color1) 1743 if m.color2: 1744 if m.color2 in ccolors: 1745 ccolors.remove(m.color2) 1746 else: 1747 canticolors.append(m.color2) 1748 for index in mcolors[:]: 1749 if index in ccolors: 1750 mcolors.remove(index) 1751 ccolors.remove(index) 1752 for index in manticolors[:]: 1753 if index in canticolors: 1754 manticolors.remove(index) 1755 canticolors.remove(index) 1756 1757 if mcolors != []: 1758 #only case is a epsilon_ijk structure. 1759 if len(canticolors) + len(mcolors) != 3: 1760 logger.critical(str(self)) 1761 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs]) 1762 else: 1763 popup_index += canticolors 1764 elif manticolors != []: 1765 #only case is a epsilon_ijk structure. 1766 if len(ccolors) + len(manticolors) != 3: 1767 logger.critical(str(self)) 1768 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs]) 1769 else: 1770 popup_index += ccolors 1771 1772 # Check that color popup (from epsilon_ijk) are raised only once 1773 if len(popup_index) != len(set(popup_index)): 1774 logger.critical(self) 1775 raise Exception, "Wrong color flow: identical poping-up index, %s" % (popup_index)
1776
1777 - def __eq__(self, other):
1778 """two event are the same if they have the same momentum. other info are ignored""" 1779 1780 if other is None: 1781 return False 1782 1783 for i,p in enumerate(self): 1784 if p.E != other[i].E: 1785 return False 1786 elif p.pz != other[i].pz: 1787 return False 1788 elif p.px != other[i].px: 1789 return False 1790 elif p.py != other[i].py: 1791 return False 1792 return True
1793 1794
1795 - def __str__(self, event_id=''):
1796 """return a correctly formatted LHE event""" 1797 1798 out="""<event%(event_flag)s> 1799 %(scale)s 1800 %(particles)s 1801 %(comments)s 1802 %(tag)s 1803 %(reweight)s 1804 </event> 1805 """ 1806 if event_id not in ['', None]: 1807 self.eventflag['event'] = str(event_id) 1808 1809 if self.eventflag: 1810 event_flag = ' %s' % ' '.join('%s="%s"' % (k,v) for (k,v) in self.eventflag.items()) 1811 else: 1812 event_flag = '' 1813 1814 if self.nexternal: 1815 scale_str = "%2d %6d %+13.7e %14.8e %14.8e %14.8e" % \ 1816 (self.nexternal,self.ievent,self.wgt,self.scale,self.aqed,self.aqcd) 1817 else: 1818 scale_str = '' 1819 1820 if self.reweight_data: 1821 # check that all key have an order if not add them at the end 1822 if set(self.reweight_data.keys()) != set(self.reweight_order): 1823 self.reweight_order += [k for k in self.reweight_data.keys() \ 1824 if k not in self.reweight_order] 1825 1826 reweight_str = '<rwgt>\n%s\n</rwgt>' % '\n'.join( 1827 '<wgt id=\'%s\'> %+13.7e </wgt>' % (i, float(self.reweight_data[i])) 1828 for i in self.reweight_order) 1829 else: 1830 reweight_str = '' 1831 1832 tag_str = self.tag 1833 if self.matched_scale_data: 1834 tmp_scale = ' '.join(['pt_clust_%i=\"%s\"' % (i+1,v) 1835 for i,v in enumerate(self.matched_scale_data) 1836 if v!=-1]) 1837 if tmp_scale: 1838 tag_str = "<scales %s></scales>%s" % (tmp_scale, self.tag) 1839 1840 if self.syscalc_data: 1841 keys= ['rscale', 'asrwt', ('pdfrwt', 'beam', '1'), ('pdfrwt', 'beam', '2'), 1842 'matchscale', 'totfact'] 1843 sys_str = "<mgrwt>\n" 1844 template = """<%(key)s%(opts)s>%(values)s</%(key)s>\n""" 1845 for k in keys: 1846 if k not in self.syscalc_data: 1847 continue 1848 replace = {} 1849 replace['values'] = self.syscalc_data[k] 1850 if isinstance(k, str): 1851 replace['key'] = k 1852 replace['opts'] = '' 1853 else: 1854 replace['key'] = k[0] 1855 replace['opts'] = ' %s=\"%s\"' % (k[1],k[2]) 1856 sys_str += template % replace 1857 sys_str += "</mgrwt>\n" 1858 reweight_str = sys_str + reweight_str 1859 1860 out = out % {'event_flag': event_flag, 1861 'scale': scale_str, 1862 'particles': '\n'.join([str(p) for p in self]), 1863 'tag': tag_str, 1864 'comments': self.comment, 1865 'reweight': reweight_str} 1866 1867 return re.sub('[\n]+', '\n', out)
1868
1869 - def get_momenta(self, get_order, allow_reversed=True):
1870 """return the momenta vector in the order asked for""" 1871 1872 #avoid to modify the input 1873 order = [list(get_order[0]), list(get_order[1])] 1874 out = [''] *(len(order[0])+len(order[1])) 1875 for i, part in enumerate(self): 1876 if part.status == 1: #final 1877 try: 1878 ind = order[1].index(part.pid) 1879 except ValueError, error: 1880 if not allow_reversed: 1881 raise error 1882 else: 1883 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 1884 try: 1885 return self.get_momenta_str(order, False) 1886 except ValueError: 1887 raise error 1888 position = len(order[0]) + ind 1889 order[1][ind] = 0 1890 elif part.status == -1: 1891 try: 1892 ind = order[0].index(part.pid) 1893 except ValueError, error: 1894 if not allow_reversed: 1895 raise error 1896 else: 1897 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 1898 try: 1899 return self.get_momenta_str(order, False) 1900 except ValueError: 1901 raise error 1902 1903 position = ind 1904 order[0][ind] = 0 1905 else: #intermediate 1906 continue 1907 1908 out[position] = (part.E, part.px, part.py, part.pz) 1909 1910 return out
1911 1912 1913
1914 - def get_ht_scale(self, prefactor=1):
1915 1916 scale = 0 1917 for particle in self: 1918 if particle.status != 1: 1919 continue 1920 p=FourMomentum(particle) 1921 scale += math.sqrt(p.mass_sqr + p.pt**2) 1922 1923 return prefactor * scale
1924 1925
1926 - def get_et_scale(self, prefactor=1):
1927 1928 scale = 0 1929 for particle in self: 1930 if particle.status != 1: 1931 continue 1932 p = FourMomentum(particle) 1933 pt = p.pt 1934 if (pt>0): 1935 scale += p.E*pt/math.sqrt(pt**2+p.pz**2) 1936 1937 return prefactor * scale
1938
1939 - def get_sqrts_scale(self, prefactor=1):
1940 1941 scale = 0 1942 init = [] 1943 for particle in self: 1944 if particle.status == -1: 1945 init.append(FourMomentum(particle)) 1946 if len(init) == 1: 1947 return init[0].mass 1948 elif len(init)==2: 1949 return math.sqrt((init[0]+init[1])**2)
1950 1951 1952 1953
1954 - def get_momenta_str(self, get_order, allow_reversed=True):
1955 """return the momenta str in the order asked for""" 1956 1957 out = self.get_momenta(get_order, allow_reversed) 1958 #format 1959 format = '%.12f' 1960 format_line = ' '.join([format]*4) + ' \n' 1961 out = [format_line % one for one in out] 1962 out = ''.join(out).replace('e','d') 1963 return out
1964
1965 -class WeightFile(EventFile):
1966 """A class to allow to read both gzip and not gzip file. 1967 containing only weight from pythia --generated by SysCalc""" 1968
1969 - def __new__(self, path, mode='r', *args, **opt):
1970 if path.endswith(".gz"): 1971 try: 1972 return gzip.GzipFile.__new__(WeightFileGzip, path, mode, *args, **opt) 1973 except IOError, error: 1974 raise 1975 except Exception, error: 1976 if mode == 'r': 1977 misc.gunzip(path) 1978 return file.__new__(WeightFileNoGzip, path[:-3], mode, *args, **opt) 1979 else: 1980 return file.__new__(WeightFileNoGzip, path, mode, *args, **opt)
1981 1982
1983 - def __init__(self, path, mode='r', *args, **opt):
1984 """open file and read the banner [if in read mode]""" 1985 1986 super(EventFile, self).__init__(path, mode, *args, **opt) 1987 self.banner = '' 1988 if mode == 'r': 1989 line = '' 1990 while '</header>' not in line.lower(): 1991 try: 1992 line = super(EventFile, self).next() 1993 except StopIteration: 1994 self.seek(0) 1995 self.banner = '' 1996 break 1997 if "<event" in line.lower(): 1998 self.seek(0) 1999 self.banner = '' 2000 break 2001 2002 self.banner += line
2003
2004 2005 -class WeightFileGzip(WeightFile, EventFileGzip):
2006 pass
2007
2008 -class WeightFileNoGzip(WeightFile, EventFileNoGzip):
2009 pass
2010
2011 2012 -class FourMomentum(object):
2013 """a convenient object for 4-momenta operation""" 2014
2015 - def __init__(self, obj=0, px=0, py=0, pz=0, E=0):
2016 """initialize the four momenta""" 2017 2018 if obj is 0 and E: 2019 obj = E 2020 2021 if isinstance(obj, (FourMomentum, Particle)): 2022 px = obj.px 2023 py = obj.py 2024 pz = obj.pz 2025 E = obj.E 2026 elif isinstance(obj, (list, tuple)): 2027 assert len(obj) ==4 2028 E = obj[0] 2029 px = obj[1] 2030 py = obj[2] 2031 pz = obj[3] 2032 elif isinstance(obj, str): 2033 obj = [float(i) for i in obj.split()] 2034 assert len(obj) ==4 2035 E = obj[0] 2036 px = obj[1] 2037 py = obj[2] 2038 pz = obj[3] 2039 else: 2040 E =obj 2041 2042 2043 self.E = float(E) 2044 self.px = float(px) 2045 self.py = float(py) 2046 self.pz = float(pz)
2047 2048 @property
2049 - def mass(self):
2050 """return the mass""" 2051 return math.sqrt(max(self.E**2 - self.px**2 - self.py**2 - self.pz**2,0))
2052 2053 @property
2054 - def mass_sqr(self):
2055 """return the mass square""" 2056 return max(self.E**2 - self.px**2 - self.py**2 - self.pz**2,0)
2057 2058 @property
2059 - def pt(self):
2060 return math.sqrt(max(0, self.pt2()))
2061 2062 @property
2063 - def pseudorapidity(self):
2064 norm = math.sqrt(self.px**2 + self.py**2+self.pz**2) 2065 return 0.5* math.log((norm - self.pz) / (norm + self.pz))
2066 2067 @property
2068 - def rapidity(self):
2069 return 0.5* math.log((self.E +self.pz) / (self.E - self.pz))
2070 2071 2072
2073 - def pt2(self):
2074 """ return the pt square """ 2075 2076 return self.px**2 + self.py**2
2077
2078 - def __add__(self, obj):
2079 2080 assert isinstance(obj, FourMomentum) 2081 new = FourMomentum(self.E+obj.E, 2082 self.px + obj.px, 2083 self.py + obj.py, 2084 self.pz + obj.pz) 2085 return new
2086
2087 - def __iadd__(self, obj):
2088 """update the object with the sum""" 2089 self.E += obj.E 2090 self.px += obj.px 2091 self.py += obj.py 2092 self.pz += obj.pz 2093 return self
2094
2095 - def __sub__(self, obj):
2096 2097 assert isinstance(obj, FourMomentum) 2098 new = FourMomentum(self.E-obj.E, 2099 self.px - obj.px, 2100 self.py - obj.py, 2101 self.pz - obj.pz) 2102 return new
2103
2104 - def __isub__(self, obj):
2105 """update the object with the sum""" 2106 self.E -= obj.E 2107 self.px -= obj.px 2108 self.py -= obj.py 2109 self.pz -= obj.pz 2110 return self
2111
2112 - def __mul__(self, obj):
2113 if isinstance(obj, FourMomentum): 2114 return self.E*obj.E - self.px *obj.px - self.py * obj.py - self.pz * obj.pz 2115 elif isinstance(obj, (float, int)): 2116 return FourMomentum(obj*self.E,obj*self.px,obj*self.py,obj*self.pz ) 2117 else: 2118 raise NotImplemented
2119 __rmul__ = __mul__ 2120
2121 - def __pow__(self, power):
2122 assert power in [1,2] 2123 2124 if power == 1: 2125 return FourMomentum(self) 2126 elif power == 2: 2127 return self.mass_sqr
2128
2129 - def __repr__(self):
2130 return 'FourMomentum(%s,%s,%s,%s)' % (self.E, self.px, self.py,self.pz)
2131
2132 - def get_tuple(self):
2133 return (self.E, self.px, self.py,self.pz)
2134
2135 - def boost(self, mom):
2136 """mom 4-momenta is suppose to be given in the rest frame of this 4-momenta. 2137 the output is the 4-momenta in the frame of this 4-momenta 2138 function copied from HELAS routine.""" 2139 2140 2141 pt = self.px**2 + self.py**2 + self.pz**2 2142 if pt: 2143 s3product = self.px * mom.px + self.py * mom.py + self.pz * mom.pz 2144 mass = self.mass 2145 lf = (mom.E + (self.E - mass) * s3product / pt ) / mass 2146 return FourMomentum(E=(self.E*mom.E+s3product)/mass, 2147 px=mom.px + self.px * lf, 2148 py=mom.py + self.py * lf, 2149 pz=mom.pz + self.pz * lf) 2150 else: 2151 return FourMomentum(mom)
2152
2153 - def zboost(self, pboost=None, E=0, pz=0):
2154 """Both momenta should be in the same frame. 2155 The boost perform correspond to the boost required to set pboost at 2156 rest (only z boost applied). 2157 """ 2158 if isinstance(pboost, FourMomentum): 2159 E = pboost.E 2160 pz = pboost.pz 2161 2162 #beta = pz/E 2163 gamma = E / math.sqrt(E**2-pz**2) 2164 gammabeta = pz / math.sqrt(E**2-pz**2) 2165 2166 out = FourMomentum([gamma*self.E - gammabeta*self.pz, 2167 self.px, 2168 self.py, 2169 gamma*self.pz - gammabeta*self.E]) 2170 2171 if abs(out.pz) < 1e-6 * out.E: 2172 out.pz = 0 2173 return out
2174
2175 - def boost_to_restframe(self, pboost):
2176 """apply the boost transformation such that pboost is at rest in the new frame. 2177 First apply a rotation to allign the pboost to the z axis and then use 2178 zboost routine (see above) 2179 """ 2180 2181 2182 2183 2184 # write pboost as (E, p cosT sinF, p sinT sinF, p cosF) 2185 # rotation such that it become (E, 0 , 0 , p ) is 2186 # cosT sinF , -sinT , cosT sinF 2187 # sinT cosF , cosT , sinT sinF 2188 # -sinT , 0 , cosF 2189 p = math.sqrt( pboost.px**2 + pboost.py**2+ pboost.pz**2) 2190 cosF = pboost.pz / p 2191 sinF = math.sqrt(1-cosF**2) 2192 sinT = pboost.py/p/sinF 2193 cosT = pboost.px/p/sinF 2194 2195 out=FourMomentum([self.E, 2196 self.px*cosT*cosF + self.py*sinT*cosF-self.pz*sinF, 2197 -self.px*sinT+ self.py*cosT, 2198 self.px*cosT*sinF + self.py*sinT*sinF + self.pz*cosF 2199 ]) 2200 out = out.zboost(E=pboost.E,pz=p) 2201 return out
2202
2203 2204 2205 2206 -class OneNLOWeight(object):
2207
2208 - def __init__(self, input, real_type=(1,11)):
2209 """ """ 2210 2211 self.real_type = real_type 2212 if isinstance(input, str): 2213 self.parse(input)
2214
2215 - def __str__(self):
2216 2217 out = """ pwgt: %(pwgt)s 2218 born, real : %(born)s %(real)s 2219 pdgs : %(pdgs)s 2220 bjks : %(bjks)s 2221 scales**2, gs: %(scales2)s %(gs)s 2222 born/real related : %(born_related)s %(real_related)s 2223 type / nfks : %(type)s %(nfks)s 2224 to merge : %(to_merge_pdg)s in %(merge_new_pdg)s 2225 ref_wgt : %(ref_wgt)s""" % self.__dict__ 2226 return out
2227 2228
2229 - def parse(self, text):
2230 """parse the line and create the related object""" 2231 #0.546601845792D+00 0.000000000000D+00 0.000000000000D+00 0.119210435309D+02 0.000000000000D+00 5 -1 2 -11 12 21 0 0.24546101D-01 0.15706890D-02 0.12586055D+04 0.12586055D+04 0.12586055D+04 1 2 2 2 5 2 2 0.539995789976D+04 2232 #0.274922677249D+01 0.000000000000D+00 0.000000000000D+00 0.770516514633D+01 0.113763730192D+00 5 21 2 -11 12 1 2 0.52500539D-02 0.30205908D+00 0.45444066D+04 0.45444066D+04 0.45444066D+04 0.12520062D+01 1 2 1 3 5 1 -1 0.110944218997D+05 2233 # below comment are from Rik description email 2234 data = text.split() 2235 # 1. The first three doubles are, as before, the 'wgt', i.e., the overall event of this 2236 # contribution, and the ones multiplying the log[mu_R/QES] and the log[mu_F/QES] 2237 # stripped of alpha_s and the PDFs. 2238 # from example: 0.274922677249D+01 0.000000000000D+00 0.000000000000D+00 2239 self.pwgt = [float(f) for f in data[:3]] 2240 # 2. The next two doubles are the values of the (corresponding) Born and 2241 # real-emission matrix elements. You can either use these values to check 2242 # that the newly computed original matrix element weights are correct, 2243 # or directly use these so that you don't have to recompute the original weights. 2244 # For contributions for which the real-emission matrix elements were 2245 # not computed, the 2nd of these numbers is zero. The opposite is not true, 2246 # because each real-emission phase-space configuration has an underlying Born one 2247 # (this is not unique, but on our code we made a specific choice here). 2248 # This latter information is useful if the real-emission matrix elements 2249 # are unstable; you can then reweight with the Born instead. 2250 # (see also point 9 below, where the momentum configurations are assigned). 2251 # I don't think this instability is real problem when reweighting the real-emission 2252 # with tree-level matrix elements (as we generally would do), but is important 2253 # when reweighting with loop-squared contributions as we have been doing for gg->H. 2254 # (I'm not sure that reweighting tree-level with loop^2 is something that 2255 # we can do in general, because we don't really know what to do with the 2256 # virtual matrix elements because we cannot generate 2-loop diagrams.) 2257 # from example: 0.770516514633D+01 0.113763730192D+00 2258 self.born = float(data[3]) 2259 self.real = float(data[4]) 2260 # 3. integer: number of external particles of the real-emission configuration (as before) 2261 # from example: 5 2262 self.nexternal = int(data[5]) 2263 # 4. PDG codes corresponding to the real-emission configuration (as before) 2264 # from example: 21 2 -11 12 1 2 2265 self.pdgs = [int(i) for i in data[6:6+self.nexternal]] 2266 flag = 6+self.nexternal # new starting point for the position 2267 # 5. next integer is the power of g_strong in the matrix elements (as before) 2268 # from example: 2 2269 self.qcdpower = int(data[flag]) 2270 # 6. 2 doubles: The bjorken x's used for this contribution (as before) 2271 # from example: 0.52500539D-02 0.30205908D+00 2272 self.bjks = [float(f) for f in data[flag+1:flag+3]] 2273 # 7. 3 doubles: The Ellis-sexton scale, the renormalisation scale and the factorisation scale, all squared, used for this contribution (as before) 2274 # from example: 0.45444066D+04 0.45444066D+04 0.45444066D+04 2275 self.scales2 = [float(f) for f in data[flag+3:flag+6]] 2276 # 8.the value of g_strong 2277 # from example: 0.12520062D+01 2278 self.gs = float(data[flag+6]) 2279 # 9. 2 integers: the corresponding Born and real-emission type kinematics. (in the list of momenta) 2280 # Note that also the Born-kinematics has n+1 particles, with, in general, 2281 # one particle with zero momentum (this is not ALWAYS the case, 2282 # there could also be 2 particles with perfectly collinear momentum). 2283 # To convert this from n+1 to a n particles, you have to sum the momenta 2284 # of the two particles that 'merge', see point 12 below. 2285 # from example: 1 2 2286 self.born_related = int(data[flag+7]) 2287 self.real_related = int(data[flag+8]) 2288 # 10. 1 integer: the 'type'. This is the information you should use to determine 2289 # if to reweight with Born, virtual or real-emission matrix elements. 2290 # (Apart from the possible problems with complicated real-emission matrix elements 2291 # that need to be computed very close to the soft/collinear limits, see point 2 above. 2292 # I guess that for tree-level this is always okay, but when reweighting 2293 # a tree-level contribution with a one-loop squared one, as we do 2294 # for gg->Higgs, this is important). 2295 # type=1 : real-emission: 2296 # type=2 : Born: 2297 # type=3 : integrated counter terms: 2298 # type=4 : soft counter-term: 2299 # type=5 : collinear counter-term: 2300 # type=6 : soft-collinear counter-term: 2301 # type=7 : O(alphaS) expansion of Sudakov factor for NNLL+NLO: 2302 # type=8 : soft counter-term (with n+1-body kin.): 2303 # type=9 : collinear counter-term (with n+1-body kin.): 2304 # type=10: soft-collinear counter-term (with n+1-body kin.): 2305 # type=11: real-emission (with n-body kin.): 2306 # type=12: MC subtraction with n-body kin.: 2307 # type=13: MC subtraction with n+1-body kin.: 2308 # type=14: virtual corrections minus approximate virtual 2309 # type=15: approximate virtual corrections: 2310 # from example: 1 2311 self.type = int(data[flag+9]) 2312 # 11. 1 integer: The FKS configuration for this contribution (not really 2313 # relevant for anything, but is used in checking the reweighting to 2314 # get scale & PDF uncertainties). 2315 # from example: 3 2316 self.nfks = int(data[flag+10]) 2317 # 12. 2 integers: the two particles that should be merged to form the 2318 # born contribution from the real-emission one. Remove these two particles 2319 # from the (ordered) list of PDG codes, and insert a newly created particle 2320 # at the location of the minimum of the two particles removed. 2321 # I.e., if you merge particles 2 and 4, you have to insert the new particle 2322 # as the 2nd particle. And particle 5 and above will be shifted down by one. 2323 # from example: 5 1 2324 self.to_merge_pdg = [int (f) for f in data[flag+11:flag+13]] 2325 # 13. 1 integer: the PDG code of the particle that is created after merging the two particles at point 12. 2326 # from example -1 2327 self.merge_new_pdg = int(data[flag+13]) 2328 # 14. 1 double: the reference number that one should be able to reconstruct 2329 # form the weights (point 1 above) and the rest of the information of this line. 2330 # This is really the contribution to this event as computed by the code 2331 # (and is passed to the integrator). It contains everything. 2332 # from example: 0.110944218997D+05 2333 self.ref_wgt = float(data[flag+14]) 2334 2335 #check the momenta configuration linked to the event 2336 if self.type in self.real_type: 2337 self.momenta_config = self.real_related 2338 else: 2339 self.momenta_config = self.born_related
2340
2341 2342 -class NLO_PARTIALWEIGHT(object):
2343
2344 - class BasicEvent(list):
2345
2346 - def __init__(self, momenta, wgts, event, real_type=(1,11)):
2347 list.__init__(self, momenta) 2348 2349 assert self 2350 self.wgts = wgts 2351 self.pdgs = list(wgts[0].pdgs) 2352 self.event = event 2353 self.real_type = real_type 2354 2355 if wgts[0].momenta_config == wgts[0].born_related: 2356 # need to remove one momenta. 2357 ind1, ind2 = [ind-1 for ind in wgts[0].to_merge_pdg] 2358 if ind1> ind2: 2359 ind1, ind2 = ind2, ind1 2360 if ind1 >= sum(1 for p in event if p.status==-1): 2361 new_p = self[ind1] + self[ind2] 2362 else: 2363 new_p = self[ind1] - self[ind2] 2364 self.pop(ind1) 2365 self.insert(ind1, new_p) 2366 self.pop(ind2) 2367 self.pdgs.pop(ind1) 2368 self.pdgs.insert(ind1, wgts[0].merge_new_pdg ) 2369 self.pdgs.pop(ind2) 2370 # DO NOT update the pdgs of the partial weight! 2371 elif any(w.type in self.real_type for w in wgts): 2372 if any(w.type not in self.real_type for w in wgts): 2373 raise Exception 2374 # check if this is too soft/colinear if so use the born 2375 ind1, ind2 = [ind-1 for ind in wgts[0].to_merge_pdg] 2376 if ind1> ind2: 2377 ind1, ind2 = ind2, ind1 2378 if ind1 >= sum(1 for p in event if p.status==-1): 2379 new_p = self[ind1] + self[ind2] 2380 else: 2381 new_p = self[ind1] - self[ind2] 2382 2383 if __debug__: 2384 ptot = FourMomentum() 2385 for i in xrange(len(self)): 2386 if i <2: 2387 ptot += self[i] 2388 else: 2389 ptot -= self[i] 2390 if ptot.mass_sqr > 1e-16: 2391 misc.sprint(ptot, ptot.mass_sqr) 2392 2393 inv_mass = new_p.mass_sqr 2394 shat = (self[0]+self[1]).mass_sqr 2395 if (abs(inv_mass)/shat < 1e-6): 2396 # misc.sprint(abs(inv_mass)/shat) 2397 self.pop(ind1) 2398 self.insert(ind1, new_p) 2399 self.pop(ind2) 2400 self.pdgs.pop(ind1) 2401 self.pdgs.insert(ind1, wgts[0].merge_new_pdg ) 2402 self.pdgs.pop(ind2) 2403 # DO NOT update the pdgs of the partial weight! 2404 2405 if __debug__: 2406 ptot = FourMomentum() 2407 for i in xrange(len(self)): 2408 if i <2: 2409 ptot += self[i] 2410 else: 2411 ptot -= self[i] 2412 if ptot.mass_sqr > 1e-16: 2413 misc.sprint(ptot, ptot.mass_sqr) 2414 # raise Exception 2415 else: 2416 raise Exception
2417
2418 - def get_pdg_code(self):
2419 return self.pdgs
2420
2421 - def get_tag_and_order(self):
2422 """ return the tag and order for this basic event""" 2423 (initial, _), _ = self.event.get_tag_and_order() 2424 order = self.get_pdg_code() 2425 2426 2427 initial, out = order[:len(initial)], order[len(initial):] 2428 initial.sort() 2429 out.sort() 2430 return (tuple(initial), tuple(out)), order
2431
2432 - def get_momenta(self, get_order, allow_reversed=True):
2433 """return the momenta vector in the order asked for""" 2434 2435 #avoid to modify the input 2436 order = [list(get_order[0]), list(get_order[1])] 2437 out = [''] *(len(order[0])+len(order[1])) 2438 pdgs = self.get_pdg_code() 2439 for pos, part in enumerate(self): 2440 if pos < len(get_order[0]): #initial 2441 try: 2442 ind = order[0].index(pdgs[pos]) 2443 except ValueError, error: 2444 if not allow_reversed: 2445 raise error 2446 else: 2447 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 2448 try: 2449 return self.get_momenta(order, False) 2450 except ValueError: 2451 raise error 2452 2453 2454 position = ind 2455 order[0][ind] = 0 2456 else: #final 2457 try: 2458 ind = order[1].index(pdgs[pos]) 2459 except ValueError, error: 2460 if not allow_reversed: 2461 raise error 2462 else: 2463 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 2464 try: 2465 return self.get_momenta(order, False) 2466 except ValueError: 2467 raise error 2468 position = len(order[0]) + ind 2469 order[1][ind] = 0 2470 2471 out[position] = (part.E, part.px, part.py, part.pz) 2472 2473 return out
2474 2475
2476 - def get_helicity(self, *args):
2477 return [9] * len(self)
2478 2479 @property
2480 - def aqcd(self):
2481 return self.event.aqcd
2482
2483 - def get_ht_scale(self, prefactor=1):
2484 2485 scale = 0 2486 for particle in self: 2487 p = particle 2488 scale += math.sqrt(max(0, p.mass_sqr + p.pt**2)) 2489 2490 return prefactor * scale
2491
2492 - def get_et_scale(self, prefactor=1):
2493 2494 scale = 0 2495 for particle in self: 2496 p = particle 2497 pt = p.pt 2498 if (pt>0): 2499 scale += p.E*pt/math.sqrt(pt**2+p.pz**2) 2500 2501 return prefactor * scale
2502 2503
2504 - def get_sqrts_scale(self, event,prefactor=1):
2505 2506 scale = 0 2507 nb_init = 0 2508 for particle in event: 2509 if particle.status == -1: 2510 nb_init+=1 2511 if nb_init == 1: 2512 return self[0].mass 2513 elif nb_init==2: 2514 return math.sqrt((self[0]+self[1])**2)
2515 2516 2517 2518
2519 - def __init__(self, input, event, real_type=(1,11)):
2520 2521 self.real_type = real_type 2522 self.event = event 2523 if isinstance(input, str): 2524 self.parse(input)
2525 2526
2527 - def parse(self, text):
2528 """create the object from the string information (see example below)""" 2529 #0.2344688900d+00 8 2 0 2530 #0.4676614699d+02 0.0000000000d+00 0.0000000000d+00 0.4676614699d+02 2531 #0.4676614699d+02 0.0000000000d+00 0.0000000000d+00 -.4676614699d+02 2532 #0.4676614699d+02 0.2256794794d+02 0.4332148227d+01 0.4073073437d+02 2533 #0.4676614699d+02 -.2256794794d+02 -.4332148227d+01 -.4073073437d+02 2534 #0.0000000000d+00 -.0000000000d+00 -.0000000000d+00 -.0000000000d+00 2535 #0.4780341163d+02 0.0000000000d+00 0.0000000000d+00 0.4780341163d+02 2536 #0.4822581633d+02 0.0000000000d+00 0.0000000000d+00 -.4822581633d+02 2537 #0.4729127470d+02 0.2347155377d+02 0.5153455534d+01 0.4073073437d+02 2538 #0.4627255267d+02 -.2167412893d+02 -.3519736379d+01 -.4073073437d+02 2539 #0.2465400591d+01 -.1797424844d+01 -.1633719155d+01 -.4224046944d+00 2540 #0.473706252575d-01 0.000000000000d+00 0.000000000000d+00 5 -3 3 -11 11 21 0 0.11849903d-02 0.43683926d-01 0.52807978d+03 0.52807978d+03 0.52807978d+03 1 2 1 0.106660059627d+03 2541 #-.101626389492d-02 0.000000000000d+00 -.181915673961d-03 5 -3 3 -11 11 21 2 0.11849903d-02 0.43683926d-01 0.52807978d+03 0.52807978d+03 0.52807978d+03 1 3 1 -.433615206719d+01 2542 #0.219583436285d-02 0.000000000000d+00 0.000000000000d+00 5 -3 3 -11 11 21 2 0.11849903d-02 0.43683926d-01 0.52807978d+03 0.52807978d+03 0.52807978d+03 1 15 1 0.936909375537d+01 2543 #0.290043597283d-03 0.000000000000d+00 0.000000000000d+00 5 -3 3 -11 11 21 2 0.12292838d-02 0.43683926d-01 0.58606724d+03 0.58606724d+03 0.58606724d+03 1 12 1 0.118841547979d+01 2544 #-.856330613460d-01 0.000000000000d+00 0.000000000000d+00 5 -3 3 -11 11 21 2 0.11849903d-02 0.43683926d-01 0.52807978d+03 0.52807978d+03 0.52807978d+03 1 4 1 -.365375546483d+03 2545 #0.854918237609d-01 0.000000000000d+00 0.000000000000d+00 5 -3 3 -11 11 21 2 0.12112732d-02 0.45047393d-01 0.58606724d+03 0.58606724d+03 0.58606724d+03 2 11 1 0.337816057347d+03 2546 #0.359257891118d-05 0.000000000000d+00 0.000000000000d+00 5 21 3 -11 11 3 2 0.12292838d-02 0.43683926d-01 0.58606724d+03 0.58606724d+03 0.58606724d+03 1 12 3 0.334254554762d+00 2547 #0.929944817736d-03 0.000000000000d+00 0.000000000000d+00 5 21 3 -11 11 3 2 0.12112732d-02 0.45047393d-01 0.58606724d+03 0.58606724d+03 0.58606724d+03 2 11 3 0.835109616010d+02 2548 2549 text = text.lower().replace('d','e') 2550 all_line = text.split('\n') 2551 #get global information 2552 first_line ='' 2553 while not first_line.strip(): 2554 first_line = all_line.pop(0) 2555 2556 wgt, nb_wgt, nb_event, _ = first_line.split() 2557 nb_wgt, nb_event = int(nb_wgt), int(nb_event) 2558 2559 momenta = [] 2560 wgts = [] 2561 for line in all_line: 2562 data = line.split() 2563 if len(data) == 4: 2564 p = FourMomentum(data) 2565 momenta.append(p) 2566 elif len(data)>0: 2567 wgt = OneNLOWeight(line, real_type=self.real_type) 2568 wgts.append(wgt) 2569 2570 assert len(wgts) == int(nb_wgt) 2571 2572 get_weights_for_momenta = {} 2573 size_momenta = 0 2574 for wgt in wgts: 2575 if wgt.momenta_config in get_weights_for_momenta: 2576 get_weights_for_momenta[wgt.momenta_config].append(wgt) 2577 else: 2578 if size_momenta == 0: size_momenta = wgt.nexternal 2579 assert size_momenta == wgt.nexternal 2580 get_weights_for_momenta[wgt.momenta_config] = [wgt] 2581 2582 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) 2583 2584 2585 2586 self.cevents = [] 2587 for key in range(1, nb_event+1): 2588 if key in get_weights_for_momenta: 2589 wgt = get_weights_for_momenta[key] 2590 evt = self.BasicEvent(momenta[:size_momenta], get_weights_for_momenta[key], self.event, self.real_type) 2591 self.cevents.append(evt) 2592 momenta = momenta[size_momenta:] 2593 2594 nb_wgt_check = 0 2595 for cevt in self.cevents: 2596 nb_wgt_check += len(cevt.wgts) 2597 assert nb_wgt_check == int(nb_wgt)
2598 2599 2600 2601 if '__main__' == __name__: 2602 2603 lhe = EventFile('unweighted_events.lhe.gz') 2604 #lhe.parsing = False 2605 start = time.time() 2606 for event in lhe: 2607 event.parse_lo_weight() 2608 print 'old method -> ', time.time()-start 2609 lhe = EventFile('unweighted_events.lhe.gz') 2610 #lhe.parsing = False 2611 start = time.time() 2612 for event in lhe: 2613 event.parse_lo_weight_test() 2614 print 'new method -> ', time.time()-start 2615 2616 2617 # Example 1: adding some missing information to the event (here distance travelled) 2618 if False: 2619 start = time 2620 lhe = EventFile('unweighted_events.lhe.gz') 2621 output = open('output_events.lhe', 'w') 2622 #write the banner to the output file 2623 output.write(lhe.banner) 2624 # Loop over all events 2625 for event in lhe: 2626 for particle in event: 2627 # modify particle attribute: here remove the mass 2628 particle.mass = 0 2629 particle.vtim = 2 # The one associate to distance travelled by the particle. 2630 2631 #write this modify event 2632 output.write(str(event)) 2633 output.write('</LesHouchesEvent>\n') 2634 2635 # Example 3: Plotting some variable 2636 if False: 2637 lhe = EventFile('unweighted_events.lhe.gz') 2638 import matplotlib.pyplot as plt 2639 import matplotlib.gridspec as gridspec 2640 nbins = 100 2641 2642 nb_pass = 0 2643 data = [] 2644 for event in lhe: 2645 etaabs = 0 2646 etafinal = 0 2647 for particle in event: 2648 if particle.status==1: 2649 p = FourMomentum(particle) 2650 eta = p.pseudorapidity 2651 if abs(eta) > etaabs: 2652 etafinal = eta 2653 etaabs = abs(eta) 2654 if etaabs < 4: 2655 data.append(etafinal) 2656 nb_pass +=1 2657 2658 2659 print nb_pass 2660 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1]) 2661 gs1.update(wspace=0, hspace=0) # set the spacing between axes. 2662 ax = plt.subplot(gs1[0]) 2663 2664 n, bins, patches = ax.hist(data, nbins, histtype='step', label='original') 2665 ax_c = ax.twinx() 2666 ax_c.set_ylabel('MadGraph5_aMC@NLO') 2667 ax_c.yaxis.set_label_coords(1.01, 0.25) 2668 ax_c.set_yticks(ax.get_yticks()) 2669 ax_c.set_yticklabels([]) 2670 ax.set_xlim([-4,4]) 2671 print "bin value:", n 2672 print "start/end point of bins", bins 2673 plt.axis('on') 2674 plt.xlabel('weight ratio') 2675 plt.show() 2676 2677 2678 # Example 4: More complex plotting example (with ratio plot) 2679 if False: 2680 lhe = EventFile('unweighted_events.lhe') 2681 import matplotlib.pyplot as plt 2682 import matplotlib.gridspec as gridspec 2683 nbins = 100 2684 2685 #mtau, wtau = 45, 5.1785e-06 2686 mtau, wtau = 1.777, 4.027000e-13 2687 nb_pass = 0 2688 data, data2, data3 = [], [], [] 2689 for event in lhe: 2690 nb_pass +=1 2691 if nb_pass > 10000: 2692 break 2693 tau1 = FourMomentum() 2694 tau2 = FourMomentum() 2695 for part in event: 2696 if part.pid in [-12,11,16]: 2697 momenta = FourMomentum(part) 2698 tau1 += momenta 2699 elif part.pid == 15: 2700 tau2 += FourMomentum(part) 2701 2702 if abs((mtau-tau2.mass())/wtau)<1e6 and tau2.mass() >1: 2703 data.append((tau1.mass()-mtau)/wtau) 2704 data2.append((tau2.mass()-mtau)/wtau) 2705 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1]) 2706 gs1.update(wspace=0, hspace=0) # set the spacing between axes. 2707 ax = plt.subplot(gs1[0]) 2708 2709 n, bins, patches = ax.hist(data2, nbins, histtype='step', label='original') 2710 n2, bins2, patches2 = ax.hist(data, bins=bins, histtype='step',label='reconstructed') 2711 import cmath 2712 2713 breit = lambda m : math.sqrt(4*math.pi)*1/(((m)**2-mtau**2)**2+(mtau*wtau)**2)*wtau 2714 2715 data3 = [breit(mtau + x*wtau)*wtau*16867622.6624*50 for x in bins] 2716 2717 ax.plot(bins, data3,label='breit-wigner') 2718 # add the legend 2719 ax.legend() 2720 # add on the right program tag 2721 ax_c = ax.twinx() 2722 ax_c.set_ylabel('MadGraph5_aMC@NLO') 2723 ax_c.yaxis.set_label_coords(1.01, 0.25) 2724 ax_c.set_yticks(ax.get_yticks()) 2725 ax_c.set_yticklabels([]) 2726 2727 plt.title('invariant mass of tau LHE/reconstructed') 2728 plt.axis('on') 2729 ax.set_xticklabels([]) 2730 # ratio plot 2731 ax = plt.subplot(gs1[1]) 2732 data4 = [n[i]/(data3[i]) for i in range(nbins)] 2733 ax.plot(bins, data4 + [0] , 'b') 2734 data4 = [n2[i]/(data3[i]) for i in range(nbins)] 2735 ax.plot(bins, data4 + [0] , 'g') 2736 ax.set_ylim([0,2]) 2737 #remove last y tick to avoid overlap with above plot: 2738 tick = ax.get_yticks() 2739 ax.set_yticks(tick[:-1]) 2740 2741 2742 plt.axis('on') 2743 plt.xlabel('(M - Mtau)/Wtau') 2744 plt.show() 2745