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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
|
import ROOT
import os.path
import numpy as np
ROOT.gROOT.ProcessLine('.L '+os.path.join(os.path.dirname(__file__), 'root.C+g'))
import ROOT
import chroma.event as event
def tvector3_to_ndarray(vec):
'''Convert a ROOT.TVector3 into a numpy np.float32 array'''
return np.array((vec.X(), vec.Y(), vec.Z()), dtype=np.float32)
def make_photon_with_arrays(size):
'''Returns a new chroma.event.Photons object for `size` number of
photons with empty arrays set for all the photon attributes.'''
return event.Photons(pos=np.empty((size,3), dtype=np.float32),
dir=np.empty((size,3), dtype=np.float32),
pol=np.empty((size,3), dtype=np.float32),
wavelengths=np.empty(size, dtype=np.float32),
t=np.empty(size, dtype=np.float32),
flags=np.empty(size, dtype=np.uint32),
last_hit_triangles=np.empty(size, dtype=np.int32))
def root_vertex_to_python_vertex(vertex):
"Returns a chroma.event.Vertex object from a root Vertex object."
return event.Vertex(str(vertex.particle_name),
tvector3_to_ndarray(vertex.pos),
tvector3_to_ndarray(vertex.dir),
tvector3_to_ndarray(vertex.pol),
vertex.ke,
vertex.t0)
def root_event_to_python_event(ev):
'''Returns a new chroma.event.Event object created from the
contents of the ROOT event `ev`.'''
pyev = event.Event(ev.id)
pyev.primary_vertex = root_vertex_to_python_vertex(ev.primary_vertex)
for vertex in ev.vertices:
pyev.vertices.append(root_vertex_to_python_vertex(vertex))
# photon begin
if ev.photons_beg.size() > 0:
photons = make_photon_with_arrays(ev.photons_beg.size())
ROOT.get_photons(ev.photons_beg,
photons.pos.ravel(),
photons.dir.ravel(),
photons.pol.ravel(),
photons.wavelengths,
photons.t,
photons.last_hit_triangles,
photons.flags)
pyev.photons_beg = photons
# photon end
if ev.photons_end.size() > 0:
photons = make_photon_with_arrays(ev.photons_end.size())
ROOT.get_photons(ev.photons_end,
photons.pos.ravel(),
photons.dir.ravel(),
photons.pol.ravel(),
photons.wavelengths,
photons.t,
photons.last_hit_triangles,
photons.flags)
pyev.photons_end = photons
# channels
hit = np.empty(ev.nchannels, dtype=np.int32)
t = np.empty(ev.nchannels, dtype=np.float32)
q = np.empty(ev.nchannels, dtype=np.float32)
flags = np.empty(ev.nchannels, dtype=np.uint32)
ROOT.get_channels(ev, hit, t, q, flags)
pyev.channels = event.Channels(hit.astype(bool), t, q, flags)
return pyev
class RootReader(object):
'''Reader of Chroma events from a ROOT file. This class can be used to
navigate up and down the file linearly or in a random access fashion.
All returned events are instances of the chroma.event.Event class.
It implements the iterator protocol, so you can do
for ev in RootReader('electron.root'):
# process event here
'''
def __init__(self, filename):
'''Open ROOT file named `filename` containing TTree `T`.'''
self.f = ROOT.TFile(filename)
self.T = self.f.T
self.i = -1
def __len__(self):
'''Returns number of events in this file.'''
return self.T.GetEntries()
def next(self):
'''Return the next event in the file. Raises StopIteration
when you get to the end.'''
if self.i + 1 >= len(self):
raise StopIteration
self.i += 1
self.T.GetEntry(self.i)
return root_event_to_python_event(self.T.ev)
def prev(self):
'''Return the next event in the file. Raises StopIteration if
that would go past the beginning.'''
if self.i <= 0:
self.i = -1
raise StopIteration
self.i -= 1
self.T.GetEntry(self.i)
return root_event_to_python_event(self.T.ev)
def current(self):
'''Return the current event in the file.'''
self.T.GetEntry(self.i) # just in case?
return root_event_to_python_event(self.T.ev)
def jump_to(self, index):
'''Return the event at `index`. Updates current location.'''
if index < 0 or index >= len(self):
raise IndexError
self.T.GetEntry(self.i)
return root_event_to_python_event(self.T.ev)
def index(self):
'''Return the current event index'''
return self.i
class RootWriter(object):
def __init__(self, filename):
self.filename = filename
self.file = ROOT.TFile(filename, 'RECREATE')
self.T = ROOT.TTree('T', 'Chroma events')
self.ev = ROOT.Event()
self.T.Branch('ev', self.ev)
def write_event(self, pyev):
"Write an event.Event object to the ROOT tree as a ROOT.Event object."
self.ev.id = pyev.id
self.ev.primary_vertex.particle_name = pyev.primary_vertex.particle_name
self.ev.primary_vertex.pos.SetXYZ(*pyev.primary_vertex.pos)
self.ev.primary_vertex.dir.SetXYZ(*pyev.primary_vertex.dir)
if pyev.primary_vertex.pol is not None:
self.ev.primary_vertex.pol.SetXYZ(*pyev.primary_vertex.pol)
self.ev.primary_vertex.ke = pyev.primary_vertex.ke
if pyev.photons_beg is not None:
photons = pyev.photons_beg
ROOT.fill_photons(self.ev.photons_beg,
len(photons.pos),
photons.pos.ravel(),
photons.dir.ravel(),
photons.pol.ravel(),
photons.wavelengths, photons.t,
photons.last_hit_triangles, photons.flags)
if pyev.photons_end is not None:
photons = pyev.photons_end
ROOT.fill_photons(self.ev.photons_end,
len(photons.pos),
photons.pos.ravel(),
photons.dir.ravel(),
photons.pol.ravel(),
photons.wavelengths, photons.t,
photons.last_hit_triangles, photons.flags)
self.ev.vertices.resize(0)
if pyev.vertices is not None:
self.ev.vertices.resize(len(pyev.vertices))
for i, vertex in enumerate(pyev.vertices):
self.ev.vertices[i].particle_name = vertex.particle_name
self.ev.vertices[i].pos.SetXYZ(*vertex.pos)
self.ev.vertices[i].dir.SetXYZ(*vertex.dir)
if vertex.pol is not None:
self.ev.vertices[i].pol.SetXYZ(*vertex.pol)
self.ev.vertices[i].ke = vertex.ke
self.ev.vertices[i].t0 = vertex.t0
if pyev.channels is not None:
ROOT.fill_channels(self.ev, np.count_nonzero(pyev.channels.hit), np.arange(len(pyev.channels.t))[pyev.channels.hit].astype(np.int32), pyev.channels.t, pyev.channels.q, pyev.channels.flags, len(pyev.channels.hit))
self.T.Fill()
def close(self):
self.T.Write()
self.file.Close()
|
