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)