Package madgraph :: Package various :: Module lhe_parser
[hide private]
[frames] | no frames]

Source Code for Module madgraph.various.lhe_parser

  1  import collections 
  2  import re 
  3  import misc 
  4   
  5  import logging 
  6  logger = logging.getLogger("madgraph.lhe_parser") 
  7   
8 -class Particle(object):
9 """ """ 10 pattern=re.compile(r'''^\s* 11 (?P<pid>-?\d+)\s+ #PID 12 (?P<status>-?\d+)\s+ #status (1 for output particle) 13 (?P<mother1>-?\d+)\s+ #mother 14 (?P<mother2>-?\d+)\s+ #mother 15 (?P<color1>[+-e.\d]*)\s+ #color1 16 (?P<color2>[+-e.\d]*)\s+ #color2 17 (?P<px>[+-e.\d]*)\s+ #px 18 (?P<py>[+-e.\d]*)\s+ #py 19 (?P<pz>[+-e.\d]*)\s+ #pz 20 (?P<E>[+-e.\d]*)\s+ #E 21 (?P<mass>[+-e.\d]*)\s+ #mass 22 (?P<vtim>[+-e.\d]*)\s+ #displace vertex 23 (?P<helicity>[+-e.\d]*)\s* #helicity 24 ($|(?P<comment>\#[\d|D]*)) #comment/end of string 25 ''',66) #verbose+ignore case 26 27 28
29 - def __init__(self, line=None, event=None):
30 """ """ 31 32 self.event = event 33 self.event_id = len(event) #not yet in the event 34 # LHE information 35 self.pid = 0 36 self.status = 0 37 self.mother1 = None 38 self.mother2 = None 39 self.color1 = 0 40 self.color2 = None 41 self.px = 0 42 self.py = 0 43 self.pz = 0 44 self.E = 0 45 self.mass = 0 46 self.vtim = 0 47 self.helicity = 9 48 self.comment = '' 49 50 if line: 51 self.parse(line)
52
53 - def parse(self, line):
54 """parse the line""" 55 56 obj = self.pattern.search(line) 57 if not obj: 58 raise Exception, 'the line\n%s\n is not a valid format for LHE particle' % line 59 for key, value in obj.groupdict().items(): 60 if key not in ['comment','pid']: 61 setattr(self, key, float(value)) 62 elif key in ['pid']: 63 setattr(self, key, int(value)) 64 else: 65 self.comment = value 66 # assign the mother: 67 if self.mother1: 68 try: 69 self.mother1 = self.event[int(self.mother1) -1] 70 except KeyError: 71 raise Exception, 'Wrong Events format: a daughter appears before it\'s mother' 72 if self.mother2: 73 try: 74 self.mother2 = self.event[int(self.mother2) -1] 75 except KeyError: 76 raise Exception, 'Wrong Events format: a daughter appears before it\'s mother'
77 78 79 80
81 - def __str__(self):
82 """string representing the particles""" 83 return " %8d %2d %4d %4d %4d %4d %+13.7e %+13.7e %+13.7e %14.8e %14.8e %10.4e %10.4e" \ 84 % (self.pid, 85 self.status, 86 self.mother1.event_id+1 if self.mother1 else 0, 87 self.mother2.event_id+1 if self.mother2 else 0, 88 self.color1, 89 self.color2, 90 self.px, 91 self.py, 92 self.pz, 93 self.E, 94 self.mass, 95 self.vtim, 96 self.helicity)
97
98 - def __eq__(self, other):
99 100 if self.pid == other.pid and \ 101 self.status == other.status and \ 102 self.mother1 == other.mother1 and \ 103 self.mother2 == other.mother2 and \ 104 self.color1 == other.color1 and \ 105 self.color2 == other.color2 and \ 106 self.px == other.px and \ 107 self.py == other.py and \ 108 self.pz == other.pz and \ 109 self.E == other.E and \ 110 self.mass == other.mass and \ 111 self.vtim == other.vtim and \ 112 self.helicity == other.helicity: 113 return True 114 return False
115 116 117 118
119 - def __repr__(self):
120 return 'Particle("%s", event=%s)' % (str(self), self.event)
121
122 -class EventFile(file):
123 """ """ 124
125 - def __init__(self, path, mode='r', *args, **opt):
126 """open file and read the banner [if in read mode]""" 127 128 file.__init__(self, path, mode, *args, **opt) 129 self.banner = '' 130 if mode == 'r': 131 line = '' 132 while '</init>' not in line.lower(): 133 try: 134 line = file.next(self) 135 except StopIteration: 136 self.seek(0) 137 self.banner = '' 138 break 139 if "<event>" in line.lower(): 140 self.seek(0) 141 self.banner = '' 142 break 143 144 self.banner += line
145
146 - def get_banner(self):
147 """return a banner object""" 148 import madgraph.various.banner as banner 149 output = banner.Banner() 150 output.read_banner(self.banner.split('\n')) 151 return output
152 153
154 - def next(self):
155 """get next event""" 156 text = '' 157 line = '' 158 mode = 0 159 while '</event>' not in line: 160 line = file.next(self).lower() 161 if '<event>' in line: 162 mode = 1 163 if mode: 164 text += line 165 return Event(text)
166 167
168 -class Event(list):
169 """Class storing a single event information (list of particles + global information)""" 170
171 - def __init__(self, text=None):
172 """The initialization of an empty Event (or one associate to a text file)""" 173 list.__init__(self) 174 175 # First line information 176 self.nexternal = 0 177 self.ievent = 0 178 self.wgt = 0 179 self.aqcd = 0 180 self.scale = 0 181 self.aqed = 0 182 self.aqcd = 0 183 # Weight information 184 self.tag = '' 185 self.comment = '' 186 self.reweight_data ={} 187 188 if text: 189 self.parse(text)
190
191 - def parse(self, text):
192 """Take the input file and create the structured information""" 193 194 text = re.sub(r'</?event>', '', text) # remove pointless tag 195 status = 'first' 196 for line in text.split('\n'): 197 line = line.strip() 198 if not line: 199 continue 200 if line.startswith('#'): 201 self.comment += '%s\n' % line 202 continue 203 if 'first' == status: 204 self.assign_scale_line(line) 205 status = 'part' 206 continue 207 208 if '<' in line: 209 status = 'tag' 210 211 if 'part' == status: 212 self.append(Particle(line, event=self)) 213 else: 214 self.tag += '%s\n' % line
215
216 - def parse_reweight(self):
217 """Parse the re-weight information in order to return a dictionary 218 {key: value}. If no group is define group should be '' """ 219 220 self.reweight_data = {} 221 self.reweight_order = [] 222 start, stop = self.tag.find('<rwgt>'), self.tag.find('</rwgt>') 223 if start != -1 != stop : 224 pattern = re.compile(r'''<\s*wgt id=\'(?P<id>[^\']+)\'\s*>\s*(?P<val>[\ded+-.]*)\s*</wgt>''') 225 data = pattern.findall(self.tag) 226 try: 227 self.reweight_data = dict([(pid, float(value)) for (pid, value) in data 228 if not self.reweight_order.append(pid)]) 229 # the if is to create the order file on the flight 230 except ValueError, error: 231 raise Exception, 'Event File has unvalid weight. %s' % error 232 self.tag = self.tag[:start] + self.tag[stop+7:]
233
234 - def check(self):
235 """check various property of the events""" 236 237 #1. Check that the 4-momenta are conserved 238 E, px, py, pz = 0,0,0,0 239 absE, abspx, abspy, abspz = 0,0,0,0 240 for particle in self: 241 coeff = 1 242 if particle.status == -1: 243 coeff = -1 244 elif particle.status != 1: 245 continue 246 E += coeff * particle.E 247 absE += abs(particle.E) 248 px += coeff * particle.px 249 py += coeff * particle.py 250 pz += coeff * particle.pz 251 abspx += abs(particle.px) 252 abspy += abs(particle.py) 253 abspz += abs(particle.pz) 254 # check that relative error is under control 255 threshold = 5e-11 256 if E/absE > threshold: 257 logger.critical(self) 258 raise Exception, "Do not conserve Energy %s, %s" % (E/absE, E) 259 if px/abspx > threshold: 260 logger.critical(self) 261 raise Exception, "Do not conserve Px %s, %s" % (px/abspx, px) 262 if py/abspy > threshold: 263 logger.critical(self) 264 raise Exception, "Do not conserve Py %s, %s" % (py/abspy, py) 265 if pz/abspz > threshold: 266 logger.critical(self) 267 raise Exception, "Do not conserve Pz %s, %s" % (pz/abspz, pz) 268 269 #2. check the color of the event 270 self.check_color_structure()
271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293
294 - def assign_scale_line(self, line):
295 """read the line corresponding to global event line 296 format of the line is: 297 Nexternal IEVENT WEIGHT SCALE AEW AS 298 """ 299 inputs = line.split() 300 assert len(inputs) == 6 301 self.nexternal=int(inputs[0]) 302 self.ievent=int(inputs[1]) 303 self.wgt=float(inputs[2]) 304 self.scale=float(inputs[3]) 305 self.aqed=float(inputs[4]) 306 self.aqcd=float(inputs[5])
307
308 - def get_tag_and_order(self):
309 """Return the unique tag identifying the SubProcesses for the generation. 310 Usefull for program like MadSpin and Reweight module.""" 311 312 initial, final, order = [], [], [[], []] 313 for particle in self: 314 if particle.status == -1: 315 initial.append(particle.pid) 316 order[0].append(particle.pid) 317 elif particle.status == 1: 318 final.append(particle.pid) 319 order[1].append(particle.pid) 320 initial.sort(), final.sort() 321 tag = (tuple(initial), tuple(final)) 322 return tag, order
323
324 - def check_color_structure(self):
325 """check the validity of the color structure""" 326 327 #1. check that each color is raised only once. 328 color_index = collections.defaultdict(int) 329 for particle in self: 330 if particle.status in [-1,1]: 331 if particle.color1: 332 color_index[particle.color1] +=1 333 if particle.color2: 334 color_index[particle.color2] +=1 335 336 for key,value in color_index.items(): 337 if value > 2: 338 print self 339 print key, value 340 raise Exception, 'Wrong color_flow' 341 342 #2. check that each parent present have coherent color-structure 343 check = [] 344 popup_index = [] #check that the popup index are created in a unique way 345 for particle in self: 346 mothers = [] 347 childs = [] 348 if particle.mother1: 349 mothers.append(particle.mother1) 350 if particle.mother2 and particle.mother2 is not particle.mother1: 351 mothers.append(particle.mother2) 352 if not mothers: 353 continue 354 if (particle.mother1.event_id, particle.mother2.event_id) in check: 355 continue 356 check.append((particle.mother1.event_id, particle.mother2.event_id)) 357 358 childs = [p for p in self if p.mother1 is particle.mother1 and \ 359 p.mother2 is particle.mother2] 360 361 mcolors = [] 362 manticolors = [] 363 for m in mothers: 364 if m.color1: 365 if m.color1 in manticolors: 366 manticolors.remove(m.color1) 367 else: 368 mcolors.append(m.color1) 369 if m.color2: 370 if m.color2 in mcolors: 371 mcolors.remove(m.color2) 372 else: 373 manticolors.append(m.color2) 374 ccolors = [] 375 canticolors = [] 376 for m in childs: 377 if m.color1: 378 if m.color1 in canticolors: 379 canticolors.remove(m.color1) 380 else: 381 ccolors.append(m.color1) 382 if m.color2: 383 if m.color2 in ccolors: 384 ccolors.remove(m.color2) 385 else: 386 canticolors.append(m.color2) 387 for index in mcolors[:]: 388 if index in ccolors: 389 mcolors.remove(index) 390 ccolors.remove(index) 391 for index in manticolors[:]: 392 if index in canticolors: 393 manticolors.remove(index) 394 canticolors.remove(index) 395 396 if mcolors != []: 397 #only case is a epsilon_ijk structure. 398 if len(canticolors) + len(mcolors) != 3: 399 logger.critical(str(self)) 400 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs]) 401 else: 402 popup_index += canticolors 403 elif manticolors != []: 404 #only case is a epsilon_ijk structure. 405 if len(ccolors) + len(manticolors) != 3: 406 logger.critical(str(self)) 407 raise Exception, "Wrong color flow for %s -> %s" ([m.pid for m in mothers], [c.pid for c in childs]) 408 else: 409 popup_index += ccolors 410 411 # Check that color popup (from epsilon_ijk) are raised only once 412 if len(popup_index) != len(set(popup_index)): 413 logger.critical(self) 414 raise Exception, "Wrong color flow: identical poping-up index, %s" % (popup_index)
415 416 417 418 419 420
421 - def __str__(self):
422 """return a correctly formatted LHE event""" 423 424 out="""<event> 425 %(scale)s 426 %(particles)s 427 %(comments)s 428 %(tag)s 429 %(reweight)s 430 </event> 431 """ 432 433 scale_str = "%2d %6d %+13.7e %14.8e %14.8e %14.8e" % \ 434 (self.nexternal,self.ievent,self.wgt,self.scale,self.aqed,self.aqcd) 435 if self.reweight_data: 436 # check that all key have an order if not add them at the end 437 if set(self.reweight_data.keys()) != set(self.reweight_order): 438 self.reweight_order += [k for k in self.reweight_data.keys() \ 439 if k not in self.reweight_order] 440 441 442 reweight_str = '<rwgt>\n%s\n</rwgt>' % '\n'.join( 443 '<wgt id=\'%s\'> %+13.7e </wgt>' % (i, float(self.reweight_data[i])) 444 for i in self.reweight_order) 445 else: 446 reweight_str = '' 447 out = out % {'scale': scale_str, 448 'particles': '\n'.join([str(p) for p in self]), 449 'tag': self.tag, 450 'comments': self.comment, 451 'reweight': reweight_str} 452 return re.sub('[\n]+', '\n', out)
453
454 - def get_momenta_str(self, get_order, allow_reversed=True):
455 """return the momenta str in the order asked for""" 456 457 458 #avoid to modify the input 459 order = [list(get_order[0]), list(get_order[1])] 460 out = [''] *(len(order[0])+len(order[1])) 461 for i, part in enumerate(self): 462 if part.status == 1: #final 463 try: 464 ind = order[1].index(part.pid) 465 except ValueError, error: 466 if not allow_reversed: 467 raise error 468 else: 469 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 470 try: 471 return self.get_momenta_str(order, False) 472 except ValueError: 473 raise error 474 position = len(order[0]) + ind 475 order[1][ind] = 0 476 elif part.status == -1: 477 try: 478 ind = order[0].index(part.pid) 479 except ValueError, error: 480 if not allow_reversed: 481 raise error 482 else: 483 order = [[-i for i in get_order[0]],[-i for i in get_order[1]]] 484 try: 485 return self.get_momenta_str(order, False) 486 except ValueError: 487 raise error 488 489 position = ind 490 order[0][ind] = 0 491 else: #intermediate 492 continue 493 out[position] = '%g %g %g %g \n'% (part.E, part.px, part.py, part.pz) 494 495 out = ''.join(out).replace('e','d') 496 return out
497 498 499 500 if '__main__' == __name__: 501 lhe = EventFile('unweighted_events.lhe') 502 output = open('output_events.lhe', 'w') 503 #write the banner to the output file 504 output.write(lhe.banner) 505 # Loop over all events 506 for event in lhe: 507 for particle in event: 508 # modify particle attribute: here remove the mass 509 particle.mass = 0 510 particle.vtim = 2 # The one associate to distance travelled by the particle. 511 512 #write this modify event 513 output.write(str(event)) 514