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
48 self.nb_flow_list = 0
49 self.flows = []
50
51 if text:
52 self.parse(text, event)
53
54 @property
57
58 pid = pdg_code
59
60 @property
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])
70 self.pdg = int(data[2])
71 self.px = float(data[3])
72 self.py = float(data[4])
73 self.pz = float(data[5])
74 self.E = float(data[6])
75 self.mass = float(data[7])
76 self.status = int(data[8])
77 self.polarization_theta = float(data[9])
78 self.polarization_phi = float(data[10])
79 self.vertex_barcode = float(data[11])
80 self.nb_flow_list = int(data[12])
81 self.flows = [(int(data[13+2*i]),int(data[13+2*i+1]))
82 for i in range(self.nb_flow_list)]
83
84 if event:
85 event.curr_vertex.add_outcoming(self)
86
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
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
140 self.incoming.append(particle)
141
143 self.outcoming.append(particle)
144
146
148 """The initialization of an empty Event (or one associate to a text file)"""
149
150 self.particles = {}
151 self.vertex = {}
152
153
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
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
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
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
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
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
238 """just keep the information so far"""
239 self.N = '%s\n' % line
241 self.U = '%s\n' % line
243 self.H = '%s\n' % line
245 self.F = '%s\n' % line
247 self.C = '%s\n' % line
248
250 return list(self.particles.values()).__iter__()
251
252
253
254
255
257 self.vertex[V.barcode] = V
258
259 - def add_particle(self, P):
260 self.particles[P.barcode] = P
261
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
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
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
324
326 try:
327 self.file.close()
328 except Exception:
329 pass
330
344
345 - def close(self,*args, **opts):
350
351
356
357 __next__ = next
358
359
360
361
362
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
400
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
411
412 @property
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