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