def get(self):
t = self.t.get()
q = self.q.get()
# For now, assume all channels with small
# enough hit time were hit.
return event.Channels(t < 1e8, t, q, self.flags.get())
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
if ev.nchannels > 0:
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)
else:
pyev.channels = None
return pyev
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)
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,
photons.channel)
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,
photons.channel)
pyev.photons_end = photons
# photon tracks
if ev.photon_tracks.size() > 0:
photon_tracks = []
for i in range(ev.photon_tracks.size()):
photons = make_photon_with_arrays(ev.photon_tracks[i].size())
ROOT.get_photons(ev.photon_tracks[i],
photons.pos.ravel(),
photons.dir.ravel(),
photons.pol.ravel(),
photons.wavelengths,
photons.t,
photons.last_hit_triangles,
photons.flags,
photons.channel)
photon_tracks.append(photons)
pyev.photon_tracks = photon_tracks
pyev.photon_parent_trackids = np.asarray(ev.photon_parent_trackids).copy()
# hits
if ev.hits.size() > 0:
pyev.hits = {}
for hit in ev.hits:
photons = make_photon_with_arrays(hit.second.size())
ROOT.get_photons(hit.second,
photons.pos.ravel(),
photons.dir.ravel(),
photons.pol.ravel(),
photons.wavelengths, photons.t,
photons.last_hit_triangles, photons.flags,
photons.channel)
pyev.hits[hit.first] = photons
# flat_hits
if ev.flat_hits.size() > 0:
photons = make_photon_with_arrays(ev.flat_hits.size())
ROOT.get_photons(ev.flat_hits,
photons.pos.ravel(),
photons.dir.ravel(),
photons.pol.ravel(),
photons.wavelengths, photons.t,
photons.last_hit_triangles, photons.flags,
photons.channel)
pyev.flat_hits = 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,
photons.channel)
pyev.photons_end = photons
# channels
if ev.nchannels > 0:
hit = np.zeros(ev.nchannels, dtype=np.int32)
t = np.zeros(ev.nchannels, dtype=np.float32)
q = np.zeros(ev.nchannels, dtype=np.float32)
flags = np.zeros(ev.nchannels, dtype=np.uint32)
ROOT.get_channels(ev, hit, t, q, flags)
pyev.channels = event.Channels(hit.astype(bool), t, q, flags)
else:
pyev.channels = None
return pyev
def test_file_write_and_read(self):
ev = event.Event(
1,
event.Vertex('e-',
pos=(0, 0, 1),
dir=(1, 0, 0),
ke=15.0,
pol=(0, 1, 0),
t0=40.0))
photons_beg = root.make_photon_with_arrays(1)
photons_beg.pos[0] = (1, 2, 3)
photons_beg.dir[0] = (4, 5, 6)
photons_beg.pol[0] = (7, 8, 9)
photons_beg.wavelengths[0] = 400.0
photons_beg.t[0] = 100.0
photons_beg.last_hit_triangles[0] = 5
photons_beg.flags[0] = 20
ev.photons_beg = photons_beg
photons_end = root.make_photon_with_arrays(1)
photons_end.pos[0] = (1, 2, 3)
photons_end.dir[0] = (4, 5, 6)
photons_end.pol[0] = (7, 8, 9)
photons_end.wavelengths[0] = 400.0
photons_end.t[0] = 100.0
photons_end.last_hit_triangles[0] = 5
photons_end.flags[0] = 20
ev.photons_end = photons_end
ev.vertices = [ev.primary_vertex]
channels = event.Channels(hit=np.array([True, False]),
t=np.array([20.0, 1e9], dtype=np.float32),
q=np.array([2.0, 0.0], dtype=np.float32),
flags=np.array([8, 32], dtype=np.uint32))
ev.channels = channels
filename = '/tmp/chroma-filewritertest.root'
writer = root.RootWriter(filename)
writer.write_event(ev)
writer.close()
# Exercise the RootReader methods
reader = root.RootReader(filename)
self.assertEquals(len(reader), 1)
self.assertRaises(StopIteration, reader.prev)
reader.next()
self.assertEqual(reader.index(), 0)
self.assertRaises(StopIteration, reader.next)
reader.jump_to(0)
# Enough screwing around, let's get the one event in the file
newev = reader.current()
# Now check if everything is correct in the event
for attribute in ['id']:
self.assertEqual(getattr(ev, attribute), getattr(newev, attribute),
'compare %s' % attribute)
for attribute in ['pos', 'dir', 'pol', 'ke', 't0']:
self.assertTrue(
np.allclose(getattr(ev.primary_vertex, attribute),
getattr(newev.primary_vertex, attribute)),
'compare %s' % attribute)
for i in range(len(ev.vertices)):
self.assertTrue(
np.allclose(getattr(ev.vertices[i], attribute),
getattr(newev.vertices[i], attribute)),
'compare %s' % attribute)
for attribute in [
'pos', 'dir', 'pol', 'wavelengths', 't', 'last_hit_triangles',
'flags'
]:
self.assertTrue(
np.allclose(getattr(ev.photons_beg, attribute),
getattr(newev.photons_beg, attribute)),
'compare %s' % attribute)
self.assertTrue(
np.allclose(getattr(ev.photons_end, attribute),
getattr(newev.photons_end, attribute)),
'compare %s' % attribute)