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