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),
pos=tvector3_to_ndarray(vertex.pos),
dir=tvector3_to_ndarray(vertex.dir),
ke=vertex.ke,
t0=vertex.t0,
pol=tvector3_to_ndarray(vertex.pol),
trackid=vertex.trackid,
pdgcode=vertex.pdgcode)
def particle_gun(particle_name_iter,
pos_iter,
dir_iter,
ke_iter,
t0_iter=constant(0.0),
start_id=0):
for i, particle_name, pos, dir, ke, t0 in izip(count(start_id),
particle_name_iter,
pos_iter, dir_iter, ke_iter,
t0_iter):
dir = dir / norm(dir)
vertex = event.Vertex(particle_name, pos, dir, ke, t0=t0)
ev_vertex = event.Event(i, vertex, [vertex])
yield ev_vertex
def pi0_gun(pos_iter, dir_iter, ke_iter, t0_iter=constant(0.0), start_id=0, gamma1_dir_iter=None):
if gamma1_dir_iter is None:
gamma1_dir_iter = isotropic()
for i, pos, dir, ke, t0, gamma1_dir in zip(count(start_id), pos_iter, dir_iter, ke_iter, t0_iter, gamma1_dir_iter):
dir = dir/norm(dir)
primary_vertex = event.Vertex('pi0', pos, dir, ke, t0=t0)
# In rest frame
theta_rest = np.arccos(gamma1_dir[2])
phi_rest = np.arctan2(gamma1_dir[1], gamma1_dir[0])
# In lab frame
(gamma1_e, gamma1_dir), (gamma2_e, gamma2_dir) = \
pi0_decay(ke+134.9766, dir, theta_rest, phi_rest)
gamma1_vertex = event.Vertex('gamma', pos, gamma1_dir, gamma1_e, t0=t0)
gamma2_vertex = event.Vertex('gamma', pos, gamma2_dir, gamma2_e, t0=t0)
ev_vertex = event.Event(i, primary_vertex, [gamma1_vertex, gamma2_vertex])
yield ev_vertex
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)