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
27 self.nb_flow_list = 0
28 self.flows = []
29
30 if text:
31 self.parse(text, event)
32
33 @property
36
37 pid = pdg_code
38
39 @property
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])
49 self.pdg = int(data[2])
50 self.px = float(data[3])
51 self.py = float(data[4])
52 self.pz = float(data[5])
53 self.E = float(data[6])
54 self.mass = float(data[7])
55 self.status = int(data[8])
56 self.polarization_theta = float(data[9])
57 self.polarization_phi = float(data[10])
58 self.vertex_barcode = float(data[11])
59 self.nb_flow_list = int(data[12])
60 self.flows = [(int(data[13+2*i]),int(data[13+2*i+1]))
61 for i in range(self.nb_flow_list)]
62
63 if event:
64 event.curr_vertex.add_outcoming(self)
65
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
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
119 self.incoming.append(particle)
120
122 self.outcoming.append(particle)
123
125
127 """The initialization of an empty Event (or one associate to a text file)"""
128
129 self.particles = {}
130 self.vertex = {}
131
132
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
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
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
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
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
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
217 """just keep the information so far"""
218 self.N = '%s\n' % line
220 self.U = '%s\n' % line
222 self.H = '%s\n' % line
224 self.F = '%s\n' % line
226 self.C = '%s\n' % line
227
229 return self.particles.values().__iter__()
230
231
232
233
234
236 self.vertex[V.barcode] = V
237
238 - def add_particle(self, P):
239 self.particles[P.barcode] = P
240
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
268
269 self.parsing = True
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):
297
298
299
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
323 """A way to read/write a gzipped lhef event"""
324
326 currpos = super(HEPMC_EventFileGzip, self).tell()
327 if not currpos:
328 currpos = self.size
329 return currpos
330
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
340 """A way to read a standard event file"""
341
342 - def close(self,*args, **opts):
347
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