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