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

Source Code for Module madgraph.various.hepmc_parser

  1  from __future__ import division 
  2   
  3  from __future__ import absolute_import 
  4  from __future__ import print_function 
  5  import gzip 
  6  from six.moves import range 
  7  import os  
  8   
  9  if '__main__' == __name__: 
 10      import sys 
 11      root = os.path.dirname(__file__) 
 12      sys.path.append(os.path.realpath(os.path.join(root, '../../'))) 
 13      if os.path.basename(root) == 'internal': 
 14              __package__ = "internal" 
 15              sys.path.append(os.path.dirname(root)) 
 16              import internal 
 17      else: 
 18          __package__ = "madgraph.various" 
 19       
 20   
 21  try:  
 22      import madgraph 
 23  except ImportError as error: 
 24      print( sys.path) 
 25      print(error) 
 26      from . import misc 
 27  else: 
 28      import madgraph.various.misc as misc 
 29       
 30  import os 
 31  import logging 
32 33 -class HEPMC_Particle(object):
34
35 - def __init__(self, text=None, event=None):
36 37 self.barcode = 0 38 self.pdg = 0 39 self.px = 0 40 self.py = 0 41 self.pz = 0 42 self.E = 0 43 self.mass = 0 44 self.status = 0 45 self.polarization_theta = 0 46 self.polarization_phi = 0 47 self.vertex_barcode = 0 #vertex on which this particle is incoming 48 self.nb_flow_list = 0 49 self.flows = [] 50 51 if text: 52 self.parse(text, event)
53 54 @property
55 - def pdg_code(self):
56 return self.pdg
57 58 pid = pdg_code 59 60 @property
61 - def helicity(self):
62 return 9
63
64 - def parse(self,line=None, event=None):
65 """ P 3 -2 0 0 3.0332529367341937e+01 3.0332529367341937e+01 0 21 0 0 -3 1 2 501""" 66 67 data = line.split() 68 69 self.barcode = int(data[1]) # 3 70 self.pdg = int(data[2]) #-2 71 self.px = float(data[3]) #0 72 self.py = float(data[4]) #0 73 self.pz = float(data[5]) #30.3 74 self.E = float(data[6]) # 30.3 75 self.mass = float(data[7]) # 0 76 self.status = int(data[8]) # 21 77 self.polarization_theta = float(data[9]) #0 78 self.polarization_phi = float(data[10]) #0 79 self.vertex_barcode = float(data[11]) #-3 vertex on which this particle is incoming 80 self.nb_flow_list = int(data[12]) # 1 81 self.flows = [(int(data[13+2*i]),int(data[13+2*i+1])) 82 for i in range(self.nb_flow_list)] # 2 501 83 84 if event: 85 event.curr_vertex.add_outcoming(self)
86
87 - def __str__(self):
88 """P 3 -2 0 0 3.0332529367341937e+01 3.0332529367341937e+01 0 21 0 0 -3 1 2 501""" 89 90 start = """P %i %i %17.16e %17.16e %17.16e %17.16e %17.16e %i %17.16e %17.16e %i %i %s\n""" %\ 91 (self.barcode, self.pdg, self.px, self.py, self.pz, self.E, self.mass, 92 self.status, self.polarization_theta, self.polarization_phi, 93 self.vertex_barcode, self.nb_flow_list, ' '.join("%i %i" % f for f in self.flows)) 94 95 96 return start.replace("%17.16e" % 0, '0')
97
98 99 100 101 -class HEPMC_Vertex(object):
102
103 - def __init__(self, text=None, event=None):
104 105 self.barcode = 0 106 self.id = 0 107 self.x = 0 108 self.y = 0 109 self.z = 0 110 self.ctau = 0 111 self.nb_orphan = 0 112 self.nb_outgoing = 0 113 self.nb_weight = 0 114 self.weights = [] 115 self.incoming = [] 116 self.outcoming = [] 117 118 119 if text: 120 self.parse(text,event)
121
122 - def parse(self, line, event=None):
123 """V -8 0 0 0 0 0 0 2 0""" 124 125 data = line.split() 126 self.barcode = int(data[1]) 127 self.id = float(data[2]) 128 self.x = float(data[3]) 129 self.y = float(data[4]) 130 self.z = float(data[5]) 131 self.ctau = float(data[6]) 132 self.nb_orphan = int(data[7]) 133 self.nb_outgoing = int(data[8]) 134 self.nb_weight = int(data[9]) 135 self.weights = [float(data[10+i]) for i in range(self.nb_weight)] 136 if event: 137 event.vertex[self.barcode] = self
138
139 - def add_incoming(self, particle):
140 self.incoming.append(particle)
141
142 - def add_outcoming(self, particle):
143 self.outcoming.append(particle)
144
145 -class HEPMC_Event(object):
146
147 - def __init__(self, text=None):
148 """The initialization of an empty Event (or one associate to a text file)""" 149 # 150 self.particles = {} #barcode to object 151 self.vertex = {} #barcode to object 152 153 # First line information (E line) 154 self.event_id = 0 155 self.nb_interaction = 0 156 self.scale = 0. 157 self.alphas = 0. 158 self.alphaew = 0. 159 self.process_id = 0 160 self.barcode_vertex =0 161 self.nb_vertex = 0 162 self.barcode_beam1 = 0 163 self.barcode_beam2 = 0 164 self.nb_random_state = 0 165 self.randoms = [] 166 self.nb_weight = 0 167 self.weights = [] 168 169 # not parse container (so far) 170 self.N = '' 171 self.U = '' 172 self.C = '' 173 self.H = '' 174 self.F = '' 175 176 if text: 177 self.parse(text)
178 179 @property
180 - def wgt(self):
181 if self.weights: 182 return self.weights[0] 183 else: 184 return 0.
185 @wgt.setter
186 - def wgt(self, value):
187 self.nb_weight = 1 188 self.weights = [value]
189 190
191 - def parse(self, text):
192 193 for line in text.split('\n'): 194 if not line: 195 continue 196 if line[0] == 'P': 197 P = HEPMC_Particle(line, self) 198 self.add_particle(P) 199 elif line[0] == 'V': 200 V = HEPMC_Vertex(line, self) 201 self.curr_vertex = V 202 self.add_vertex(V) 203 elif line[0] in ['E', 'N', 'U', 'H','F','C']: 204 getattr(self, 'parse_%s' % line[0])(line) 205 else: 206 self.comment = '%s%s\n' % (self.comment,line) 207 208 # add the information about incoming particle 209 for particle in self: 210 try: 211 self.vertex[particle.vertex_barcode].add_incoming(particle) 212 except KeyError: 213 if particle.vertex_barcode == 0: 214 continue 215 raise
216
217 - def parse_E(self,line):
218 """E 249 -1 -1.0000000000000000e+00 -1.0000000000000000e+00 -1.0000000000000000e+00 0 0 462 1 2 0 1 8.2247251000000005e-22""" 219 220 data = line.split() 221 self.event_id = int(data[1]) 222 self.nb_interaction = int(data[2]) 223 self.scale = float(data[3]) 224 self.alphas = float(data[4]) 225 self.alphaew = float(data[5]) 226 self.process_id = int(data[6]) 227 self.barcode_vertex= int(data[7]) 228 self.nb_vertex = int(data[8]) 229 self.barcode_beam1 = int(data[9]) 230 self.barcode_beam2 = int(data[10]) 231 self.nb_random_state = int(data[11]) 232 self.randoms = [float(data[12+i]) for i in range(self.nb_random_state)] 233 self.nb_weight = int(data[12+self.nb_random_state]) 234 self.weights = [float(data[13+self.nb_random_state+i]) 235 for i in range(self.nb_weight)]
236
237 - def parse_N(self,line):
238 """just keep the information so far""" 239 self.N = '%s\n' % line
240 - def parse_U(self,line):
241 self.U = '%s\n' % line
242 - def parse_H(self,line):
243 self.H = '%s\n' % line
244 - def parse_F(self,line):
245 self.F = '%s\n' % line
246 - def parse_C(self,line):
247 self.C = '%s\n' % line
248
249 - def __iter__(self):
250 return list(self.particles.values()).__iter__()
251 252 #def __next__(self): 253 # 254 # self.particles.__next__() 255
256 - def add_vertex(self, V):
257 self.vertex[V.barcode] = V
258
259 - def add_particle(self, P):
260 self.particles[P.barcode] = P
261
262 -class HEPMC_EventFile(object):
263 264
265 - def __init__(self, path, mode='r', *args, **opt):
266 """open file and read the banner [if in read mode]""" 267 268 if mode in ['r','rb']: 269 mode ='r' 270 self.mode = mode 271 272 if not path.endswith(".gz"): 273 self.file = open(path, mode, *args, **opt) 274 elif mode == 'r' and not os.path.exists(path) and os.path.exists(path[:-3]): 275 self.file = open(path[:-3], mode, *args, **opt) 276 path = path[:-3] 277 else: 278 try: 279 self.file = gzip.GzipFile(path, mode, *args, **opt) 280 self.zip_mode =True 281 except IOError as error: 282 raise 283 except Exception as error: 284 misc.sprint(error) 285 if mode == 'r': 286 misc.gunzip(path) 287 else: 288 self.to_zip = True 289 self.file = open(path[:-3], mode, *args, **opt) 290 path = path[:-3] 291 292 self.parsing = True # check if/when we need to parse the event. 293 self.eventgroup = False 294 295 self.header = '' 296 if mode == 'r': 297 line = '' 298 while 'HepMC::IO_GenEvent-START_EVENT_LISTING' not in line: 299 line = self.file.readline() 300 if not line: 301 self.seek(0) 302 self.banner = '' 303 break 304 if 'b' in mode or self.zip_mode: 305 line = str(line.decode()) 306 self.header += line 307 self.start_event = ''
308
309 - def seek(self, *args, **opts):
310 self.start_event = "" 311 return self.file.seek(*args, **opts)
312
313 - def tell(self):
314 if self.zipmode: 315 currpos = self.file.tell() 316 if not currpos: 317 currpos = self.size 318 return currpos 319 else: 320 self.file.tell()
321
322 - def __iter__(self):
323 return self
324
325 - def __del__(self):
326 try: 327 self.file.close() 328 except Exception: 329 pass
330
331 - def __len__(self):
332 if self.file.closed: 333 return 0 334 if hasattr(self,"len"): 335 return self.len 336 self.seek(0) 337 nb_event=0 338 with misc.TMP_variable(self, 'parsing', False): 339 for _ in self: 340 nb_event +=1 341 self.len = nb_event 342 self.seek(0) 343 return self.len
344
345 - def close(self,*args, **opts):
346 347 out = self.file.close(*args, **opts) 348 if self.to_zip: 349 misc.gzip(self.path)
350 351
352 - def next(self):
353 354 355 return self.next_event()
356 357 __next__ = next 358 359 360 361 362
363 - def next_event(self):
364 """get next event""" 365 text = self.start_event 366 line = '' 367 while 1: 368 line = self.file.readline() 369 if not line: 370 raise StopIteration 371 if 'b' in self.mode or self.zip_mode: 372 line = str(line.decode()) 373 if line.startswith('E'): 374 self.start_event = line 375 if text: 376 return HEPMC_Event(text) 377 else: 378 text += line 379 380 elif line.lstrip().startswith('HepMC::IO_GenEvent-END_EVENT_LISTING'): 381 if text: 382 return HEPMC_Event(text) 383 elif line.lstrip().startswith('HepMC::IO_GenEvent-START_EVENT_LISTING'): 384 text = '' 385 else: 386 text += line
387
388 - def getfilesize(self):
389 if self.zip_mode: 390 self.file.seek(-4, 2) 391 r = self.file.read() 392 self.file.seek() 393 import struct 394 return struct.unpack('<I', r)[0] 395 else: 396 self.file.seek(0,2) 397 pos = self.file.tell() 398 self.file.seek() 399 return pos
400
401 - def write(self, text):
402 403 if self.zip_mode or 'b' in self.mode: 404 self.file.write(text.encode()) 405 else: 406 self.file.write(text)
407 408 @property
409 - def name(self):
410 return self.file.name
411 412 @property
413 - def closed(self):
414 return self.file.closed
415 416 417 418 if "__main__" == __name__: 419 path = "/Users/omattelaer/Documents/eclipse/2.7.1/PROC_sm_43/Events/run_01/tag_1_pythia8_events.hepmc.gz" 420 evts = HEPMC_EventFile(path) 421 nb_event = 0 422 nb_p = 0 423 for event in evts: 424 nb_event +=1 425 for p in event: 426 nb_p+=1 427 print(nb_event, nb_p) 428