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):
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 if init_information["idbmup1"] == 0: 861 event = self.next() 862 init_information["idbmup1"]= event[0].pdg 863 if init_information["idbmup2"] == 0: 864 init_information["idbmup2"]= event[1].pdg 865 self.seek(0) 866 if init_information["idbmup2"] == 0: 867 event = self.next() 868 init_information["idbmup2"] = event[1].pdg 869 self.seek(0) 870 871 init_information["nprup"] = nb_group 872 873 if run_card["lhe_version"] < 3: 874 init_information["generator_info"] = "" 875 else: 876 init_information["generator_info"] = "<generator name='MadGraph5_aMC@NLO' version='%s'>please cite 1405.0301 </generator>\n" \ 877 % misc.get_pkg_info()['version'] 878 879 # cross_information: 880 cross_info = "%(cross)e %(error)e %(wgt)e %(id)i" 881 init_information["cross_info"] = [] 882 for id in grouped_cross: 883 conv = {"id": int(id), "cross": grouped_cross[id], "error": math.sqrt(grouped_error[id]), 884 "wgt": wgt} 885 init_information["cross_info"].append( cross_info % conv) 886 init_information["cross_info"] = '\n'.join(init_information["cross_info"]) 887 init_information['lha_stra'] = -1 * abs(lha_strategy) 888 889 template_init =\ 890 """ %(idbmup1)i %(idbmup2)i %(ebmup1)e %(ebmup2)e %(pdfgup1)i %(pdfgup2)i %(pdfsup1)i %(pdfsup2)i %(lha_stra)i %(nprup)i 891 %(cross_info)s 892 %(generator_info)s 893 """ 894 895 self.banner["init"] = template_init % init_information
896 897 898
899 - def initialize_unweighting(self, getwgt, trunc_error):
900 """ scan once the file to return 901 - the list of the hightest weight (of size trunc_error*NB_EVENT 902 - the cross-section by type of process 903 - the total number of events in the files 904 In top of that it initialise the information for the next routine 905 to determine how to choose which file to read 906 """ 907 self.seek(0) 908 all_wgt = [] 909 total_event = 0 910 sum_cross = collections.defaultdict(int) 911 for i,f in enumerate(self.files): 912 nb_event = 0 913 # We need to loop over the event file to get some information about the 914 # new cross-section/ wgt of event. 915 cross = collections.defaultdict(int) 916 new_wgt =[] 917 for event in f: 918 nb_event += 1 919 total_event += 1 920 event.sample_scale = 1 921 wgt = getwgt(event) 922 cross['all'] += wgt 923 cross['abs'] += abs(wgt) 924 cross[event.ievent] += wgt 925 new_wgt.append(abs(wgt)) 926 # avoid all_wgt to be too large 927 if nb_event % 20000 == 0: 928 new_wgt.sort() 929 # drop the lowest weight 930 nb_keep = max(20, int(nb_event*trunc_error*15)) 931 new_wgt = new_wgt[-nb_keep:] 932 if nb_event == 0: 933 raise Exception 934 # store the information 935 self.initial_nb_events[i] = nb_event 936 self.scales[i] = self.across[i]/cross['abs'] if self.across[i] else 1 937 #misc.sprint("sum of wgt in event %s is %s. Should be %s => scale %s (nb_event: %s)" 938 # % (i, cross['all'], self.allcross[i], self.scales[i], nb_event)) 939 for key in cross: 940 sum_cross[key] += cross[key]* self.scales[i] 941 all_wgt +=[self.scales[i] * w for w in new_wgt] 942 all_wgt.sort() 943 nb_keep = max(20, int(total_event*trunc_error*10)) 944 all_wgt = all_wgt[-nb_keep:] 945 946 self.total_event_in_files = total_event 947 #final selection of the interesting weight to keep 948 all_wgt.sort() 949 # drop the lowest weight 950 nb_keep = max(20, int(total_event*trunc_error*10)) 951 all_wgt = all_wgt[-nb_keep:] 952 self.seek(0) 953 self._configure = True 954 return all_wgt, sum_cross, total_event
955
956 - def configure(self):
957 958 self._configure = True 959 for i,f in enumerate(self.files): 960 self.initial_nb_events[i] = len(f) 961 self.total_event_in_files = sum(self.initial_nb_events)
962
963 - def __len__(self):
964 965 return len(self.files)
966
967 - def seek(self, pos):
968 """ """ 969 970 if pos !=0: 971 raise Exception 972 for i in range(len(self)): 973 self.curr_nb_events[i] = 0 974 for f in self.files: 975 f.seek(pos)
976
977 - def unweight(self, outputpath, get_wgt, **opts):
978 """unweight the current file according to wgt information wgt. 979 which can either be a fct of the event or a tag in the rwgt list. 980 max_wgt allow to do partial unweighting. 981 trunc_error allow for dynamical partial unweighting 982 event_target reweight for that many event with maximal trunc_error. 983 (stop to write event when target is reached) 984 """ 985 986 if isinstance(get_wgt, str): 987 unwgt_name =get_wgt 988 def get_wgt_multi(event): 989 event.parse_reweight() 990 return event.reweight_data[unwgt_name] * event.sample_scale
991 else: 992 unwgt_name = get_wgt.func_name 993 get_wgt_multi = lambda event: get_wgt(event) * event.sample_scale 994 #define the weighting such that we have built-in the scaling 995 996 997 if 'event_target' in opts and opts['event_target']: 998 if 'normalization' in opts: 999 if opts['normalization'] == 'sum': 1000 new_wgt = sum(self.across)/opts['event_target'] 1001 strategy = 3 1002 elif opts['normalization'] == 'average': 1003 strategy = 4 1004 new_wgt = sum(self.across) 1005 elif opts['normalization'] == 'unit': 1006 strategy =3 1007 new_wgt = 1. 1008 else: 1009 strategy = 4 1010 new_wgt = sum(self.across) 1011 self.define_init_banner(new_wgt, strategy) 1012 self.written_weight = new_wgt 1013 elif 'write_init' in opts and opts['write_init']: 1014 self.define_init_banner(0,0) 1015 del opts['write_init'] 1016 return super(MultiEventFile, self).unweight(outputpath, get_wgt_multi, **opts)
1017
1018 - def write(self, path, random=False, banner=None, get_info=False):
1019 """ """ 1020 1021 if isinstance(path, str): 1022 out = EventFile(path, 'w') 1023 if self.parsefile and not banner: 1024 banner = self.files[0].banner 1025 elif not banner: 1026 firstlhe = EventFile(self.files[0]) 1027 banner = firstlhe.banner 1028 else: 1029 out = path 1030 if banner: 1031 out.write(banner) 1032 nb_event = 0 1033 info = collections.defaultdict(float) 1034 if random and self.open: 1035 for event in self: 1036 nb_event +=1 1037 out.write(event) 1038 if get_info: 1039 event.parse_reweight() 1040 for key, value in event.reweight_data.items(): 1041 info[key] += value 1042 info['central'] += event.wgt 1043 elif not random: 1044 for i,f in enumerate(self.files): 1045 #check if we need to parse the file or not 1046 if not self.parsefile: 1047 if i==0: 1048 try: 1049 lhe = firstlhe 1050 except: 1051 lhe = EventFile(f) 1052 else: 1053 lhe = EventFile(f) 1054 else: 1055 lhe = f 1056 for event in lhe: 1057 nb_event +=1 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 out.write(str(event)) 1064 lhe.close() 1065 out.write("</LesHouchesEvents>\n") 1066 return nb_event, info
1067
1068 - def remove(self):
1069 """ """ 1070 if self.parsefile: 1071 for f in self.files: 1072 os.remove(f.name) 1073 else: 1074 for f in self.files: 1075 os.remove(f)
1076
1077 1078 1079 -class Event(list):
1080 """Class storing a single event information (list of particles + global information)""" 1081 1082 warning_order = True # raise a warning if the order of the particle are not in accordance of child/mother 1083
1084 - def __init__(self, text=None):
1085 """The initialization of an empty Event (or one associate to a text file)""" 1086 list.__init__(self) 1087 1088 # First line information 1089 self.nexternal = 0 1090 self.ievent = 0 1091 self.wgt = 0 1092 self.aqcd = 0 1093 self.scale = 0 1094 self.aqed = 0 1095 self.aqcd = 0 1096 # Weight information 1097 self.tag = '' 1098 self.eventflag = {} # for information in <event > 1099 self.comment = '' 1100 self.reweight_data = {} 1101 self.matched_scale_data = None 1102 self.syscalc_data = {} 1103 if text: 1104 self.parse(text)
1105 1106 1107
1108 - def parse(self, text):
1109 """Take the input file and create the structured information""" 1110 #text = re.sub(r'</?event>', '', text) # remove pointless tag 1111 status = 'first' 1112 for line in text.split('\n'): 1113 line = line.strip() 1114 if not line: 1115 continue 1116 elif line[0] == '#': 1117 self.comment += '%s\n' % line 1118 continue 1119 elif line.startswith('<event'): 1120 if '=' in line: 1121 found = re.findall(r"""(\w*)=(?:(?:['"])([^'"]*)(?=['"])|(\S*))""",line) 1122 #for '<event line=4 value=\'3\' error="5" test=" 1 and 2">\n' 1123 #return [('line', '', '4'), ('value', '3', ''), ('error', '5', ''), ('test', ' 1 and 2', '')] 1124 self.eventflag = dict((n, a1) if a1 else (n,a2) for n,a1,a2 in found) 1125 # return {'test': ' 1 and 2', 'line': '4', 'value': '3', 'error': '5'} 1126 continue 1127 1128 elif 'first' == status: 1129 if '<rwgt>' in line: 1130 status = 'tag' 1131 else: 1132 self.assign_scale_line(line) 1133 status = 'part' 1134 continue 1135 if '<' in line: 1136 status = 'tag' 1137 1138 if 'part' == status: 1139 part = Particle(line, event=self) 1140 if part.E != 0: 1141 self.append(part) 1142 elif self.nexternal: 1143 self.nexternal-=1 1144 else: 1145 if '</event>' in line: 1146 line = line.replace('</event>','',1) 1147 self.tag += '%s\n' % line 1148 1149 self.assign_mother()
1150
1151 - def assign_mother(self):
1152 # assign the mother: 1153 for i,particle in enumerate(self): 1154 if i < particle.mother1 or i < particle.mother2: 1155 if self.warning_order: 1156 logger.warning("Order of particle in the event did not agree with parent/child order. This might be problematic for some code.") 1157 Event.warning_order = False 1158 self.reorder_mother_child() 1159 return self.assign_mother() 1160 1161 if particle.mother1: 1162 try: 1163 particle.mother1 = self[int(particle.mother1) -1] 1164 except Exception: 1165 logger.warning("WRONG MOTHER INFO %s", self) 1166 particle.mother1 = 0 1167 if particle.mother2: 1168 try: 1169 particle.mother2 = self[int(particle.mother2) -1] 1170 except Exception: 1171 logger.warning("WRONG MOTHER INFO %s", self) 1172 particle.mother2 = 0
1173
1174 - def rescale_weights(self, ratio):
1175 """change all the weights by a given ratio""" 1176 1177 self.wgt *= ratio 1178 self.parse_reweight() 1179 for key in self.reweight_data: 1180 self.reweight_data[key] *= ratio 1181 return self.wgt
1182
1183 - def reorder_mother_child(self):
1184 """check and correct the mother/child position. 1185 only correct one order by call (but this is a recursive call)""" 1186 1187 tomove, position = None, None 1188 for i,particle in enumerate(self): 1189 if i < particle.mother1: 1190 # move i after particle.mother1 1191 tomove, position = i, particle.mother1-1 1192 break 1193 if i < particle.mother2: 1194 tomove, position = i, particle.mother2-1 1195 1196 # nothing to change -> we are done 1197 if not tomove: 1198 return 1199 1200 # move the particles: 1201 particle = self.pop(tomove) 1202 self.insert(int(position), particle) 1203 1204 #change the mother id/ event_id in the event. 1205 for i, particle in enumerate(self): 1206 particle.event_id = i 1207 #misc.sprint( i, particle.event_id) 1208 m1, m2 = particle.mother1, particle.mother2 1209 if m1 == tomove +1: 1210 particle.mother1 = position+1 1211 elif tomove < m1 <= position +1: 1212 particle.mother1 -= 1 1213 if m2 == tomove +1: 1214 particle.mother2 = position+1 1215 elif tomove < m2 <= position +1: 1216 particle.mother2 -= 1 1217 # re-call the function for the next potential change 1218 return self.reorder_mother_child()
1219 1220 1221 1222 1223 1224
1225 - def parse_reweight(self):
1226 """Parse the re-weight information in order to return a dictionary 1227 {key: value}. If no group is define group should be '' """ 1228 if self.reweight_data: 1229 return self.reweight_data 1230 self.reweight_data = {} 1231 self.reweight_order = [] 1232 start, stop = self.tag.find('<rwgt>'), self.tag.find('</rwgt>') 1233 if start != -1 != stop : 1234 pattern = re.compile(r'''<\s*wgt id=(?:\'|\")(?P<id>[^\'\"]+)(?:\'|\")\s*>\s*(?P<val>[\ded+-.]*)\s*</wgt>''',re.I) 1235 data = pattern.findall(self.tag[start:stop]) 1236 try: 1237 self.reweight_data = dict([(pid, float(value)) for (pid, value) in data 1238 if not self.reweight_order.append(pid)]) 1239 # the if is to create the order file on the flight 1240 except ValueError, error: 1241 raise Exception, 'Event File has unvalid weight. %s' % error 1242 self.tag = self.tag[:start] + self.tag[stop+7:] 1243 return self.reweight_data
1244
1245 - def parse_nlo_weight(self, real_type=(1,11)):
1246 """ """ 1247 if hasattr(self, 'nloweight'): 1248 return self.nloweight 1249 1250 start, stop = self.tag.find('<mgrwgt>'), self.tag.find('</mgrwgt>') 1251 if start != -1 != stop : 1252 1253 text = self.tag[start+8:stop] 1254 self.nloweight = NLO_PARTIALWEIGHT(text, self, real_type=real_type) 1255 return self.nloweight
1256
1257 - def parse_lo_weight(self):
1258 """ """ 1259 1260 1261 if hasattr(self, 'loweight'): 1262 return self.loweight 1263 1264 if not hasattr(Event, 'loweight_pattern'): 1265 Event.loweight_pattern = re.compile('''<rscale>\s*(?P<nqcd>\d+)\s+(?P<ren_scale>[\d.e+-]+)\s*</rscale>\s*\n\s* 1266 <asrwt>\s*(?P<asrwt>[\s\d.+-]+)\s*</asrwt>\s*\n\s* 1267 <pdfrwt\s+beam=["']?1["']?\>\s*(?P<beam1>[\s\d.e+-]*)\s*</pdfrwt>\s*\n\s* 1268 <pdfrwt\s+beam=["']?2["']?\>\s*(?P<beam2>[\s\d.e+-]*)\s*</pdfrwt>\s*\n\s* 1269 <totfact>\s*(?P<totfact>[\d.e+-]*)\s*</totfact> 1270 ''',re.X+re.I+re.M) 1271 1272 start, stop = self.tag.find('<mgrwt>'), self.tag.find('</mgrwt>') 1273 1274 if start != -1 != stop : 1275 text = self.tag[start+8:stop] 1276 1277 info = Event.loweight_pattern.search(text) 1278 self.loweight={} 1279 self.loweight['n_qcd'] = int(info.group('nqcd')) 1280 self.loweight['ren_scale'] = float(info.group('ren_scale')) 1281 self.loweight['asrwt'] =[float(x) for x in info.group('asrwt').split()[1:]] 1282 self.loweight['tot_fact'] = float(info.group('totfact')) 1283 1284 args = info.group('beam1').split() 1285 npdf = int(args[0]) 1286 self.loweight['n_pdfrw1'] = npdf 1287 self.loweight['pdf_pdg_code1'] = [int(i) for i in args[1:1+npdf]] 1288 self.loweight['pdf_x1'] = [float(i) for i in args[1+npdf:1+2*npdf]] 1289 self.loweight['pdf_q1'] = [float(i) for i in args[1+2*npdf:1+3*npdf]] 1290 args = info.group('beam2').split() 1291 npdf = int(args[0]) 1292 self.loweight['n_pdfrw2'] = npdf 1293 self.loweight['pdf_pdg_code2'] = [int(i) for i in args[1:1+npdf]] 1294 self.loweight['pdf_x2'] = [float(i) for i in args[1+npdf:1+2*npdf]] 1295 self.loweight['pdf_q2'] = [float(i) for i in args[1+2*npdf:1+3*npdf]] 1296 1297 else: 1298 return None 1299 return self.loweight
1300 1301
1302 - def parse_matching_scale(self):
1303 """Parse the line containing the starting scale for the shower""" 1304 1305 if self.matched_scale_data is not None: 1306 return self.matched_scale_data 1307 1308 self.matched_scale_data = [] 1309 1310 1311 pattern = re.compile("<scales\s|</scales>") 1312 data = re.split(pattern,self.tag) 1313 if len(data) == 1: 1314 return [] 1315 else: 1316 tmp = {} 1317 start,content, end = data 1318 self.tag = "%s%s" % (start, end) 1319 pattern = re.compile("pt_clust_(\d*)=\"([\de+-.]*)\"") 1320 for id,value in pattern.findall(content): 1321 tmp[int(id)] = float(value) 1322 for i in range(1, len(self)+1): 1323 if i in tmp: 1324 self.matched_scale_data.append(tmp[i]) 1325 else: 1326 self.matched_scale_data.append(-1) 1327 return self.matched_scale_data
1328
1329 - def parse_syscalc_info(self):
1330 """ parse the flag for syscalc between <mgrwt></mgrwt> 1331 <mgrwt> 1332 <rscale> 3 0.26552898E+03</rscale> 1333 <asrwt>0</asrwt> 1334 <pdfrwt beam="1"> 1 21 0.14527945E+00 0.26552898E+03</pdfrwt> 1335 <pdfrwt beam="2"> 1 21 0.15249110E-01 0.26552898E+03</pdfrwt> 1336 <totfact> 0.10344054E+04</totfact> 1337 </mgrwt> 1338 """ 1339 if self.syscalc_data: 1340 return self.syscalc_data 1341 1342 pattern = re.compile("<mgrwt>|</mgrwt>") 1343 pattern2 = re.compile("<(?P<tag>[\w]*)(?:\s*(\w*)=[\"'](.*)[\"']\s*|\s*)>(.*)</(?P=tag)>") 1344 data = re.split(pattern,self.tag) 1345 if len(data) == 1: 1346 return [] 1347 else: 1348 tmp = {} 1349 start,content, end = data 1350 self.tag = "%s%s" % (start, end) 1351 for tag, key, keyval, tagval in pattern2.findall(content): 1352 if key: 1353 self.syscalc_data[(tag, key, keyval)] = tagval 1354 else: 1355 self.syscalc_data[tag] = tagval 1356 return self.syscalc_data
1357 1358
1359 - def add_decay_to_particle(self, position, decay_event):
1360 """define the decay of the particle id by the event pass in argument""" 1361 1362 this_particle = self[position] 1363 #change the status to internal particle 1364 this_particle.status = 2 1365 this_particle.helicity = 0 1366 1367 # some usefull information 1368 decay_particle = decay_event[0] 1369 this_4mom = FourMomentum(this_particle) 1370 nb_part = len(self) #original number of particle 1371 1372 thres = decay_particle.E*1e-10 1373 assert max(decay_particle.px, decay_particle.py, decay_particle.pz) < thres,\ 1374 "not on rest particle %s %s %s %s" % (decay_particle.E, decay_particle.px,decay_particle.py,decay_particle.pz) 1375 1376 self.nexternal += decay_event.nexternal -1 1377 old_scales = list(self.parse_matching_scale()) 1378 if old_scales: 1379 jet_position = sum(1 for i in range(position) if self[i].status==1) 1380 initial_pos = sum(1 for i in range(position) if self[i].status==-1) 1381 self.matched_scale_data.pop(initial_pos+jet_position) 1382 # add the particle with only handling the 4-momenta/mother 1383 # color information will be corrected later. 1384 for particle in decay_event[1:]: 1385 # duplicate particle to avoid border effect 1386 new_particle = Particle(particle, self) 1387 new_particle.event_id = len(self) 1388 self.append(new_particle) 1389 if old_scales: 1390 self.matched_scale_data.append(old_scales[initial_pos+jet_position]) 1391 # compute and assign the new four_momenta 1392 new_momentum = this_4mom.boost(FourMomentum(new_particle)) 1393 new_particle.set_momentum(new_momentum) 1394 # compute the new mother 1395 for tag in ['mother1', 'mother2']: 1396 mother = getattr(particle, tag) 1397 if isinstance(mother, Particle): 1398 mother_id = getattr(particle, tag).event_id 1399 if mother_id == 0: 1400 setattr(new_particle, tag, this_particle) 1401 else: 1402 try: 1403 setattr(new_particle, tag, self[nb_part + mother_id -1]) 1404 except Exception, error: 1405 print error 1406 misc.sprint( self) 1407 misc.sprint(nb_part + mother_id -1) 1408 misc.sprint(tag) 1409 misc.sprint(position, decay_event) 1410 misc.sprint(particle) 1411 misc.sprint(len(self), nb_part + mother_id -1) 1412 raise 1413 elif tag == "mother2" and isinstance(particle.mother1, Particle): 1414 new_particle.mother2 = this_particle 1415 else: 1416 raise Exception, "Something weird happens. Please report it for investigation" 1417 # Need to correct the color information of the particle 1418 # first find the first available color index 1419 max_color=501 1420 for particle in self[:nb_part]: 1421 max_color=max(max_color, particle.color1, particle.color2) 1422 1423 # define a color mapping and assign it: 1424 color_mapping = {} 1425 color_mapping[decay_particle.color1] = this_particle.color1 1426 color_mapping[decay_particle.color2] = this_particle.color2 1427 for particle in self[nb_part:]: 1428 if particle.color1: 1429 if particle.color1 not in color_mapping: 1430 max_color +=1 1431 color_mapping[particle.color1] = max_color 1432 particle.color1 = max_color 1433 else: 1434 particle.color1 = color_mapping[particle.color1] 1435 if particle.color2: 1436 if particle.color2 not in color_mapping: 1437 max_color +=1 1438 color_mapping[particle.color2] = max_color 1439 particle.color2 = max_color 1440 else: 1441 particle.color2 = color_mapping[particle.color2]
1442
1443 - def add_decays(self, pdg_to_decay):
1444 """use auto-recursion""" 1445 1446 pdg_to_decay = dict(pdg_to_decay) 1447 1448 for i,particle in enumerate(self): 1449 if particle.status != 1: 1450 continue 1451 if particle.pdg in pdg_to_decay and pdg_to_decay[particle.pdg]: 1452 one_decay = pdg_to_decay[particle.pdg].pop() 1453 self.add_decay_to_particle(i, one_decay) 1454 return self.add_decays(pdg_to_decay) 1455 return self
1456 1457 1458
1459 - def remove_decay(self, pdg_code=0, event_id=None):
1460 1461 to_remove = [] 1462 if event_id is not None: 1463 to_remove.append(self[event_id]) 1464 1465 if pdg_code: 1466 for particle in self: 1467 if particle.pid == pdg_code: 1468 to_remove.append(particle) 1469 1470 new_event = Event() 1471 # copy first line information + ... 1472 for tag in ['nexternal', 'ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']: 1473 setattr(new_event, tag, getattr(self, tag)) 1474 1475 for particle in self: 1476 if isinstance(particle.mother1, Particle) and particle.mother1 in to_remove: 1477 to_remove.append(particle) 1478 if particle.status == 1: 1479 new_event.nexternal -= 1 1480 continue 1481 elif isinstance(particle.mother2, Particle) and particle.mother2 in to_remove: 1482 to_remove.append(particle) 1483 if particle.status == 1: 1484 new_event.nexternal -= 1 1485 continue 1486 else: 1487 new_event.append(Particle(particle)) 1488 1489 #ensure that the event_id is correct for all_particle 1490 # and put the status to 1 for removed particle 1491 for pos, particle in enumerate(new_event): 1492 particle.event_id = pos 1493 if particle in to_remove: 1494 particle.status = 1 1495 return new_event
1496
1497 - def get_decay(self, pdg_code=0, event_id=None):
1498 1499 to_start = [] 1500 if event_id is not None: 1501 to_start.append(self[event_id]) 1502 1503 elif pdg_code: 1504 for particle in self: 1505 if particle.pid == pdg_code: 1506 to_start.append(particle) 1507 break 1508 1509 new_event = Event() 1510 # copy first line information + ... 1511 for tag in ['ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']: 1512 setattr(new_event, tag, getattr(self, tag)) 1513 1514 # Add the decaying particle 1515 old2new = {} 1516 new_decay_part = Particle(to_start[0]) 1517 new_decay_part.mother1 = None 1518 new_decay_part.mother2 = None 1519 new_decay_part.status = -1 1520 old2new[new_decay_part.event_id] = len(old2new) 1521 new_event.append(new_decay_part) 1522 1523 1524 # add the other particle 1525 for particle in self: 1526 if isinstance(particle.mother1, Particle) and particle.mother1.event_id in old2new\ 1527 or isinstance(particle.mother2, Particle) and particle.mother2.event_id in old2new: 1528 old2new[particle.event_id] = len(old2new) 1529 new_event.append(Particle(particle)) 1530 1531 #ensure that the event_id is correct for all_particle 1532 # and correct the mother1/mother2 by the new reference 1533 nexternal = 0 1534 for pos, particle in enumerate(new_event): 1535 particle.event_id = pos 1536 if particle.mother1: 1537 particle.mother1 = new_event[old2new[particle.mother1.event_id]] 1538 if particle.mother2: 1539 particle.mother2 = new_event[old2new[particle.mother2.event_id]] 1540 if particle.status in [-1,1]: 1541 nexternal +=1 1542 new_event.nexternal = nexternal 1543 1544 return new_event
1545 1546
1547 - def check(self):
1548 """check various property of the events""" 1549 1550 #1. Check that the 4-momenta are conserved 1551 E, px, py, pz = 0,0,0,0 1552 absE, abspx, abspy, abspz = 0,0,0,0 1553 for particle in self: 1554 coeff = 1 1555 if particle.status == -1: 1556 coeff = -1 1557 elif particle.status != 1: 1558 continue 1559 E += coeff * particle.E 1560 absE += abs(particle.E) 1561 px += coeff * particle.px 1562 py += coeff * particle.py 1563 pz += coeff * particle.pz 1564 abspx += abs(particle.px) 1565 abspy += abs(particle.py) 1566 abspz += abs(particle.pz) 1567 # check that relative error is under control 1568 threshold = 5e-7 1569 if E/absE > threshold: 1570 logger.critical(self) 1571 raise Exception, "Do not conserve Energy %s, %s" % (E/absE, E) 1572 if px/abspx > threshold: 1573 logger.critical(self) 1574 raise Exception, "Do not conserve Px %s, %s" % (px/abspx, px) 1575 if py/abspy > threshold: 1576 logger.critical(self) 1577 raise Exception, "Do not conserve Py %s, %s" % (py/abspy, py) 1578 if pz/abspz > threshold: 1579 logger.critical(self) 1580 raise Exception, "Do not conserve Pz %s, %s" % (pz/abspz, pz) 1581 1582 #2. check the color of the event 1583 self.check_color_structure()
1584
1585 - def assign_scale_line(self, line):
1586 """read the line corresponding to global event line 1587 format of the line is: 1588 Nexternal IEVENT WEIGHT SCALE AEW AS 1589 """ 1590 inputs = line.split() 1591 assert len(inputs) == 6 1592 self.nexternal=int(inputs[0]) 1593 self.ievent=int(inputs[1]) 1594 self.wgt=float(inputs[2]) 1595 self.scale=float(inputs[3]) 1596 self.aqed=float(inputs[4]) 1597 self.aqcd=float(inputs[5])
1598
1599 - def get_tag_and_order(self):
1600 """Return the unique tag identifying the SubProcesses for the generation. 1601 Usefull for program like MadSpin and Reweight module.""" 1602 1603 initial, final, order = [], [], [[], []] 1604 for particle in self: 1605 if particle.status == -1: 1606 initial.append(particle.pid) 1607 order[0].append(particle.pid) 1608 elif particle.status == 1: 1609 final.append(particle.pid) 1610 order[1].append(particle.pid) 1611 initial.sort(), final.sort() 1612 tag = (tuple(initial), tuple(final)) 1613 return tag, order
1614
1615 - def get_helicity(self, get_order, allow_reversed=True):
1616 """return a list with the helicities in the order asked for""" 1617 1618 #avoid to modify the input 1619 order = [list(get_order[0]), list(get_order[1])] 1620 out = [9] *(len(order[0])+len(order[1])) 1621 for i, part in enumerate(self): 1622 if part.status == 1: #final 1623 try: 1624 ind = order[1].index(part.pid) 1625 except ValueError, error: 1626 if not allow_reversed: 1627 raise error 1628 else: 1629 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 1630 try: 1631 return self.get_helicity(order, False) 1632 except ValueError: 1633 raise error 1634 position = len(order[0]) + ind 1635 order[1][ind] = 0 1636 elif part.status == -1: 1637 try: 1638 ind = order[0].index(part.pid) 1639 except ValueError, error: 1640 if not allow_reversed: 1641 raise error 1642 else: 1643 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 1644 try: 1645 return self.get_helicity(order, False) 1646 except ValueError: 1647 raise error 1648 1649 position = ind 1650 order[0][ind] = 0 1651 else: #intermediate 1652 continue 1653 out[position] = int(part.helicity) 1654 return out
1655 1656
1657 - def check_color_structure(self):
1658 """check the validity of the color structure""" 1659 1660 #1. check that each color is raised only once. 1661 color_index = collections.defaultdict(int) 1662 for particle in self: 1663 if particle.status in [-1,1]: 1664 if particle.color1: 1665 color_index[particle.color1] +=1 1666 if -7 < particle.pdg < 0: 1667 raise Exception, "anti-quark with color tag" 1668 if particle.color2: 1669 color_index[particle.color2] +=1 1670 if 7 > particle.pdg > 0: 1671 raise Exception, "quark with anti-color tag" 1672 1673 1674 for key,value in color_index.items(): 1675 if value > 2: 1676 print self 1677 print key, value 1678 raise Exception, 'Wrong color_flow' 1679 1680 1681 #2. check that each parent present have coherent color-structure 1682 check = [] 1683 popup_index = [] #check that the popup index are created in a unique way 1684 for particle in self: 1685 mothers = [] 1686 childs = [] 1687 if particle.mother1: 1688 mothers.append(particle.mother1) 1689 if particle.mother2 and particle.mother2 is not particle.mother1: 1690 mothers.append(particle.mother2) 1691 if not mothers: 1692 continue 1693 if (particle.mother1.event_id, particle.mother2.event_id) in check: 1694 continue 1695 check.append((particle.mother1.event_id, particle.mother2.event_id)) 1696 1697 childs = [p for p in self if p.mother1 is particle.mother1 and \ 1698 p.mother2 is particle.mother2] 1699 1700 mcolors = [] 1701 manticolors = [] 1702 for m in mothers: 1703 if m.color1: 1704 if m.color1 in manticolors: 1705 manticolors.remove(m.color1) 1706 else: 1707 mcolors.append(m.color1) 1708 if m.color2: 1709 if m.color2 in mcolors: 1710 mcolors.remove(m.color2) 1711 else: 1712 manticolors.append(m.color2) 1713 ccolors = [] 1714 canticolors = [] 1715 for m in childs: 1716 if m.color1: 1717 if m.color1 in canticolors: 1718 canticolors.remove(m.color1) 1719 else: 1720 ccolors.append(m.color1) 1721 if m.color2: 1722 if m.color2 in ccolors: 1723 ccolors.remove(m.color2) 1724 else: 1725 canticolors.append(m.color2) 1726 for index in mcolors[:]: 1727 if index in ccolors: 1728 mcolors.remove(index) 1729 ccolors.remove(index) 1730 for index in manticolors[:]: 1731 if index in canticolors: 1732 manticolors.remove(index) 1733 canticolors.remove(index) 1734 1735 if mcolors != []: 1736 #only case is a epsilon_ijk structure. 1737 if len(canticolors) + len(mcolors) != 3: 1738 logger.critical(str(self)) 1739 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs]) 1740 else: 1741 popup_index += canticolors 1742 elif manticolors != []: 1743 #only case is a epsilon_ijk structure. 1744 if len(ccolors) + len(manticolors) != 3: 1745 logger.critical(str(self)) 1746 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs]) 1747 else: 1748 popup_index += ccolors 1749 1750 # Check that color popup (from epsilon_ijk) are raised only once 1751 if len(popup_index) != len(set(popup_index)): 1752 logger.critical(self) 1753 raise Exception, "Wrong color flow: identical poping-up index, %s" % (popup_index)
1754
1755 - def __eq__(self, other):
1756 """two event are the same if they have the same momentum. other info are ignored""" 1757 1758 if other is None: 1759 return False 1760 1761 for i,p in enumerate(self): 1762 if p.E != other[i].E: 1763 return False 1764 elif p.pz != other[i].pz: 1765 return False 1766 elif p.px != other[i].px: 1767 return False 1768 elif p.py != other[i].py: 1769 return False 1770 return True
1771 1772
1773 - def __str__(self, event_id=''):
1774 """return a correctly formatted LHE event""" 1775 1776 out="""<event%(event_flag)s> 1777 %(scale)s 1778 %(particles)s 1779 %(comments)s 1780 %(tag)s 1781 %(reweight)s 1782 </event> 1783 """ 1784 if event_id not in ['', None]: 1785 self.eventflag['event'] = str(event_id) 1786 1787 if self.eventflag: 1788 event_flag = ' %s' % ' '.join('%s="%s"' % (k,v) for (k,v) in self.eventflag.items()) 1789 else: 1790 event_flag = '' 1791 1792 if self.nexternal: 1793 scale_str = "%2d %6d %+13.7e %14.8e %14.8e %14.8e" % \ 1794 (self.nexternal,self.ievent,self.wgt,self.scale,self.aqed,self.aqcd) 1795 else: 1796 scale_str = '' 1797 1798 if self.reweight_data: 1799 # check that all key have an order if not add them at the end 1800 if set(self.reweight_data.keys()) != set(self.reweight_order): 1801 self.reweight_order += [k for k in self.reweight_data.keys() \ 1802 if k not in self.reweight_order] 1803 1804 reweight_str = '<rwgt>\n%s\n</rwgt>' % '\n'.join( 1805 '<wgt id=\'%s\'> %+13.7e </wgt>' % (i, float(self.reweight_data[i])) 1806 for i in self.reweight_order) 1807 else: 1808 reweight_str = '' 1809 1810 tag_str = self.tag 1811 if self.matched_scale_data: 1812 tmp_scale = ' '.join(['pt_clust_%i=\"%s\"' % (i+1,v) 1813 for i,v in enumerate(self.matched_scale_data) 1814 if v!=-1]) 1815 if tmp_scale: 1816 tag_str = "<scales %s></scales>%s" % (tmp_scale, self.tag) 1817 1818 if self.syscalc_data: 1819 keys= ['rscale', 'asrwt', ('pdfrwt', 'beam', '1'), ('pdfrwt', 'beam', '2'), 1820 'matchscale', 'totfact'] 1821 sys_str = "<mgrwt>\n" 1822 template = """<%(key)s%(opts)s>%(values)s</%(key)s>\n""" 1823 for k in keys: 1824 if k not in self.syscalc_data: 1825 continue 1826 replace = {} 1827 replace['values'] = self.syscalc_data[k] 1828 if isinstance(k, str): 1829 replace['key'] = k 1830 replace['opts'] = '' 1831 else: 1832 replace['key'] = k[0] 1833 replace['opts'] = ' %s=\"%s\"' % (k[1],k[2]) 1834 sys_str += template % replace 1835 sys_str += "</mgrwt>\n" 1836 reweight_str = sys_str + reweight_str 1837 1838 out = out % {'event_flag': event_flag, 1839 'scale': scale_str, 1840 'particles': '\n'.join([str(p) for p in self]), 1841 'tag': tag_str, 1842 'comments': self.comment, 1843 'reweight': reweight_str} 1844 1845 return re.sub('[\n]+', '\n', out)
1846
1847 - def get_momenta(self, get_order, allow_reversed=True):
1848 """return the momenta vector in the order asked for""" 1849 1850 #avoid to modify the input 1851 order = [list(get_order[0]), list(get_order[1])] 1852 out = [''] *(len(order[0])+len(order[1])) 1853 for i, part in enumerate(self): 1854 if part.status == 1: #final 1855 try: 1856 ind = order[1].index(part.pid) 1857 except ValueError, error: 1858 if not allow_reversed: 1859 raise error 1860 else: 1861 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 1862 try: 1863 return self.get_momenta_str(order, False) 1864 except ValueError: 1865 raise error 1866 position = len(order[0]) + ind 1867 order[1][ind] = 0 1868 elif part.status == -1: 1869 try: 1870 ind = order[0].index(part.pid) 1871 except ValueError, error: 1872 if not allow_reversed: 1873 raise error 1874 else: 1875 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 1876 try: 1877 return self.get_momenta_str(order, False) 1878 except ValueError: 1879 raise error 1880 1881 position = ind 1882 order[0][ind] = 0 1883 else: #intermediate 1884 continue 1885 1886 out[position] = (part.E, part.px, part.py, part.pz) 1887 1888 return out
1889 1890 1891
1892 - def get_ht_scale(self, prefactor=1):
1893 1894 scale = 0 1895 for particle in self: 1896 if particle.status != 1: 1897 continue 1898 p=FourMomentum(particle) 1899 scale += math.sqrt(p.mass_sqr + p.pt**2) 1900 1901 return prefactor * scale
1902 1903
1904 - def get_et_scale(self, prefactor=1):
1905 1906 scale = 0 1907 for particle in self: 1908 if particle.status != 1: 1909 continue 1910 p = FourMomentum(particle) 1911 pt = p.pt 1912 if (pt>0): 1913 scale += p.E*pt/math.sqrt(pt**2+p.pz**2) 1914 1915 return prefactor * scale
1916
1917 - def get_sqrts_scale(self, prefactor=1):
1918 1919 scale = 0 1920 init = [] 1921 for particle in self: 1922 if particle.status == -1: 1923 init.append(FourMomentum(particle)) 1924 if len(init) == 1: 1925 return init[0].mass 1926 elif len(init)==2: 1927 return math.sqrt((init[0]+init[1])**2)
1928 1929 1930 1931
1932 - def get_momenta_str(self, get_order, allow_reversed=True):
1933 """return the momenta str in the order asked for""" 1934 1935 out = self.get_momenta(get_order, allow_reversed) 1936 #format 1937 format = '%.12f' 1938 format_line = ' '.join([format]*4) + ' \n' 1939 out = [format_line % one for one in out] 1940 out = ''.join(out).replace('e','d') 1941 return out
1942
1943 -class WeightFile(EventFile):
1944 """A class to allow to read both gzip and not gzip file. 1945 containing only weight from pythia --generated by SysCalc""" 1946
1947 - def __new__(self, path, mode='r', *args, **opt):
1948 if path.endswith(".gz"): 1949 try: 1950 return gzip.GzipFile.__new__(WeightFileGzip, path, mode, *args, **opt) 1951 except IOError, error: 1952 raise 1953 except Exception, error: 1954 if mode == 'r': 1955 misc.gunzip(path) 1956 return file.__new__(WeightFileNoGzip, path[:-3], mode, *args, **opt) 1957 else: 1958 return file.__new__(WeightFileNoGzip, path, mode, *args, **opt)
1959 1960
1961 - def __init__(self, path, mode='r', *args, **opt):
1962 """open file and read the banner [if in read mode]""" 1963 1964 super(EventFile, self).__init__(path, mode, *args, **opt) 1965 self.banner = '' 1966 if mode == 'r': 1967 line = '' 1968 while '</header>' not in line.lower(): 1969 try: 1970 line = super(EventFile, self).next() 1971 except StopIteration: 1972 self.seek(0) 1973 self.banner = '' 1974 break 1975 if "<event" in line.lower(): 1976 self.seek(0) 1977 self.banner = '' 1978 break 1979 1980 self.banner += line
1981
1982 1983 -class WeightFileGzip(WeightFile, EventFileGzip):
1984 pass
1985
1986 -class WeightFileNoGzip(WeightFile, EventFileNoGzip):
1987 pass
1988
1989 1990 -class FourMomentum(object):
1991 """a convenient object for 4-momenta operation""" 1992
1993 - def __init__(self, obj=0, px=0, py=0, pz=0, E=0):
1994 """initialize the four momenta""" 1995 1996 if obj is 0 and E: 1997 obj = E 1998 1999 if isinstance(obj, (FourMomentum, Particle)): 2000 px = obj.px 2001 py = obj.py 2002 pz = obj.pz 2003 E = obj.E 2004 elif isinstance(obj, (list, tuple)): 2005 assert len(obj) ==4 2006 E = obj[0] 2007 px = obj[1] 2008 py = obj[2] 2009 pz = obj[3] 2010 elif isinstance(obj, str): 2011 obj = [float(i) for i in obj.split()] 2012 assert len(obj) ==4 2013 E = obj[0] 2014 px = obj[1] 2015 py = obj[2] 2016 pz = obj[3] 2017 else: 2018 E =obj 2019 2020 2021 self.E = float(E) 2022 self.px = float(px) 2023 self.py = float(py) 2024 self.pz = float(pz)
2025 2026 @property
2027 - def mass(self):
2028 """return the mass""" 2029 return math.sqrt(self.E**2 - self.px**2 - self.py**2 - self.pz**2)
2030 2031 @property
2032 - def mass_sqr(self):
2033 """return the mass square""" 2034 return self.E**2 - self.px**2 - self.py**2 - self.pz**2
2035 2036 @property
2037 - def pt(self):
2038 return math.sqrt(max(0, self.pt2()))
2039 2040 @property
2041 - def pseudorapidity(self):
2042 norm = math.sqrt(self.px**2 + self.py**2+self.pz**2) 2043 return 0.5* math.log((norm - self.pz) / (norm + self.pz))
2044 2045 @property
2046 - def rapidity(self):
2047 return 0.5* math.log((self.E +self.pz) / (self.E - self.pz))
2048 2049 2050
2051 - def pt2(self):
2052 """ return the pt square """ 2053 2054 return self.px**2 + self.py**2
2055
2056 - def __add__(self, obj):
2057 2058 assert isinstance(obj, FourMomentum) 2059 new = FourMomentum(self.E+obj.E, 2060 self.px + obj.px, 2061 self.py + obj.py, 2062 self.pz + obj.pz) 2063 return new
2064
2065 - def __iadd__(self, obj):
2066 """update the object with the sum""" 2067 self.E += obj.E 2068 self.px += obj.px 2069 self.py += obj.py 2070 self.pz += obj.pz 2071 return self
2072
2073 - def __sub__(self, obj):
2074 2075 assert isinstance(obj, FourMomentum) 2076 new = FourMomentum(self.E-obj.E, 2077 self.px - obj.px, 2078 self.py - obj.py, 2079 self.pz - obj.pz) 2080 return new
2081
2082 - def __isub__(self, obj):
2083 """update the object with the sum""" 2084 self.E -= obj.E 2085 self.px -= obj.px 2086 self.py -= obj.py 2087 self.pz -= obj.pz 2088 return self
2089
2090 - def __mul__(self, obj):
2091 if isinstance(obj, FourMomentum): 2092 return self.E*obj.E - self.px *obj.px - self.py * obj.py - self.pz * obj.pz 2093 elif isinstance(obj, (float, int)): 2094 return FourMomentum(obj*self.E,obj*self.px,obj*self.py,obj*self.pz ) 2095 else: 2096 raise NotImplemented
2097 __rmul__ = __mul__ 2098
2099 - def __pow__(self, power):
2100 assert power in [1,2] 2101 2102 if power == 1: 2103 return FourMomentum(self) 2104 elif power == 2: 2105 return self.mass_sqr
2106
2107 - def __repr__(self):
2108 return 'FourMomentum(%s,%s,%s,%s)' % (self.E, self.px, self.py,self.pz)
2109
2110 - def get_tuple(self):
2111 return (self.E, self.px, self.py,self.pz)
2112
2113 - def boost(self, mom):
2114 """mom 4-momenta is suppose to be given in the rest frame of this 4-momenta. 2115 the output is the 4-momenta in the frame of this 4-momenta 2116 function copied from HELAS routine.""" 2117 2118 2119 pt = self.px**2 + self.py**2 + self.pz**2 2120 if pt: 2121 s3product = self.px * mom.px + self.py * mom.py + self.pz * mom.pz 2122 mass = self.mass 2123 lf = (mom.E + (self.E - mass) * s3product / pt ) / mass 2124 return FourMomentum(E=(self.E*mom.E+s3product)/mass, 2125 px=mom.px + self.px * lf, 2126 py=mom.py + self.py * lf, 2127 pz=mom.pz + self.pz * lf) 2128 else: 2129 return FourMomentum(mom)
2130
2131 - def zboost(self, pboost=None, E=0, pz=0):
2132 """Both momenta should be in the same frame. 2133 The boost perform correspond to the boost required to set pboost at 2134 rest (only z boost applied). 2135 """ 2136 if isinstance(pboost, FourMomentum): 2137 E = pboost.E 2138 pz = pboost.pz 2139 2140 #beta = pz/E 2141 gamma = E / math.sqrt(E**2-pz**2) 2142 gammabeta = pz / math.sqrt(E**2-pz**2) 2143 2144 out = FourMomentum([gamma*self.E - gammabeta*self.pz, 2145 self.px, 2146 self.py, 2147 gamma*self.pz - gammabeta*self.E]) 2148 2149 if abs(out.pz) < 1e-6 * out.E: 2150 out.pz = 0 2151 return out
2152
2153 - def boost_to_restframe(self, pboost):
2154 """apply the boost transformation such that pboost is at rest in the new frame. 2155 First apply a rotation to allign the pboost to the z axis and then use 2156 zboost routine (see above) 2157 """ 2158 2159 2160 2161 2162 # write pboost as (E, p cosT sinF, p sinT sinF, p cosF) 2163 # rotation such that it become (E, 0 , 0 , p ) is 2164 # cosT sinF , -sinT , cosT sinF 2165 # sinT cosF , cosT , sinT sinF 2166 # -sinT , 0 , cosF 2167 p = math.sqrt( pboost.px**2 + pboost.py**2+ pboost.pz**2) 2168 cosF = pboost.pz / p 2169 sinF = math.sqrt(1-cosF**2) 2170 sinT = pboost.py/p/sinF 2171 cosT = pboost.px/p/sinF 2172 2173 out=FourMomentum([self.E, 2174 self.px*cosT*cosF + self.py*sinT*cosF-self.pz*sinF, 2175 -self.px*sinT+ self.py*cosT, 2176 self.px*cosT*sinF + self.py*sinT*sinF + self.pz*cosF 2177 ]) 2178 out = out.zboost(E=pboost.E,pz=p) 2179 return out
2180
2181 2182 2183 2184 -class OneNLOWeight(object):
2185
2186 - def __init__(self, input, real_type=(1,11)):
2187 """ """ 2188 2189 self.real_type = real_type 2190 if isinstance(input, str): 2191 self.parse(input)
2192
2193 - def __str__(self):
2194 2195 out = """ pwgt: %(pwgt)s 2196 born, real : %(born)s %(real)s 2197 pdgs : %(pdgs)s 2198 bjks : %(bjks)s 2199 scales**2, gs: %(scales2)s %(gs)s 2200 born/real related : %(born_related)s %(real_related)s 2201 type / nfks : %(type)s %(nfks)s 2202 to merge : %(to_merge_pdg)s in %(merge_new_pdg)s 2203 ref_wgt : %(ref_wgt)s""" % self.__dict__ 2204 return out
2205 2206
2207 - def parse(self, text):
2208 """parse the line and create the related object""" 2209 #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 2210 #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 2211 # below comment are from Rik description email 2212 data = text.split() 2213 # 1. The first three doubles are, as before, the 'wgt', i.e., the overall event of this 2214 # contribution, and the ones multiplying the log[mu_R/QES] and the log[mu_F/QES] 2215 # stripped of alpha_s and the PDFs. 2216 # from example: 0.274922677249D+01 0.000000000000D+00 0.000000000000D+00 2217 self.pwgt = [float(f) for f in data[:3]] 2218 # 2. The next two doubles are the values of the (corresponding) Born and 2219 # real-emission matrix elements. You can either use these values to check 2220 # that the newly computed original matrix element weights are correct, 2221 # or directly use these so that you don't have to recompute the original weights. 2222 # For contributions for which the real-emission matrix elements were 2223 # not computed, the 2nd of these numbers is zero. The opposite is not true, 2224 # because each real-emission phase-space configuration has an underlying Born one 2225 # (this is not unique, but on our code we made a specific choice here). 2226 # This latter information is useful if the real-emission matrix elements 2227 # are unstable; you can then reweight with the Born instead. 2228 # (see also point 9 below, where the momentum configurations are assigned). 2229 # I don't think this instability is real problem when reweighting the real-emission 2230 # with tree-level matrix elements (as we generally would do), but is important 2231 # when reweighting with loop-squared contributions as we have been doing for gg->H. 2232 # (I'm not sure that reweighting tree-level with loop^2 is something that 2233 # we can do in general, because we don't really know what to do with the 2234 # virtual matrix elements because we cannot generate 2-loop diagrams.) 2235 # from example: 0.770516514633D+01 0.113763730192D+00 2236 self.born = float(data[3]) 2237 self.real = float(data[4]) 2238 # 3. integer: number of external particles of the real-emission configuration (as before) 2239 # from example: 5 2240 self.nexternal = int(data[5]) 2241 # 4. PDG codes corresponding to the real-emission configuration (as before) 2242 # from example: 21 2 -11 12 1 2 2243 self.pdgs = [int(i) for i in data[6:6+self.nexternal]] 2244 flag = 6+self.nexternal # new starting point for the position 2245 # 5. next integer is the power of g_strong in the matrix elements (as before) 2246 # from example: 2 2247 self.qcdpower = int(data[flag]) 2248 # 6. 2 doubles: The bjorken x's used for this contribution (as before) 2249 # from example: 0.52500539D-02 0.30205908D+00 2250 self.bjks = [float(f) for f in data[flag+1:flag+3]] 2251 # 7. 3 doubles: The Ellis-sexton scale, the renormalisation scale and the factorisation scale, all squared, used for this contribution (as before) 2252 # from example: 0.45444066D+04 0.45444066D+04 0.45444066D+04 2253 self.scales2 = [float(f) for f in data[flag+3:flag+6]] 2254 # 8.the value of g_strong 2255 # from example: 0.12520062D+01 2256 self.gs = float(data[flag+6]) 2257 # 9. 2 integers: the corresponding Born and real-emission type kinematics. (in the list of momenta) 2258 # Note that also the Born-kinematics has n+1 particles, with, in general, 2259 # one particle with zero momentum (this is not ALWAYS the case, 2260 # there could also be 2 particles with perfectly collinear momentum). 2261 # To convert this from n+1 to a n particles, you have to sum the momenta 2262 # of the two particles that 'merge', see point 12 below. 2263 # from example: 1 2 2264 self.born_related = int(data[flag+7]) 2265 self.real_related = int(data[flag+8]) 2266 # 10. 1 integer: the 'type'. This is the information you should use to determine 2267 # if to reweight with Born, virtual or real-emission matrix elements. 2268 # (Apart from the possible problems with complicated real-emission matrix elements 2269 # that need to be computed very close to the soft/collinear limits, see point 2 above. 2270 # I guess that for tree-level this is always okay, but when reweighting 2271 # a tree-level contribution with a one-loop squared one, as we do 2272 # for gg->Higgs, this is important). 2273 # type=1 : real-emission: 2274 # type=2 : Born: 2275 # type=3 : integrated counter terms: 2276 # type=4 : soft counter-term: 2277 # type=5 : collinear counter-term: 2278 # type=6 : soft-collinear counter-term: 2279 # type=7 : O(alphaS) expansion of Sudakov factor for NNLL+NLO: 2280 # type=8 : soft counter-term (with n+1-body kin.): 2281 # type=9 : collinear counter-term (with n+1-body kin.): 2282 # type=10: soft-collinear counter-term (with n+1-body kin.): 2283 # type=11: real-emission (with n-body kin.): 2284 # type=12: MC subtraction with n-body kin.: 2285 # type=13: MC subtraction with n+1-body kin.: 2286 # type=14: virtual corrections minus approximate virtual 2287 # type=15: approximate virtual corrections: 2288 # from example: 1 2289 self.type = int(data[flag+9]) 2290 # 11. 1 integer: The FKS configuration for this contribution (not really 2291 # relevant for anything, but is used in checking the reweighting to 2292 # get scale & PDF uncertainties). 2293 # from example: 3 2294 self.nfks = int(data[flag+10]) 2295 # 12. 2 integers: the two particles that should be merged to form the 2296 # born contribution from the real-emission one. Remove these two particles 2297 # from the (ordered) list of PDG codes, and insert a newly created particle 2298 # at the location of the minimum of the two particles removed. 2299 # I.e., if you merge particles 2 and 4, you have to insert the new particle 2300 # as the 2nd particle. And particle 5 and above will be shifted down by one. 2301 # from example: 5 1 2302 self.to_merge_pdg = [int (f) for f in data[flag+11:flag+13]] 2303 # 13. 1 integer: the PDG code of the particle that is created after merging the two particles at point 12. 2304 # from example -1 2305 self.merge_new_pdg = int(data[flag+13]) 2306 # 14. 1 double: the reference number that one should be able to reconstruct 2307 # form the weights (point 1 above) and the rest of the information of this line. 2308 # This is really the contribution to this event as computed by the code 2309 # (and is passed to the integrator). It contains everything. 2310 # from example: 0.110944218997D+05 2311 self.ref_wgt = float(data[flag+14]) 2312 2313 #check the momenta configuration linked to the event 2314 if self.type in self.real_type: 2315 self.momenta_config = self.real_related 2316 else: 2317 self.momenta_config = self.born_related
2318
2319 2320 -class NLO_PARTIALWEIGHT(object):
2321
2322 - class BasicEvent(list):
2323
2324 - def __init__(self, momenta, wgts, event, real_type=(1,11)):
2325 list.__init__(self, momenta) 2326 2327 assert self 2328 self.wgts = wgts 2329 self.pdgs = list(wgts[0].pdgs) 2330 self.event = event 2331 self.real_type = real_type 2332 2333 if wgts[0].momenta_config == wgts[0].born_related: 2334 # need to remove one momenta. 2335 ind1, ind2 = [ind-1 for ind in wgts[0].to_merge_pdg] 2336 if ind1> ind2: 2337 ind1, ind2 = ind2, ind1 2338 if ind1 >= sum(1 for p in event if p.status==-1): 2339 new_p = self[ind1] + self[ind2] 2340 else: 2341 new_p = self[ind1] - self[ind2] 2342 self.pop(ind1) 2343 self.insert(ind1, new_p) 2344 self.pop(ind2) 2345 self.pdgs.pop(ind1) 2346 self.pdgs.insert(ind1, wgts[0].merge_new_pdg ) 2347 self.pdgs.pop(ind2) 2348 # DO NOT update the pdgs of the partial weight! 2349 elif any(w.type in self.real_type for w in wgts): 2350 if any(w.type not in self.real_type for w in wgts): 2351 raise Exception 2352 # check if this is too soft/colinear if so use the born 2353 ind1, ind2 = [ind-1 for ind in wgts[0].to_merge_pdg] 2354 if ind1> ind2: 2355 ind1, ind2 = ind2, ind1 2356 if ind1 >= sum(1 for p in event if p.status==-1): 2357 new_p = self[ind1] + self[ind2] 2358 else: 2359 new_p = self[ind1] - self[ind2] 2360 2361 if __debug__: 2362 ptot = FourMomentum() 2363 for i in xrange(len(self)): 2364 if i <2: 2365 ptot += self[i] 2366 else: 2367 ptot -= self[i] 2368 if ptot.mass_sqr > 1e-16: 2369 misc.sprint(ptot, ptot.mass_sqr) 2370 2371 inv_mass = new_p.mass_sqr 2372 shat = (self[0]+self[1]).mass_sqr 2373 if (abs(inv_mass)/shat < 1e-6): 2374 # misc.sprint(abs(inv_mass)/shat) 2375 self.pop(ind1) 2376 self.insert(ind1, new_p) 2377 self.pop(ind2) 2378 self.pdgs.pop(ind1) 2379 self.pdgs.insert(ind1, wgts[0].merge_new_pdg ) 2380 self.pdgs.pop(ind2) 2381 # DO NOT update the pdgs of the partial weight! 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 # raise Exception 2393 else: 2394 raise Exception
2395
2396 - def get_pdg_code(self):
2397 return self.pdgs
2398
2399 - def get_tag_and_order(self):
2400 """ return the tag and order for this basic event""" 2401 (initial, _), _ = self.event.get_tag_and_order() 2402 order = self.get_pdg_code() 2403 2404 2405 initial, out = order[:len(initial)], order[len(initial):] 2406 initial.sort() 2407 out.sort() 2408 return (tuple(initial), tuple(out)), order
2409
2410 - def get_momenta(self, get_order, allow_reversed=True):
2411 """return the momenta vector in the order asked for""" 2412 2413 #avoid to modify the input 2414 order = [list(get_order[0]), list(get_order[1])] 2415 out = [''] *(len(order[0])+len(order[1])) 2416 pdgs = self.get_pdg_code() 2417 for pos, part in enumerate(self): 2418 if pos < len(get_order[0]): #initial 2419 try: 2420 ind = order[0].index(pdgs[pos]) 2421 except ValueError, error: 2422 if not allow_reversed: 2423 raise error 2424 else: 2425 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 2426 try: 2427 return self.get_momenta(order, False) 2428 except ValueError: 2429 raise error 2430 2431 2432 position = ind 2433 order[0][ind] = 0 2434 else: #final 2435 try: 2436 ind = order[1].index(pdgs[pos]) 2437 except ValueError, error: 2438 if not allow_reversed: 2439 raise error 2440 else: 2441 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 2442 try: 2443 return self.get_momenta(order, False) 2444 except ValueError: 2445 raise error 2446 position = len(order[0]) + ind 2447 order[1][ind] = 0 2448 2449 out[position] = (part.E, part.px, part.py, part.pz) 2450 2451 return out
2452 2453
2454 - def get_helicity(self, *args):
2455 return [9] * len(self)
2456 2457 @property
2458 - def aqcd(self):
2459 return self.event.aqcd
2460
2461 - def get_ht_scale(self, prefactor=1):
2462 2463 scale = 0 2464 for particle in self: 2465 p = particle 2466 scale += math.sqrt(max(0, p.mass_sqr + p.pt**2)) 2467 2468 return prefactor * scale
2469
2470 - def get_et_scale(self, prefactor=1):
2471 2472 scale = 0 2473 for particle in self: 2474 p = particle 2475 pt = p.pt 2476 if (pt>0): 2477 scale += p.E*pt/math.sqrt(pt**2+p.pz**2) 2478 2479 return prefactor * scale
2480 2481
2482 - def get_sqrts_scale(self, event,prefactor=1):
2483 2484 scale = 0 2485 nb_init = 0 2486 for particle in event: 2487 if particle.status == -1: 2488 nb_init+=1 2489 if nb_init == 1: 2490 return self[0].mass 2491 elif nb_init==2: 2492 return math.sqrt((self[0]+self[1])**2)
2493 2494 2495 2496
2497 - def __init__(self, input, event, real_type=(1,11)):
2498 2499 self.real_type = real_type 2500 self.event = event 2501 if isinstance(input, str): 2502 self.parse(input)
2503 2504
2505 - def parse(self, text):
2506 """create the object from the string information (see example below)""" 2507 #0.2344688900d+00 8 2 0 2508 #0.4676614699d+02 0.0000000000d+00 0.0000000000d+00 0.4676614699d+02 2509 #0.4676614699d+02 0.0000000000d+00 0.0000000000d+00 -.4676614699d+02 2510 #0.4676614699d+02 0.2256794794d+02 0.4332148227d+01 0.4073073437d+02 2511 #0.4676614699d+02 -.2256794794d+02 -.4332148227d+01 -.4073073437d+02 2512 #0.0000000000d+00 -.0000000000d+00 -.0000000000d+00 -.0000000000d+00 2513 #0.4780341163d+02 0.0000000000d+00 0.0000000000d+00 0.4780341163d+02 2514 #0.4822581633d+02 0.0000000000d+00 0.0000000000d+00 -.4822581633d+02 2515 #0.4729127470d+02 0.2347155377d+02 0.5153455534d+01 0.4073073437d+02 2516 #0.4627255267d+02 -.2167412893d+02 -.3519736379d+01 -.4073073437d+02 2517 #0.2465400591d+01 -.1797424844d+01 -.1633719155d+01 -.4224046944d+00 2518 #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 2519 #-.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 2520 #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 2521 #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 2522 #-.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 2523 #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 2524 #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 2525 #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 2526 2527 text = text.lower().replace('d','e') 2528 all_line = text.split('\n') 2529 #get global information 2530 first_line ='' 2531 while not first_line.strip(): 2532 first_line = all_line.pop(0) 2533 2534 wgt, nb_wgt, nb_event, _ = first_line.split() 2535 nb_wgt, nb_event = int(nb_wgt), int(nb_event) 2536 2537 momenta = [] 2538 wgts = [] 2539 for line in all_line: 2540 data = line.split() 2541 if len(data) == 4: 2542 p = FourMomentum(data) 2543 momenta.append(p) 2544 elif len(data)>0: 2545 wgt = OneNLOWeight(line, real_type=self.real_type) 2546 wgts.append(wgt) 2547 2548 assert len(wgts) == int(nb_wgt) 2549 2550 get_weights_for_momenta = {} 2551 size_momenta = 0 2552 for wgt in wgts: 2553 if wgt.momenta_config in get_weights_for_momenta: 2554 get_weights_for_momenta[wgt.momenta_config].append(wgt) 2555 else: 2556 if size_momenta == 0: size_momenta = wgt.nexternal 2557 assert size_momenta == wgt.nexternal 2558 get_weights_for_momenta[wgt.momenta_config] = [wgt] 2559 2560 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) 2561 2562 2563 2564 self.cevents = [] 2565 for key in range(1, nb_event+1): 2566 if key in get_weights_for_momenta: 2567 wgt = get_weights_for_momenta[key] 2568 evt = self.BasicEvent(momenta[:size_momenta], get_weights_for_momenta[key], self.event, self.real_type) 2569 self.cevents.append(evt) 2570 momenta = momenta[size_momenta:] 2571 2572 nb_wgt_check = 0 2573 for cevt in self.cevents: 2574 nb_wgt_check += len(cevt.wgts) 2575 assert nb_wgt_check == int(nb_wgt)
2576 2577 2578 2579 if '__main__' == __name__: 2580 2581 lhe = EventFile('unweighted_events.lhe.gz') 2582 #lhe.parsing = False 2583 start = time.time() 2584 for event in lhe: 2585 event.parse_lo_weight() 2586 print 'old method -> ', time.time()-start 2587 lhe = EventFile('unweighted_events.lhe.gz') 2588 #lhe.parsing = False 2589 start = time.time() 2590 for event in lhe: 2591 event.parse_lo_weight_test() 2592 print 'new method -> ', time.time()-start 2593 2594 2595 # Example 1: adding some missing information to the event (here distance travelled) 2596 if False: 2597 start = time 2598 lhe = EventFile('unweighted_events.lhe.gz') 2599 output = open('output_events.lhe', 'w') 2600 #write the banner to the output file 2601 output.write(lhe.banner) 2602 # Loop over all events 2603 for event in lhe: 2604 for particle in event: 2605 # modify particle attribute: here remove the mass 2606 particle.mass = 0 2607 particle.vtim = 2 # The one associate to distance travelled by the particle. 2608 2609 #write this modify event 2610 output.write(str(event)) 2611 output.write('</LesHouchesEvent>\n') 2612 2613 # Example 3: Plotting some variable 2614 if False: 2615 lhe = EventFile('unweighted_events.lhe.gz') 2616 import matplotlib.pyplot as plt 2617 import matplotlib.gridspec as gridspec 2618 nbins = 100 2619 2620 nb_pass = 0 2621 data = [] 2622 for event in lhe: 2623 etaabs = 0 2624 etafinal = 0 2625 for particle in event: 2626 if particle.status==1: 2627 p = FourMomentum(particle) 2628 eta = p.pseudorapidity 2629 if abs(eta) > etaabs: 2630 etafinal = eta 2631 etaabs = abs(eta) 2632 if etaabs < 4: 2633 data.append(etafinal) 2634 nb_pass +=1 2635 2636 2637 print nb_pass 2638 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1]) 2639 gs1.update(wspace=0, hspace=0) # set the spacing between axes. 2640 ax = plt.subplot(gs1[0]) 2641 2642 n, bins, patches = ax.hist(data, nbins, histtype='step', label='original') 2643 ax_c = ax.twinx() 2644 ax_c.set_ylabel('MadGraph5_aMC@NLO') 2645 ax_c.yaxis.set_label_coords(1.01, 0.25) 2646 ax_c.set_yticks(ax.get_yticks()) 2647 ax_c.set_yticklabels([]) 2648 ax.set_xlim([-4,4]) 2649 print "bin value:", n 2650 print "start/end point of bins", bins 2651 plt.axis('on') 2652 plt.xlabel('weight ratio') 2653 plt.show() 2654 2655 2656 # Example 4: More complex plotting example (with ratio plot) 2657 if False: 2658 lhe = EventFile('unweighted_events.lhe') 2659 import matplotlib.pyplot as plt 2660 import matplotlib.gridspec as gridspec 2661 nbins = 100 2662 2663 #mtau, wtau = 45, 5.1785e-06 2664 mtau, wtau = 1.777, 4.027000e-13 2665 nb_pass = 0 2666 data, data2, data3 = [], [], [] 2667 for event in lhe: 2668 nb_pass +=1 2669 if nb_pass > 10000: 2670 break 2671 tau1 = FourMomentum() 2672 tau2 = FourMomentum() 2673 for part in event: 2674 if part.pid in [-12,11,16]: 2675 momenta = FourMomentum(part) 2676 tau1 += momenta 2677 elif part.pid == 15: 2678 tau2 += FourMomentum(part) 2679 2680 if abs((mtau-tau2.mass())/wtau)<1e6 and tau2.mass() >1: 2681 data.append((tau1.mass()-mtau)/wtau) 2682 data2.append((tau2.mass()-mtau)/wtau) 2683 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1]) 2684 gs1.update(wspace=0, hspace=0) # set the spacing between axes. 2685 ax = plt.subplot(gs1[0]) 2686 2687 n, bins, patches = ax.hist(data2, nbins, histtype='step', label='original') 2688 n2, bins2, patches2 = ax.hist(data, bins=bins, histtype='step',label='reconstructed') 2689 import cmath 2690 2691 breit = lambda m : math.sqrt(4*math.pi)*1/(((m)**2-mtau**2)**2+(mtau*wtau)**2)*wtau 2692 2693 data3 = [breit(mtau + x*wtau)*wtau*16867622.6624*50 for x in bins] 2694 2695 ax.plot(bins, data3,label='breit-wigner') 2696 # add the legend 2697 ax.legend() 2698 # add on the right program tag 2699 ax_c = ax.twinx() 2700 ax_c.set_ylabel('MadGraph5_aMC@NLO') 2701 ax_c.yaxis.set_label_coords(1.01, 0.25) 2702 ax_c.set_yticks(ax.get_yticks()) 2703 ax_c.set_yticklabels([]) 2704 2705 plt.title('invariant mass of tau LHE/reconstructed') 2706 plt.axis('on') 2707 ax.set_xticklabels([]) 2708 # ratio plot 2709 ax = plt.subplot(gs1[1]) 2710 data4 = [n[i]/(data3[i]) for i in range(nbins)] 2711 ax.plot(bins, data4 + [0] , 'b') 2712 data4 = [n2[i]/(data3[i]) for i in range(nbins)] 2713 ax.plot(bins, data4 + [0] , 'g') 2714 ax.set_ylim([0,2]) 2715 #remove last y tick to avoid overlap with above plot: 2716 tick = ax.get_yticks() 2717 ax.set_yticks(tick[:-1]) 2718 2719 2720 plt.axis('on') 2721 plt.xlabel('(M - Mtau)/Wtau') 2722 plt.show() 2723