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