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