1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
|
import numpy as np
class Vertex(object):
def __init__(self, particle_name, pos, dir, pol, ke, t0=0.0):
self.particle_name = particle_name
self.pos = pos
self.dir = dir
self.pol = pol
self.ke = ke
self.t0 = t0
class Photons(object):
def __init__(self, pos, dir, pol, wavelengths, t=None, last_hit_triangles=None, flags=None):
self.pos = pos
self.dir = dir
self.pol = pol
self.wavelengths = wavelengths
if t is None:
self.t = np.zeros(len(pos), dtype=np.float32)
else:
self.t = t
if last_hit_triangles is None:
self.last_hit_triangles = np.empty(len(pos), dtype=np.int32)
self.last_hit_triangles.fill(-1)
else:
self.last_hit_triangles = last_hit_triangles
if flags is None:
self.flags = np.zeros(len(pos), dtype=np.uint32)
else:
self.flags = flags
def __add__(self, other):
pos = np.concatenate((self.pos, other.pos))
dir = np.concatenate((self.dir, other.dir))
pol = np.concatenate((self.pol, other.pol))
wavelengths = np.concatenate((self.wavelengths, other.wavelengths))
t = np.concatenate((self.t, other.t))
last_hit_triangles = np.concatenate((self.last_hit_triangles, other.last_hit_triangles))
flags = np.concatenate((self.flags, other.flags))
return Photons(pos, dir, pol, wavelengths, t, last_hit_triangles, flags)
class Channels(object):
def __init__(self, hit, t, q, flags=None):
self.hit = hit
self.t = t
self.q = q
self.flags = flags
def hit_channels(self):
return self.hit.nonzero(), self.t[self.hit], self.q[self.hit]
class Event(object):
def __init__(self, id=0, primary_vertex=None, vertices=None, photons_beg=None, photons_end=None, channels=None):
self.id = id
self.nphotons = None
self.primary_vertex = primary_vertex
if vertices is not None:
if np.iterable(vertices):
self.vertices = vertices
else:
self.vertices = [vertices]
else:
self.vertices = []
self.photons_beg = photons_beg
self.photons_end = photons_end
self.channels = channels
|
