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