def simulate(self, iterable, keep_photons_beg=False,
keep_photons_end=False, run_daq=True, max_steps=100 ):
try:
if isinstance(iterable, event.Photons):
raise TypeError # Kludge because Photons looks iterable
else:
first_element, iterable = itertoolset.peek(iterable)
except TypeError:
first_element, iterable = iterable, [iterable]
t_photon_start = time.time()
if isinstance(first_element, event.Event):
iterable = self.photon_generator.generate_events(iterable)
elif isinstance(first_element, event.Photons):
iterable = (event.Event(photons_beg=x) for x in iterable)
elif isinstance(first_element, GPUPhotons):
print "GPU Photons"
iterable = (event.Event(photons_beg=x) for x in iterable) # hacky!!!
elif isinstance(first_element, event.Vertex):
iterable = (event.Event(primary_vertex=vertex, vertices=[vertex]) for vertex in iterable)
iterable = self.photon_generator.generate_events(iterable)
t_photon_end = time.time()
print "Photon Load Time: ",t_photon_end-t_photon_start," sec"
for ev in iterable:
photons = ev.photons_beg
if isinstance(photons,event.Photons):
gpu_photons = GPUPhotons(photons,cl_context=self.context)
elif isinstance(photons,GPUPhotons):
gpu_photons = photons
ev.photons_beg = photons.get()
gpu_photons.propagate(self.gpu_geometry, self.rng_states,
nthreads_per_block=self.nthreads_per_block,
max_blocks=self.max_blocks,
max_steps=max_steps, cl_context=self.context)
ev.nphotons = len(ev.photons_beg.pos)
if not keep_photons_beg:
ev.photons_beg = None
if keep_photons_end:
ev.photons_end = gpu_photons.get()
# Skip running DAQ if we don't have one
if hasattr(self, 'gpu_daq') and run_daq:
t_daq_start = time.time()
self.gpu_daq.begin_acquire( cl_context=self.context )
self.gpu_daq.acquire(gpu_photons, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks, cl_context=self.context )
gpu_channels = self.gpu_daq.end_acquire( cl_context=self.context )
ev.channels = gpu_channels.get()
t_daq_end = time.time()
print "DAQ readout time: ",t_daq_end-t_daq_start," sec"
yield ev
def simulate(self,
iterable,
keep_photons_beg=False,
keep_photons_end=False,
keep_hits=True,
keep_flat_hits=True,
run_daq=False,
max_steps=1000,
photons_per_batch=1000000):
if isinstance(iterable, event.Photons):
first_element, iterable = iterable, [iterable]
else:
first_element, iterable = itertoolset.peek(iterable)
if isinstance(first_element, event.Event):
iterable = self.photon_generator.generate_events(iterable)
elif isinstance(first_element, event.Photons):
iterable = (event.Event(photons_beg=x) for x in iterable)
elif isinstance(first_element, event.Vertex):
iterable = (event.Event(vertices=[vertex]) for vertex in iterable)
iterable = self.photon_generator.generate_events(iterable)
nphotons = 0
batch_events = []
for ev in iterable:
ev.nphotons = len(ev.photons_beg)
ev.photons_beg.evidx[:] = len(batch_events)
nphotons += ev.nphotons
batch_events.append(ev)
#FIXME need an alternate implementation to split an event that is too large
if nphotons >= photons_per_batch:
yield from self._simulate_batch(
batch_events,
keep_photons_beg=keep_photons_beg,
keep_photons_end=keep_photons_end,
keep_hits=keep_hits,
keep_flat_hits=keep_flat_hits,
run_daq=run_daq,
max_steps=max_steps)
nphotons = 0
batch_events = []
if len(batch_events) != 0:
yield from self._simulate_batch(batch_events,
keep_photons_beg=keep_photons_beg,
keep_photons_end=keep_photons_end,
keep_hits=keep_hits,
keep_flat_hits=keep_flat_hits,
run_daq=run_daq,
max_steps=max_steps)
def simulate(self,
iterable,
keep_photons_beg=False,
keep_photons_end=False,
keep_hits=True,
run_daq=False,
max_steps=100):
if isinstance(iterable, event.Photons):
first_element, iterable = iterable, [iterable]
else:
first_element, iterable = itertoolset.peek(iterable)
if isinstance(first_element, event.Event):
iterable = self.photon_generator.generate_events(iterable)
elif isinstance(first_element, event.Photons):
iterable = (event.Event(photons_beg=x) for x in iterable)
elif isinstance(first_element, event.Vertex):
iterable = (event.Event(vertices=[vertex]) for vertex in iterable)
iterable = self.photon_generator.generate_events(iterable)
for ev in iterable:
gpu_photons = gpu.GPUPhotons(ev.photons_beg)
gpu_photons.propagate(self.gpu_geometry,
self.rng_states,
nthreads_per_block=self.nthreads_per_block,
max_blocks=self.max_blocks,
max_steps=max_steps)
ev.nphotons = len(ev.photons_beg.pos)
if not keep_photons_beg:
ev.photons_beg = None
if keep_photons_end:
ev.photons_end = gpu_photons.get()
if hasattr(self.detector, 'num_channels') and keep_hits:
ev.hits = gpu_photons.get_hits(self.gpu_geometry)
# Skip running DAQ if we don't have one
# Disabled by default because incredibly special-case
if hasattr(self, 'gpu_daq') and run_daq:
self.gpu_daq.begin_acquire()
self.gpu_daq.acquire(
gpu_photons,
self.rng_states,
nthreads_per_block=self.nthreads_per_block,
max_blocks=self.max_blocks)
gpu_channels = self.gpu_daq.end_acquire()
ev.channels = gpu_channels.get()
yield ev
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)
# 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
return pyev
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 izip(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 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 zip(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 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 eval_pdf(self,
event_channels,
iterable,
min_twidth,
trange,
min_qwidth,
qrange,
min_bin_content=100,
nreps=1,
ndaq=1,
nscatter=1,
time_only=True):
"""Returns tuple: 1D array of channel hit counts, 1D array of PDF
probability densities."""
ndaq_per_rep = 64
ndaq_reps = ndaq // ndaq_per_rep
gpu_daq = gpu.GPUDaq(self.gpu_geometry, ndaq=ndaq_per_rep)
self.gpu_pdf.setup_pdf_eval(event_channels.hit,
event_channels.t,
event_channels.q,
min_twidth,
trange,
min_qwidth,
qrange,
min_bin_content=min_bin_content,
time_only=True)
first_element, iterable = itertoolset.peek(iterable)
if isinstance(first_element, event.Event):
iterable = self.photon_generator.generate_events(iterable)
elif isinstance(first_element, event.Photons):
iterable = (event.Event(photons_beg=x) for x in iterable)
for ev in iterable:
gpu_photons_no_scatter = gpu.GPUPhotons(ev.photons_beg,
ncopies=nreps)
gpu_photons_scatter = gpu.GPUPhotons(ev.photons_beg,
ncopies=nreps * nscatter)
gpu_photons_no_scatter.propagate(
self.gpu_geometry,
self.rng_states,
nthreads_per_block=self.nthreads_per_block,
max_blocks=self.max_blocks,
use_weights=True,
scatter_first=-1,
max_steps=10)
gpu_photons_scatter.propagate(
self.gpu_geometry,
self.rng_states,
nthreads_per_block=self.nthreads_per_block,
max_blocks=self.max_blocks,
use_weights=True,
scatter_first=1,
max_steps=5)
nphotons = gpu_photons_no_scatter.true_nphotons # same for scatter
for i in range(gpu_photons_no_scatter.ncopies):
start_photon = i * nphotons
gpu_photon_no_scatter_slice = gpu_photons_no_scatter.select(
event.SURFACE_DETECT,
start_photon=start_photon,
nphotons=nphotons)
gpu_photon_scatter_slices = [
gpu_photons_scatter.select(
event.SURFACE_DETECT,
start_photon=(nscatter * i + j) * nphotons,
nphotons=nphotons) for j in range(nscatter)
]
if len(gpu_photon_no_scatter_slice) == 0:
continue
#weights = gpu_photon_slice.weights.get()
#print 'weights', weights.min(), weights.max()
for j in range(ndaq_reps):
gpu_daq.begin_acquire()
gpu_daq.acquire(gpu_photon_no_scatter_slice,
self.rng_states,
nthreads_per_block=self.nthreads_per_block,
max_blocks=self.max_blocks)
for scatter_slice in gpu_photon_scatter_slices:
gpu_daq.acquire(
scatter_slice,
self.rng_states,
nthreads_per_block=self.nthreads_per_block,
max_blocks=self.max_blocks,
weight=1.0 / nscatter)
gpu_channels = gpu_daq.end_acquire()
self.gpu_pdf.accumulate_pdf_eval(
gpu_channels, nthreads_per_block=ndaq_per_rep)
return self.gpu_pdf.get_pdf_eval()
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)
