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 math 
   6  import time 
   7  import os 
   8  import shutil 
   9   
  10  pjoin = os.path.join 
  11   
  12  if '__main__' == __name__: 
  13      import sys 
  14      sys.path.append('../../') 
  15  import misc 
  16  import logging 
  17  import gzip 
  18  logger = logging.getLogger("madgraph.lhe_parser") 
19 20 -class Particle(object):
21 """ """ 22 pattern=re.compile(r'''^\s* 23 (?P<pid>-?\d+)\s+ #PID 24 (?P<status>-?\d+)\s+ #status (1 for output particle) 25 (?P<mother1>-?\d+)\s+ #mother 26 (?P<mother2>-?\d+)\s+ #mother 27 (?P<color1>[+-e.\d]*)\s+ #color1 28 (?P<color2>[+-e.\d]*)\s+ #color2 29 (?P<px>[+-e.\d]*)\s+ #px 30 (?P<py>[+-e.\d]*)\s+ #py 31 (?P<pz>[+-e.\d]*)\s+ #pz 32 (?P<E>[+-e.\d]*)\s+ #E 33 (?P<mass>[+-e.\d]*)\s+ #mass 34 (?P<vtim>[+-e.\d]*)\s+ #displace vertex 35 (?P<helicity>[+-e.\d]*)\s* #helicity 36 ($|(?P<comment>\#[\d|D]*)) #comment/end of string 37 ''',66) #verbose+ignore case 38 39 40
41 - def __init__(self, line=None, event=None):
42 """ """ 43 44 if isinstance(line, Particle): 45 for key in line.__dict__: 46 setattr(self, key, getattr(line, key)) 47 if event: 48 self.event = event 49 return 50 51 self.event = event 52 self.event_id = len(event) #not yet in the event 53 # LHE information 54 self.pid = 0 55 self.status = 0 # -1:initial. 1:final. 2: propagator 56 self.mother1 = None 57 self.mother2 = None 58 self.color1 = 0 59 self.color2 = None 60 self.px = 0 61 self.py = 0 62 self.pz = 0 63 self.E = 0 64 self.mass = 0 65 self.vtim = 0 66 self.helicity = 9 67 self.rwgt = 0 68 self.comment = '' 69 70 if line: 71 self.parse(line)
72 73 @property
74 - def pdg(self):
75 "convenient alias" 76 return self.pid
77
78 - def parse(self, line):
79 """parse the line""" 80 81 obj = self.pattern.search(line) 82 if not obj: 83 raise Exception, 'the line\n%s\n is not a valid format for LHE particle' % line 84 for key, value in obj.groupdict().items(): 85 if key not in ['comment','pid']: 86 setattr(self, key, float(value)) 87 elif key in ['pid', 'mother1', 'mother2']: 88 setattr(self, key, int(value)) 89 else: 90 self.comment = value
91 # Note that mother1/mother2 will be modified by the Event parse function to replace the 92 # integer by a pointer to the actual particle object. 93
94 - def __str__(self):
95 """string representing the particles""" 96 return " %8d %2d %4d %4d %4d %4d %+13.10e %+13.10e %+13.10e %14.10e %14.10e %10.4e %10.4e" \ 97 % (self.pid, 98 self.status, 99 self.mother1.event_id+1 if self.mother1 else 0, 100 self.mother2.event_id+1 if self.mother2 else 0, 101 self.color1, 102 self.color2, 103 self.px, 104 self.py, 105 self.pz, 106 self.E, 107 self.mass, 108 self.vtim, 109 self.helicity)
110
111 - def __eq__(self, other):
112 113 if not isinstance(other, Particle): 114 return False 115 if self.pid == other.pid and \ 116 self.status == other.status and \ 117 self.mother1 == other.mother1 and \ 118 self.mother2 == other.mother2 and \ 119 self.color1 == other.color1 and \ 120 self.color2 == other.color2 and \ 121 self.px == other.px and \ 122 self.py == other.py and \ 123 self.pz == other.pz and \ 124 self.E == other.E and \ 125 self.mass == other.mass and \ 126 self.vtim == other.vtim and \ 127 self.helicity == other.helicity: 128 return True 129 return False
130
131 - def set_momentum(self, momentum):
132 133 self.E = momentum.E 134 self.px = momentum.px 135 self.py = momentum.py 136 self.pz = momentum.pz
137
138 - def add_decay(self, decay_event):
139 """associate to this particle the decay in the associate event""" 140 141 return self.event.add_decay_to_particle(self.event_id, decay_event)
142
143 - def __repr__(self):
144 return 'Particle("%s", event=%s)' % (str(self), self.event)
145
146 147 -class EventFile(object):
148 """A class to allow to read both gzip and not gzip file""" 149
150 - def __new__(self, path, mode='r', *args, **opt):
151 152 if path.endswith(".gz"): 153 try: 154 return gzip.GzipFile.__new__(EventFileGzip, path, mode, *args, **opt) 155 except IOError, error: 156 raise 157 except Exception, error: 158 if mode == 'r': 159 misc.gunzip(path) 160 return file.__new__(EventFileNoGzip, path[:-3], mode, *args, **opt) 161 else: 162 return file.__new__(EventFileNoGzip, path, mode, *args, **opt)
163
164 - def __init__(self, path, mode='r', *args, **opt):
165 """open file and read the banner [if in read mode]""" 166 167 super(EventFile, self).__init__(path, mode, *args, **opt) 168 self.banner = '' 169 if mode == 'r': 170 line = '' 171 while '</init>' not in line.lower(): 172 try: 173 line = super(EventFile, self).next() 174 except StopIteration: 175 self.seek(0) 176 self.banner = '' 177 break 178 if "<event" in line.lower(): 179 self.seek(0) 180 self.banner = '' 181 break 182 183 self.banner += line
184
185 - def get_banner(self):
186 """return a banner object""" 187 import madgraph.various.banner as banner 188 if isinstance(self.banner, banner.Banner): 189 return self.banner 190 191 output = banner.Banner() 192 output.read_banner(self.banner) 193 return output
194 195 @property
196 - def cross(self):
197 """return the cross-section of the file #from the banner""" 198 try: 199 return self._cross 200 except Exception: 201 pass 202 203 onebanner = self.get_banner() 204 self._cross = onebanner.get_cross() 205 return self._cross
206
207 - def __len__(self):
208 if self.closed: 209 return 0 210 if hasattr(self,"len"): 211 return self.len 212 213 init_pos = self.tell() 214 self.seek(0) 215 nb_event=0 216 for _ in self: 217 nb_event +=1 218 self.len = nb_event 219 self.seek(init_pos) 220 return self.len
221 222
223 - def next(self):
224 """get next event""" 225 text = '' 226 line = '' 227 mode = 0 228 while '</event>' not in line: 229 line = super(EventFile, self).next().lower() 230 if '<event' in line: 231 mode = 1 232 text = '' 233 if mode: 234 text += line 235 return Event(text)
236
237 - def initialize_unweighting(self, get_wgt, trunc_error):
238 """ scan once the file to return 239 - the list of the hightest weight (of size trunc_error*NB_EVENT 240 - the cross-section by type of process 241 - the total number of events in the file 242 """ 243 244 # We need to loop over the event file to get some information about the 245 # new cross-section/ wgt of event. 246 self.seek(0) 247 all_wgt = [] 248 cross = collections.defaultdict(int) 249 nb_event = 0 250 for event in self: 251 nb_event +=1 252 wgt = get_wgt(event) 253 cross['all'] += wgt 254 cross['abs'] += abs(wgt) 255 cross[event.ievent] += wgt 256 all_wgt.append(abs(wgt)) 257 # avoid all_wgt to be too large 258 if nb_event % 20000 == 0: 259 all_wgt.sort() 260 # drop the lowest weight 261 nb_keep = max(20, int(nb_event*trunc_error*15)) 262 all_wgt = all_wgt[-nb_keep:] 263 264 #final selection of the interesting weight to keep 265 all_wgt.sort() 266 # drop the lowest weight 267 nb_keep = max(20, int(nb_event*trunc_error*10)) 268 all_wgt = all_wgt[-nb_keep:] 269 self.seek(0) 270 return all_wgt, cross, nb_event
271 272
273 - def unweight(self, outputpath, get_wgt=None, max_wgt=0, trunc_error=0, event_target=0, 274 log_level=logging.INFO):
275 """unweight the current file according to wgt information wgt. 276 which can either be a fct of the event or a tag in the rwgt list. 277 max_wgt allow to do partial unweighting. 278 trunc_error allow for dynamical partial unweighting 279 event_target reweight for that many event with maximal trunc_error. 280 (stop to write event when target is reached) 281 """ 282 if not get_wgt: 283 def weight(event): 284 return event.wgt
285 get_wgt = weight 286 elif isinstance(get_wgt, str): 287 unwgt_name =get_wgt 288 def get_wgt(event): 289 event.parse_reweight() 290 return event.reweight_data[unwgt_name]
291 else: 292 unwgt_name = get_wgt.func_name 293 294 295 # check which weight to write 296 if hasattr(self, "written_weight"): 297 written_weight = lambda x: math.copysign(self.written_weight,float(x)) 298 else: 299 written_weight = lambda x: x 300 301 all_wgt, cross, nb_event = self.initialize_unweighting(get_wgt, trunc_error) 302 303 # function that need to be define on the flight 304 def max_wgt_for_trunc(trunc): 305 """find the weight with the maximal truncation.""" 306 307 xsum = 0 308 i=1 309 while (xsum - all_wgt[-i] * (i-1) <= cross['abs'] * trunc): 310 max_wgt = all_wgt[-i] 311 xsum += all_wgt[-i] 312 i +=1 313 if i == len(all_wgt): 314 break 315 316 return max_wgt 317 # end of the function 318 319 # choose the max_weight 320 if not max_wgt: 321 if trunc_error == 0 or len(all_wgt)<2 or event_target: 322 max_wgt = all_wgt[-1] 323 else: 324 max_wgt = max_wgt_for_trunc(trunc_error) 325 326 # need to modify the banner so load it to an object 327 if self.banner: 328 try: 329 import madgraph 330 except: 331 import internal.banner as banner_module 332 else: 333 import madgraph.various.banner as banner_module 334 if not isinstance(self.banner, banner_module.Banner): 335 banner = self.get_banner() 336 # 1. modify the cross-section 337 banner.modify_init_cross(cross) 338 strategy = banner.get_lha_strategy() 339 # 2. modify the lha strategy 340 if strategy >0: 341 banner.set_lha_strategy(4) 342 else: 343 banner.set_lha_strategy(-4) 344 # 3. add information about change in weight 345 banner["unweight"] = "unweighted by %s" % unwgt_name 346 else: 347 banner = self.banner 348 349 350 # Do the reweighting (up to 20 times if we have target_event) 351 nb_try = 20 352 nb_keep = 0 353 for i in range(nb_try): 354 self.seek(0) 355 if event_target: 356 if i==0: 357 max_wgt = max_wgt_for_trunc(0) 358 else: 359 #guess the correct max_wgt based on last iteration 360 efficiency = nb_keep/nb_event 361 needed_efficiency = event_target/nb_event 362 last_max_wgt = max_wgt 363 needed_max_wgt = last_max_wgt * efficiency / needed_efficiency 364 365 min_max_wgt = max_wgt_for_trunc(trunc_error) 366 max_wgt = max(min_max_wgt, needed_max_wgt) 367 max_wgt = min(max_wgt, all_wgt[-1]) 368 if max_wgt == last_max_wgt: 369 if nb_keep <= event_target: 370 logger.log(log_level+10,"fail to reach target %s", event_target) 371 break 372 else: 373 break 374 375 #create output file (here since we are sure that we have to rewrite it) 376 if outputpath: 377 outfile = EventFile(outputpath, "w") 378 # need to write banner information 379 # need to see what to do with rwgt information! 380 if self.banner and outputpath: 381 banner.write(outfile, close_tag=False) 382 383 # scan the file 384 nb_keep = 0 385 trunc_cross = 0 386 for event in self: 387 r = random.random() 388 wgt = get_wgt(event) 389 if abs(wgt) < r * max_wgt: 390 continue 391 elif wgt > 0: 392 nb_keep += 1 393 event.wgt = written_weight(max(wgt, max_wgt)) 394 395 if abs(wgt) > max_wgt: 396 trunc_cross += abs(wgt) - max_wgt 397 if event_target ==0 or nb_keep <= event_target: 398 if outputpath: 399 outfile.write(str(event)) 400 401 elif wgt < 0: 402 nb_keep += 1 403 event.wgt = -1 * max(abs(wgt), max_wgt) 404 if abs(wgt) > max_wgt: 405 trunc_cross += abs(wgt) - max_wgt 406 if outputpath and (event_target ==0 or nb_keep <= event_target): 407 outfile.write(str(event)) 408 409 if event_target and nb_keep > event_target: 410 if not outputpath: 411 #no outputpath define -> wants only the nb of unweighted events 412 continue 413 elif event_target and i != nb_try-1 and nb_keep >= event_target *1.05: 414 outfile.close() 415 # logger.log(log_level, "Found Too much event %s. Try to reduce truncation" % nb_keep) 416 continue 417 else: 418 outfile.write("</LesHouchesEvents>\n") 419 outfile.close() 420 break 421 elif event_target == 0: 422 if outputpath: 423 outfile.write("</LesHouchesEvents>\n") 424 outfile.close() 425 break 426 elif outputpath: 427 outfile.close() 428 # logger.log(log_level, "Found only %s event. Reduce max_wgt" % nb_keep) 429 430 else: 431 # pass here if event_target > 0 and all the attempt fail. 432 logger.log(log_level+10,"fail to reach target event %s (iteration=%s)", event_target,i) 433 434 # logger.log(log_level, "Final maximum weight used for final "+\ 435 # "unweighting is %s yielding %s events." % (max_wgt,nb_keep)) 436 437 if event_target: 438 nb_events_unweighted = nb_keep 439 nb_keep = min( event_target, nb_keep) 440 else: 441 nb_events_unweighted = nb_keep 442 443 logger.log(log_level, "write %i event (efficiency %.2g %%, truncation %.2g %%) after %i iteration(s)", 444 nb_keep, nb_events_unweighted/nb_event*100, trunc_cross/cross['abs']*100, i) 445 446 #correct the weight in the file if not the correct number of event 447 if nb_keep != event_target and hasattr(self, "written_weight"): 448 written_weight = lambda x: math.copysign(self.written_weight*event_target/nb_keep, float(x)) 449 startfile = EventFile(outputpath) 450 tmpname = pjoin(os.path.dirname(outputpath), "wgtcorrected_"+ os.path.basename(outputpath)) 451 outfile = EventFile(tmpname, "w") 452 outfile.write(startfile.banner) 453 for event in startfile: 454 event.wgt = written_weight(event.wgt) 455 outfile.write(str(event)) 456 outfile.write("</LesHouchesEvents>\n") 457 startfile.close() 458 outfile.close() 459 shutil.move(tmpname, outputpath) 460 461 462 self.max_wgt = max_wgt 463 return nb_keep 464
465 - def apply_fct_on_event(self, *fcts, **opts):
466 """ apply one or more fct on all event. """ 467 468 opt= {"print_step": 2000} 469 opt.update(opts) 470 471 nb_fct = len(fcts) 472 out = [] 473 for i in range(nb_fct): 474 out.append([]) 475 self.seek(0) 476 nb_event = 0 477 for event in self: 478 nb_event += 1 479 if opt["print_step"] and nb_event % opt["print_step"] == 0: 480 if hasattr(self,"len"): 481 logger.info("currently at %s/%s event" % (nb_event, self.len)) 482 else: 483 logger.info("currently at %s event" % nb_event) 484 for i in range(nb_fct): 485 out[i].append(fcts[i](event)) 486 if nb_fct == 1: 487 return out[0] 488 else: 489 return out
490
491 492 -class EventFileGzip(EventFile, gzip.GzipFile):
493 """A way to read/write a gzipped lhef event"""
494
495 -class EventFileNoGzip(EventFile, file):
496 """A way to read a standard event file"""
497
498 -class MultiEventFile(EventFile):
499 """a class to read simultaneously multiple file and read them in mixing them. 500 Unweighting can be done at the same time. 501 The number of events in each file need to be provide in advance 502 (if not provide the file is first read to find that number""" 503
504 - def __new__(cls, start_list=[]):
505 return object.__new__(MultiEventFile)
506
507 - def __init__(self, start_list=[]):
508 """if trunc_error is define here then this allow 509 to only read all the files twice and not three times.""" 510 self.files = [] 511 self.banner = '' 512 self.initial_nb_events = [] 513 self.total_event_in_files = 0 514 self.curr_nb_events = [] 515 self.allcross = [] 516 self.error = [] 517 self.across = [] 518 self.scales = [] 519 if start_list: 520 for p in start_list: 521 self.add(p) 522 self._configure = False
523
524 - def add(self, path, cross, error, across):
525 """ add a file to the pool, across allow to reweight the sum of weight 526 in the file to the given cross-section 527 """ 528 529 if across == 0: 530 # No event linked to this channel -> so no need to include it 531 return 532 533 obj = EventFile(path) 534 if len(self.files) == 0 and not self.banner: 535 self.banner = obj.banner 536 self.curr_nb_events.append(0) 537 self.initial_nb_events.append(0) 538 self.allcross.append(cross) 539 self.across.append(across) 540 self.error.append(error) 541 self.scales.append(1) 542 self.files.append(obj) 543 self._configure = False
544
545 - def __iter__(self):
546 return self
547
548 - def next(self):
549 550 if not self._configure: 551 self.configure() 552 553 remaining_event = self.total_event_in_files - sum(self.curr_nb_events) 554 if remaining_event == 0: 555 raise StopIteration 556 # determine which file need to be read 557 nb_event = random.randint(1, remaining_event) 558 sum_nb=0 559 for i, obj in enumerate(self.files): 560 sum_nb += self.initial_nb_events[i] - self.curr_nb_events[i] 561 if nb_event <= sum_nb: 562 self.curr_nb_events[i] += 1 563 event = obj.next() 564 event.sample_scale = self.scales[i] # for file reweighting 565 return event 566 else: 567 raise Exception
568 569
570 - def define_init_banner(self, wgt):
571 """define the part of the init_banner""" 572 573 if not self.banner: 574 return 575 576 # compute the cross-section of each splitted channel 577 grouped_cross = {} 578 grouped_error = {} 579 for i,ff in enumerate(self.files): 580 filename = ff.name 581 Pdir = [P for P in filename.split(os.path.sep) if P.startswith('P')][-1] 582 group = Pdir.split("_")[0][1:] 583 if group in grouped_cross: 584 grouped_cross[group] += self.allcross[i] 585 grouped_error[group] += self.error[i]**2 586 else: 587 grouped_cross[group] = self.allcross[i] 588 grouped_error[group] = self.error[i]**2 589 590 nb_group = len(grouped_cross) 591 592 # compute the information for the first line 593 try: 594 run_card = self.banner.run_card 595 except: 596 run_card = self.banner.charge_card("run_card") 597 598 init_information = run_card.get_banner_init_information() 599 init_information["nprup"] = nb_group 600 601 if run_card["lhe_version"] < 3: 602 init_information["generator_info"] = "" 603 else: 604 init_information["generator_info"] = "<generator name='MadGraph5_aMC@NLO' version='2.2.1'>please cite 1405.0301 </generator>\n" 605 606 # cross_information: 607 cross_info = "%(cross)e %(error)e %(wgt)e %(id)i" 608 init_information["cross_info"] = [] 609 for id in grouped_cross: 610 conv = {"id": int(id), "cross": grouped_cross[id], "error": math.sqrt(grouped_error[id]), 611 "wgt": wgt} 612 init_information["cross_info"].append( cross_info % conv) 613 init_information["cross_info"] = '\n'.join(init_information["cross_info"]) 614 615 616 617 template_init =\ 618 """ %(idbmup1)i %(idbmup2)i %(ebmup1)e %(ebmup2)e %(pdfgup1)i %(pdfgup2)i %(pdfsup1)i %(pdfsup2)i -3 %(nprup)i 619 %(cross_info)s 620 %(generator_info)s 621 """ 622 623 self.banner["init"] = template_init % init_information
624 625 626
627 - def initialize_unweighting(self, getwgt, trunc_error):
628 """ scan once the file to return 629 - the list of the hightest weight (of size trunc_error*NB_EVENT 630 - the cross-section by type of process 631 - the total number of events in the files 632 In top of that it initialise the information for the next routine 633 to determine how to choose which file to read 634 """ 635 self.seek(0) 636 all_wgt = [] 637 total_event = 0 638 sum_cross = collections.defaultdict(int) 639 for i,f in enumerate(self.files): 640 nb_event = 0 641 # We need to loop over the event file to get some information about the 642 # new cross-section/ wgt of event. 643 cross = collections.defaultdict(int) 644 new_wgt =[] 645 for event in f: 646 nb_event += 1 647 total_event += 1 648 event.sample_scale = 1 649 wgt = getwgt(event) 650 cross['all'] += wgt 651 cross['abs'] += abs(wgt) 652 cross[event.ievent] += wgt 653 new_wgt.append(abs(wgt)) 654 # avoid all_wgt to be too large 655 if nb_event % 20000 == 0: 656 new_wgt.sort() 657 # drop the lowest weight 658 nb_keep = max(20, int(nb_event*trunc_error*15)) 659 new_wgt = new_wgt[-nb_keep:] 660 if nb_event == 0: 661 raise Exception 662 # store the information 663 self.initial_nb_events[i] = nb_event 664 self.scales[i] = self.across[i]/cross['abs'] if self.across[i] else 1 665 #misc.sprint("sum of wgt in event %s is %s. Should be %s => scale %s (nb_event: %s)" 666 # % (i, cross['all'], self.allcross[i], self.scales[i], nb_event)) 667 for key in cross: 668 sum_cross[key] += cross[key]* self.scales[i] 669 all_wgt +=[self.scales[i] * w for w in new_wgt] 670 all_wgt.sort() 671 nb_keep = max(20, int(total_event*trunc_error*10)) 672 all_wgt = all_wgt[-nb_keep:] 673 674 self.total_event_in_files = total_event 675 #final selection of the interesting weight to keep 676 all_wgt.sort() 677 # drop the lowest weight 678 nb_keep = max(20, int(total_event*trunc_error*10)) 679 all_wgt = all_wgt[-nb_keep:] 680 self.seek(0) 681 self._configure = True 682 return all_wgt, sum_cross, total_event
683
684 - def configure(self):
685 686 self._configure = True 687 for i,f in enumerate(self.files): 688 self.initial_nb_events[i] = len(f) 689 self.total_event_in_files = sum(self.initial_nb_events)
690
691 - def __len__(self):
692 693 return len(self.files)
694
695 - def seek(self, pos):
696 """ """ 697 698 if pos !=0: 699 raise Exception 700 for i in range(len(self)): 701 self.curr_nb_events[i] = 0 702 for f in self.files: 703 f.seek(pos)
704
705 - def unweight(self, outputpath, get_wgt, **opts):
706 """unweight the current file according to wgt information wgt. 707 which can either be a fct of the event or a tag in the rwgt list. 708 max_wgt allow to do partial unweighting. 709 trunc_error allow for dynamical partial unweighting 710 event_target reweight for that many event with maximal trunc_error. 711 (stop to write event when target is reached) 712 """ 713 714 if isinstance(get_wgt, str): 715 unwgt_name =get_wgt 716 def get_wgt_multi(event): 717 event.parse_reweight() 718 return event.reweight_data[unwgt_name] * event.sample_scale
719 else: 720 unwgt_name = get_wgt.func_name 721 get_wgt_multi = lambda event: get_wgt(event) * event.sample_scale 722 #define the weighting such that we have built-in the scaling 723 724 if opts['event_target']: 725 new_wgt = sum(self.across)/opts['event_target'] 726 self.define_init_banner(new_wgt) 727 self.written_weight = new_wgt 728 729 return super(MultiEventFile, self).unweight(outputpath, get_wgt_multi, **opts)
730
731 732 -class Event(list):
733 """Class storing a single event information (list of particles + global information)""" 734 735 warning_order = True # raise a warning if the order of the particle are not in accordance of child/mother 736
737 - def __init__(self, text=None):
738 """The initialization of an empty Event (or one associate to a text file)""" 739 list.__init__(self) 740 741 # First line information 742 self.nexternal = 0 743 self.ievent = 0 744 self.wgt = 0 745 self.aqcd = 0 746 self.scale = 0 747 self.aqed = 0 748 self.aqcd = 0 749 # Weight information 750 self.tag = '' 751 self.comment = '' 752 self.reweight_data = {} 753 self.matched_scale_data = None 754 if text: 755 self.parse(text)
756 757
758 - def parse(self, text):
759 """Take the input file and create the structured information""" 760 text = re.sub(r'</?event>', '', text) # remove pointless tag 761 status = 'first' 762 for line in text.split('\n'): 763 line = line.strip() 764 if not line: 765 continue 766 if line.startswith('#'): 767 self.comment += '%s\n' % line 768 continue 769 if "<event" in line: 770 continue 771 772 if 'first' == status: 773 if '<rwgt>' in line: 774 status = 'tag' 775 776 if 'first' == status: 777 self.assign_scale_line(line) 778 status = 'part' 779 continue 780 781 if '<' in line: 782 status = 'tag' 783 784 if 'part' == status: 785 self.append(Particle(line, event=self)) 786 else: 787 self.tag += '%s\n' % line 788 789 # assign the mother: 790 for i,particle in enumerate(self): 791 if self.warning_order: 792 if i < particle.mother1 or i < particle.mother2: 793 logger.warning("Order of particle in the event did not agree with parent/child order. This might be problematic for some code.") 794 Event.warning_order = False 795 796 if particle.mother1: 797 try: 798 particle.mother1 = self[int(particle.mother1) -1] 799 except Exception: 800 logger.warning("WRONG MOTHER INFO %s", self) 801 particle.mother1 = 0 802 if particle.mother2: 803 try: 804 particle.mother2 = self[int(particle.mother2) -1] 805 except Exception: 806 logger.warning("WRONG MOTHER INFO %s", self) 807 particle.mother2 = 0
808 809
810 - def parse_reweight(self):
811 """Parse the re-weight information in order to return a dictionary 812 {key: value}. If no group is define group should be '' """ 813 if self.reweight_data: 814 return self.reweight_data 815 self.reweight_data = {} 816 self.reweight_order = [] 817 start, stop = self.tag.find('<rwgt>'), self.tag.find('</rwgt>') 818 if start != -1 != stop : 819 pattern = re.compile(r'''<\s*wgt id=(?:\'|\")(?P<id>[^\'\"]+)(?:\'|\")\s*>\s*(?P<val>[\ded+-.]*)\s*</wgt>''') 820 data = pattern.findall(self.tag) 821 try: 822 self.reweight_data = dict([(pid, float(value)) for (pid, value) in data 823 if not self.reweight_order.append(pid)]) 824 # the if is to create the order file on the flight 825 except ValueError, error: 826 raise Exception, 'Event File has unvalid weight. %s' % error 827 self.tag = self.tag[:start] + self.tag[stop+7:] 828 return self.reweight_data
829
830 - def parse_matching_scale(self):
831 """Parse the line containing the starting scale for the shower""" 832 833 if self.matched_scale_data is not None: 834 return self.matched_scale_data 835 836 self.matched_scale_data = [] 837 838 839 pattern = re.compile("<scales\s|</scales>") 840 data = re.split(pattern,self.tag) 841 if len(data) == 1: 842 return [] 843 else: 844 tmp = {} 845 start,content, end = data 846 self.tag = "%s%s" % (start, end) 847 pattern = re.compile("pt_clust_(\d*)=\"([\de+-.]*)\"") 848 for id,value in pattern.findall(content): 849 tmp[int(id)] = float(value) 850 851 for i in range(1, len(tmp)+1): 852 self.matched_scale_data.append(tmp[i]) 853 854 return self.matched_scale_data
855 856 857 858 859
860 - def add_decay_to_particle(self, position, decay_event):
861 """define the decay of the particle id by the event pass in argument""" 862 863 this_particle = self[position] 864 #change the status to internal particle 865 this_particle.status = 2 866 this_particle.helicity = 0 867 868 # some usefull information 869 decay_particle = decay_event[0] 870 this_4mom = FourMomentum(this_particle) 871 nb_part = len(self) #original number of particle 872 873 thres = decay_particle.E*1e-10 874 assert max(decay_particle.px, decay_particle.py, decay_particle.pz) < thres,\ 875 "not on rest particle %s %s %s %s" % (decay_particle.E, decay_particle.px,decay_particle.py,decay_particle.pz) 876 877 self.nexternal += decay_event.nexternal -1 878 old_scales = list(self.parse_matching_scale()) 879 if old_scales: 880 self.matched_scale_data.pop(position-2) 881 # add the particle with only handling the 4-momenta/mother 882 # color information will be corrected later. 883 for particle in decay_event[1:]: 884 # duplicate particle to avoid border effect 885 new_particle = Particle(particle, self) 886 new_particle.event_id = len(self) 887 self.append(new_particle) 888 if old_scales: 889 self.matched_scale_data.append(old_scales[position-2]) 890 # compute and assign the new four_momenta 891 new_momentum = this_4mom.boost(FourMomentum(new_particle)) 892 new_particle.set_momentum(new_momentum) 893 # compute the new mother 894 for tag in ['mother1', 'mother2']: 895 mother = getattr(particle, tag) 896 if isinstance(mother, Particle): 897 mother_id = getattr(particle, tag).event_id 898 if mother_id == 0: 899 setattr(new_particle, tag, this_particle) 900 else: 901 setattr(new_particle, tag, self[nb_part + mother_id -1]) 902 elif tag == "mother2" and isinstance(particle.mother1, Particle): 903 new_particle.mother2 = this_particle 904 else: 905 raise Exception, "Something weird happens. Please report it for investigation" 906 # Need to correct the color information of the particle 907 # first find the first available color index 908 max_color=501 909 for particle in self[:nb_part]: 910 max_color=max(max_color, particle.color1, particle.color2) 911 912 # define a color mapping and assign it: 913 color_mapping = {} 914 color_mapping[decay_particle.color1] = this_particle.color1 915 color_mapping[decay_particle.color2] = this_particle.color2 916 for particle in self[nb_part:]: 917 if particle.color1: 918 if particle.color1 not in color_mapping: 919 max_color +=1 920 color_mapping[particle.color1] = max_color 921 particle.color1 = max_color 922 else: 923 particle.color1 = color_mapping[particle.color1] 924 if particle.color2: 925 if particle.color2 not in color_mapping: 926 max_color +=1 927 color_mapping[particle.color2] = max_color 928 particle.color2 = max_color 929 else: 930 particle.color2 = color_mapping[particle.color2]
931 932 933
934 - def remove_decay(self, pdg_code=0, event_id=None):
935 936 to_remove = [] 937 if event_id is not None: 938 to_remove.append(self[event_id]) 939 940 if pdg_code: 941 for particle in self: 942 if particle.pid == pdg_code: 943 to_remove.append(particle) 944 945 new_event = Event() 946 # copy first line information + ... 947 for tag in ['nexternal', 'ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']: 948 setattr(new_event, tag, getattr(self, tag)) 949 950 for particle in self: 951 if isinstance(particle.mother1, Particle) and particle.mother1 in to_remove: 952 to_remove.append(particle) 953 if particle.status == 1: 954 new_event.nexternal -= 1 955 continue 956 elif isinstance(particle.mother2, Particle) and particle.mother2 in to_remove: 957 to_remove.append(particle) 958 if particle.status == 1: 959 new_event.nexternal -= 1 960 continue 961 else: 962 new_event.append(Particle(particle)) 963 964 #ensure that the event_id is correct for all_particle 965 # and put the status to 1 for removed particle 966 for pos, particle in enumerate(new_event): 967 particle.event_id = pos 968 if particle in to_remove: 969 particle.status = 1 970 return new_event
971
972 - def get_decay(self, pdg_code=0, event_id=None):
973 974 to_start = [] 975 if event_id is not None: 976 to_start.append(self[event_id]) 977 978 elif pdg_code: 979 for particle in self: 980 if particle.pid == pdg_code: 981 to_start.append(particle) 982 break 983 984 new_event = Event() 985 # copy first line information + ... 986 for tag in ['ievent', 'wgt', 'aqcd', 'scale', 'aqed','tag','comment']: 987 setattr(new_event, tag, getattr(self, tag)) 988 989 # Add the decaying particle 990 old2new = {} 991 new_decay_part = Particle(to_start[0]) 992 new_decay_part.mother1 = None 993 new_decay_part.mother2 = None 994 new_decay_part.status = -1 995 old2new[new_decay_part.event_id] = len(old2new) 996 new_event.append(new_decay_part) 997 998 999 # add the other particle 1000 for particle in self: 1001 if isinstance(particle.mother1, Particle) and particle.mother1.event_id in old2new\ 1002 or isinstance(particle.mother2, Particle) and particle.mother2.event_id in old2new: 1003 old2new[particle.event_id] = len(old2new) 1004 new_event.append(Particle(particle)) 1005 1006 #ensure that the event_id is correct for all_particle 1007 # and correct the mother1/mother2 by the new reference 1008 nexternal = 0 1009 for pos, particle in enumerate(new_event): 1010 particle.event_id = pos 1011 if particle.mother1: 1012 particle.mother1 = new_event[old2new[particle.mother1.event_id]] 1013 if particle.mother2: 1014 particle.mother2 = new_event[old2new[particle.mother2.event_id]] 1015 if particle.status in [-1,1]: 1016 nexternal +=1 1017 new_event.nexternal = nexternal 1018 1019 return new_event
1020 1021
1022 - def check(self):
1023 """check various property of the events""" 1024 1025 #1. Check that the 4-momenta are conserved 1026 E, px, py, pz = 0,0,0,0 1027 absE, abspx, abspy, abspz = 0,0,0,0 1028 for particle in self: 1029 coeff = 1 1030 if particle.status == -1: 1031 coeff = -1 1032 elif particle.status != 1: 1033 continue 1034 E += coeff * particle.E 1035 absE += abs(particle.E) 1036 px += coeff * particle.px 1037 py += coeff * particle.py 1038 pz += coeff * particle.pz 1039 abspx += abs(particle.px) 1040 abspy += abs(particle.py) 1041 abspz += abs(particle.pz) 1042 # check that relative error is under control 1043 threshold = 5e-7 1044 if E/absE > threshold: 1045 logger.critical(self) 1046 raise Exception, "Do not conserve Energy %s, %s" % (E/absE, E) 1047 if px/abspx > threshold: 1048 logger.critical(self) 1049 raise Exception, "Do not conserve Px %s, %s" % (px/abspx, px) 1050 if py/abspy > threshold: 1051 logger.critical(self) 1052 raise Exception, "Do not conserve Py %s, %s" % (py/abspy, py) 1053 if pz/abspz > threshold: 1054 logger.critical(self) 1055 raise Exception, "Do not conserve Pz %s, %s" % (pz/abspz, pz) 1056 1057 #2. check the color of the event 1058 self.check_color_structure()
1059
1060 - def assign_scale_line(self, line):
1061 """read the line corresponding to global event line 1062 format of the line is: 1063 Nexternal IEVENT WEIGHT SCALE AEW AS 1064 """ 1065 inputs = line.split() 1066 assert len(inputs) == 6 1067 self.nexternal=int(inputs[0]) 1068 self.ievent=int(inputs[1]) 1069 self.wgt=float(inputs[2]) 1070 self.scale=float(inputs[3]) 1071 self.aqed=float(inputs[4]) 1072 self.aqcd=float(inputs[5])
1073
1074 - def get_tag_and_order(self):
1075 """Return the unique tag identifying the SubProcesses for the generation. 1076 Usefull for program like MadSpin and Reweight module.""" 1077 1078 initial, final, order = [], [], [[], []] 1079 for particle in self: 1080 if particle.status == -1: 1081 initial.append(particle.pid) 1082 order[0].append(particle.pid) 1083 elif particle.status == 1: 1084 final.append(particle.pid) 1085 order[1].append(particle.pid) 1086 initial.sort(), final.sort() 1087 tag = (tuple(initial), tuple(final)) 1088 return tag, order
1089
1090 - def get_helicity(self, get_order, allow_reversed=True):
1091 """return a list with the helicities in the order asked for""" 1092 1093 1094 1095 #avoid to modify the input 1096 order = [list(get_order[0]), list(get_order[1])] 1097 out = [9] *(len(order[0])+len(order[1])) 1098 for i, part in enumerate(self): 1099 if part.status == 1: #final 1100 try: 1101 ind = order[1].index(part.pid) 1102 except ValueError, error: 1103 if not allow_reversed: 1104 raise error 1105 else: 1106 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 1107 try: 1108 return self.get_helicity(order, False) 1109 except ValueError: 1110 raise error 1111 position = len(order[0]) + ind 1112 order[1][ind] = 0 1113 elif part.status == -1: 1114 try: 1115 ind = order[0].index(part.pid) 1116 except ValueError, error: 1117 if not allow_reversed: 1118 raise error 1119 else: 1120 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 1121 try: 1122 return self.get_helicity(order, False) 1123 except ValueError: 1124 raise error 1125 1126 position = ind 1127 order[0][ind] = 0 1128 else: #intermediate 1129 continue 1130 out[position] = int(part.helicity) 1131 return out
1132 1133
1134 - def check_color_structure(self):
1135 """check the validity of the color structure""" 1136 1137 #1. check that each color is raised only once. 1138 color_index = collections.defaultdict(int) 1139 for particle in self: 1140 if particle.status in [-1,1]: 1141 if particle.color1: 1142 color_index[particle.color1] +=1 1143 if particle.color2: 1144 color_index[particle.color2] +=1 1145 1146 for key,value in color_index.items(): 1147 if value > 2: 1148 print self 1149 print key, value 1150 raise Exception, 'Wrong color_flow' 1151 1152 #2. check that each parent present have coherent color-structure 1153 check = [] 1154 popup_index = [] #check that the popup index are created in a unique way 1155 for particle in self: 1156 mothers = [] 1157 childs = [] 1158 if particle.mother1: 1159 mothers.append(particle.mother1) 1160 if particle.mother2 and particle.mother2 is not particle.mother1: 1161 mothers.append(particle.mother2) 1162 if not mothers: 1163 continue 1164 if (particle.mother1.event_id, particle.mother2.event_id) in check: 1165 continue 1166 check.append((particle.mother1.event_id, particle.mother2.event_id)) 1167 1168 childs = [p for p in self if p.mother1 is particle.mother1 and \ 1169 p.mother2 is particle.mother2] 1170 1171 mcolors = [] 1172 manticolors = [] 1173 for m in mothers: 1174 if m.color1: 1175 if m.color1 in manticolors: 1176 manticolors.remove(m.color1) 1177 else: 1178 mcolors.append(m.color1) 1179 if m.color2: 1180 if m.color2 in mcolors: 1181 mcolors.remove(m.color2) 1182 else: 1183 manticolors.append(m.color2) 1184 ccolors = [] 1185 canticolors = [] 1186 for m in childs: 1187 if m.color1: 1188 if m.color1 in canticolors: 1189 canticolors.remove(m.color1) 1190 else: 1191 ccolors.append(m.color1) 1192 if m.color2: 1193 if m.color2 in ccolors: 1194 ccolors.remove(m.color2) 1195 else: 1196 canticolors.append(m.color2) 1197 for index in mcolors[:]: 1198 if index in ccolors: 1199 mcolors.remove(index) 1200 ccolors.remove(index) 1201 for index in manticolors[:]: 1202 if index in canticolors: 1203 manticolors.remove(index) 1204 canticolors.remove(index) 1205 1206 if mcolors != []: 1207 #only case is a epsilon_ijk structure. 1208 if len(canticolors) + len(mcolors) != 3: 1209 logger.critical(str(self)) 1210 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs]) 1211 else: 1212 popup_index += canticolors 1213 elif manticolors != []: 1214 #only case is a epsilon_ijk structure. 1215 if len(ccolors) + len(manticolors) != 3: 1216 logger.critical(str(self)) 1217 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs]) 1218 else: 1219 popup_index += ccolors 1220 1221 # Check that color popup (from epsilon_ijk) are raised only once 1222 if len(popup_index) != len(set(popup_index)): 1223 logger.critical(self) 1224 raise Exception, "Wrong color flow: identical poping-up index, %s" % (popup_index)
1225
1226 - def __str__(self):
1227 """return a correctly formatted LHE event""" 1228 1229 out="""<event> 1230 %(scale)s 1231 %(particles)s 1232 %(comments)s 1233 %(tag)s 1234 %(reweight)s 1235 </event> 1236 """ 1237 1238 scale_str = "%2d %6d %+13.7e %14.8e %14.8e %14.8e" % \ 1239 (self.nexternal,self.ievent,self.wgt,self.scale,self.aqed,self.aqcd) 1240 if self.reweight_data: 1241 # check that all key have an order if not add them at the end 1242 if set(self.reweight_data.keys()) != set(self.reweight_order): 1243 self.reweight_order += [k for k in self.reweight_data.keys() \ 1244 if k not in self.reweight_order] 1245 1246 reweight_str = '<rwgt>\n%s\n</rwgt>' % '\n'.join( 1247 '<wgt id=\'%s\'> %+13.7e </wgt>' % (i, float(self.reweight_data[i])) 1248 for i in self.reweight_order) 1249 else: 1250 reweight_str = '' 1251 1252 tag_str = self.tag 1253 if self.matched_scale_data: 1254 tag_str = "<scales %s></scales>%s" % ( 1255 ' '.join(['pt_clust_%i=\"%s\"' % (i,v) 1256 for i,v in enumerate(self.matched_scale_data)]), 1257 self.tag) 1258 1259 out = out % {'scale': scale_str, 1260 'particles': '\n'.join([str(p) for p in self]), 1261 'tag': tag_str, 1262 'comments': self.comment, 1263 'reweight': reweight_str} 1264 return re.sub('[\n]+', '\n', out)
1265
1266 - def get_momenta_str(self, get_order, allow_reversed=True):
1267 """return the momenta str in the order asked for""" 1268 1269 1270 #avoid to modify the input 1271 order = [list(get_order[0]), list(get_order[1])] 1272 out = [''] *(len(order[0])+len(order[1])) 1273 for i, part in enumerate(self): 1274 if part.status == 1: #final 1275 try: 1276 ind = order[1].index(part.pid) 1277 except ValueError, error: 1278 if not allow_reversed: 1279 raise error 1280 else: 1281 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 1282 try: 1283 return self.get_momenta_str(order, False) 1284 except ValueError: 1285 raise error 1286 position = len(order[0]) + ind 1287 order[1][ind] = 0 1288 elif part.status == -1: 1289 try: 1290 ind = order[0].index(part.pid) 1291 except ValueError, error: 1292 if not allow_reversed: 1293 raise error 1294 else: 1295 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 1296 try: 1297 return self.get_momenta_str(order, False) 1298 except ValueError: 1299 raise error 1300 1301 position = ind 1302 order[0][ind] = 0 1303 else: #intermediate 1304 continue 1305 format = '%.12f' 1306 format_line = ' '.join([format]*4) + ' \n' 1307 out[position] = format_line % (part.E, part.px, part.py, part.pz) 1308 1309 out = ''.join(out).replace('e','d') 1310 return out
1311
1312 -class WeightFile(EventFile):
1313 """A class to allow to read both gzip and not gzip file. 1314 containing only weight from pythia --generated by SysCalc""" 1315
1316 - def __new__(self, path, mode='r', *args, **opt):
1317 if path.endswith(".gz"): 1318 try: 1319 return gzip.GzipFile.__new__(WeightFileGzip, path, mode, *args, **opt) 1320 except IOError, error: 1321 raise 1322 except Exception, error: 1323 if mode == 'r': 1324 misc.gunzip(path) 1325 return file.__new__(WeightFileNoGzip, path[:-3], mode, *args, **opt) 1326 else: 1327 return file.__new__(WeightFileNoGzip, path, mode, *args, **opt)
1328 1329
1330 - def __init__(self, path, mode='r', *args, **opt):
1331 """open file and read the banner [if in read mode]""" 1332 1333 super(EventFile, self).__init__(path, mode, *args, **opt) 1334 self.banner = '' 1335 if mode == 'r': 1336 line = '' 1337 while '</header>' not in line.lower(): 1338 try: 1339 line = super(EventFile, self).next() 1340 except StopIteration: 1341 self.seek(0) 1342 self.banner = '' 1343 break 1344 if "<event" in line.lower(): 1345 self.seek(0) 1346 self.banner = '' 1347 break 1348 1349 self.banner += line
1350
1351 1352 -class WeightFileGzip(WeightFile, EventFileGzip):
1353 pass
1354
1355 -class WeightFileNoGzip(WeightFile, EventFileNoGzip):
1356 pass
1357
1358 1359 -class FourMomentum(object):
1360 """a convenient object for 4-momenta operation""" 1361
1362 - def __init__(self, obj=0, px=0, py=0, pz=0, E=0):
1363 """initialize the four momenta""" 1364 1365 if obj is 0 and E: 1366 obj = E 1367 1368 if isinstance(obj, (FourMomentum, Particle)): 1369 px = obj.px 1370 py = obj.py 1371 pz = obj.pz 1372 E = obj.E 1373 else: 1374 E =obj 1375 1376 1377 self.E = E 1378 self.px = px 1379 self.py = py 1380 self.pz =pz
1381 1382 @property
1383 - def mass(self):
1384 """return the mass""" 1385 return math.sqrt(self.E**2 - self.px**2 - self.py**2 - self.pz**2)
1386
1387 - def mass_sqr(self):
1388 """return the mass square""" 1389 return self.E**2 - self.px**2 - self.py**2 - self.pz**2
1390 1391 @property
1392 - def pt(self):
1393 return math.sqrt(max(0, self.pt2()))
1394 1395 @property
1396 - def pseudorapidity(self):
1397 norm = math.sqrt(self.px**2 + self.py**2+self.pz**2) 1398 return 0.5* math.log((norm - self.pz) / (norm + self.pz))
1399
1400 - def pt2(self):
1401 """ return the pt square """ 1402 1403 return self.px**2 + self.py**2
1404
1405 - def __add__(self, obj):
1406 1407 assert isinstance(obj, FourMomentum) 1408 new = FourMomentum(self.E+obj.E, 1409 self.px + obj.px, 1410 self.py + obj.py, 1411 self.pz + obj.pz) 1412 return new
1413
1414 - def __iadd__(self, obj):
1415 """update the object with the sum""" 1416 self.E += obj.E 1417 self.px += obj.px 1418 self.py += obj.py 1419 self.pz += obj.pz 1420 return self
1421
1422 - def __pow__(self, power):
1423 assert power in [1,2] 1424 1425 if power == 1: 1426 return FourMomentum(self) 1427 elif power == 2: 1428 return self.mass_sqr()
1429
1430 - def boost(self, mom):
1431 """mom 4-momenta is suppose to be given in the rest frame of this 4-momenta. 1432 the output is the 4-momenta in the frame of this 4-momenta 1433 function copied from HELAS routine.""" 1434 1435 1436 pt = self.px**2 + self.py**2 + self.pz**2 1437 if pt: 1438 s3product = self.px * mom.px + self.py * mom.py + self.pz * mom.pz 1439 mass = self.mass 1440 lf = (mom.E + (self.E - mass) * s3product / pt ) / mass 1441 return FourMomentum(E=(self.E*mom.E+s3product)/mass, 1442 px=mom.px + self.px * lf, 1443 py=mom.py + self.py * lf, 1444 pz=mom.pz + self.pz * lf) 1445 else: 1446 return FourMomentum(mom)
1447 1448 1449 if '__main__' == __name__: 1450 1451 # Example 1: adding some missing information to the event (here distance travelled) 1452 if False: 1453 lhe = EventFile('unweighted_events.lhe.gz') 1454 output = open('output_events.lhe', 'w') 1455 #write the banner to the output file 1456 output.write(lhe.banner) 1457 # Loop over all events 1458 for event in lhe: 1459 for particle in event: 1460 # modify particle attribute: here remove the mass 1461 particle.mass = 0 1462 particle.vtim = 2 # The one associate to distance travelled by the particle. 1463 1464 #write this modify event 1465 output.write(str(event)) 1466 output.write('</LesHouchesEvent>\n') 1467 1468 1469 1470 # Example 3: Plotting some variable 1471 if True: 1472 lhe = EventFile('unweighted_events.lhe.gz') 1473 import matplotlib.pyplot as plt 1474 import matplotlib.gridspec as gridspec 1475 nbins = 100 1476 1477 nb_pass = 0 1478 data = [] 1479 for event in lhe: 1480 etaabs = 0 1481 etafinal = 0 1482 for particle in event: 1483 if particle.status==1: 1484 p = FourMomentum(particle) 1485 eta = p.pseudorapidity 1486 if abs(eta) > etaabs: 1487 etafinal = eta 1488 etaabs = abs(eta) 1489 if etaabs < 4: 1490 data.append(etafinal) 1491 nb_pass +=1 1492 1493 1494 print nb_pass 1495 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1]) 1496 gs1.update(wspace=0, hspace=0) # set the spacing between axes. 1497 ax = plt.subplot(gs1[0]) 1498 1499 n, bins, patches = ax.hist(data, nbins, histtype='step', label='original') 1500 ax_c = ax.twinx() 1501 ax_c.set_ylabel('MadGraph5_aMC@NLO') 1502 ax_c.yaxis.set_label_coords(1.01, 0.25) 1503 ax_c.set_yticks(ax.get_yticks()) 1504 ax_c.set_yticklabels([]) 1505 ax.set_xlim([-4,4]) 1506 print "bin value:", n 1507 print "start/end point of bins", bins 1508 plt.axis('on') 1509 plt.xlabel('weight ratio') 1510 plt.show() 1511 1512 1513 # Example 4: More complex plotting example (with ratio plot) 1514 if False: 1515 lhe = EventFile('unweighted_events.lhe') 1516 import matplotlib.pyplot as plt 1517 import matplotlib.gridspec as gridspec 1518 nbins = 100 1519 1520 #mtau, wtau = 45, 5.1785e-06 1521 mtau, wtau = 1.777, 4.027000e-13 1522 nb_pass = 0 1523 data, data2, data3 = [], [], [] 1524 for event in lhe: 1525 nb_pass +=1 1526 if nb_pass > 10000: 1527 break 1528 tau1 = FourMomentum() 1529 tau2 = FourMomentum() 1530 for part in event: 1531 if part.pid in [-12,11,16]: 1532 momenta = FourMomentum(part) 1533 tau1 += momenta 1534 elif part.pid == 15: 1535 tau2 += FourMomentum(part) 1536 1537 if abs((mtau-tau2.mass())/wtau)<1e6 and tau2.mass() >1: 1538 data.append((tau1.mass()-mtau)/wtau) 1539 data2.append((tau2.mass()-mtau)/wtau) 1540 gs1 = gridspec.GridSpec(2, 1, height_ratios=[5,1]) 1541 gs1.update(wspace=0, hspace=0) # set the spacing between axes. 1542 ax = plt.subplot(gs1[0]) 1543 1544 n, bins, patches = ax.hist(data2, nbins, histtype='step', label='original') 1545 n2, bins2, patches2 = ax.hist(data, bins=bins, histtype='step',label='reconstructed') 1546 import cmath 1547 1548 breit = lambda m : math.sqrt(4*math.pi)*1/(((m)**2-mtau**2)**2+(mtau*wtau)**2)*wtau 1549 1550 data3 = [breit(mtau + x*wtau)*wtau*16867622.6624*50 for x in bins] 1551 1552 ax.plot(bins, data3,label='breit-wigner') 1553 # add the legend 1554 ax.legend() 1555 # add on the right program tag 1556 ax_c = ax.twinx() 1557 ax_c.set_ylabel('MadGraph5_aMC@NLO') 1558 ax_c.yaxis.set_label_coords(1.01, 0.25) 1559 ax_c.set_yticks(ax.get_yticks()) 1560 ax_c.set_yticklabels([]) 1561 1562 plt.title('invariant mass of tau LHE/reconstructed') 1563 plt.axis('on') 1564 ax.set_xticklabels([]) 1565 # ratio plot 1566 ax = plt.subplot(gs1[1]) 1567 data4 = [n[i]/(data3[i]) for i in range(nbins)] 1568 ax.plot(bins, data4 + [0] , 'b') 1569 data4 = [n2[i]/(data3[i]) for i in range(nbins)] 1570 ax.plot(bins, data4 + [0] , 'g') 1571 ax.set_ylim([0,2]) 1572 #remove last y tick to avoid overlap with above plot: 1573 tick = ax.get_yticks() 1574 ax.set_yticks(tick[:-1]) 1575 1576 1577 plt.axis('on') 1578 plt.xlabel('(M - Mtau)/Wtau') 1579 plt.show() 1580