diff options
Diffstat (limited to 'sim.py')
| -rwxr-xr-x | sim.py | 283 |
1 files changed, 141 insertions, 142 deletions
@@ -3,148 +3,147 @@ import sys import time import os import numpy as np - -import detectors -import optics -import generator -from generator import constant import itertools import threading -import gpu -from fileio import root -from chroma.itertoolset import repeating_iterator -from tools import profile_if_possible, enable_debug_on_crash + +from chroma import detectors +from chroma import optics +from chroma import generator +from chroma import gpu +from chroma.fileio import root +from chroma import tools +from chroma import event +from chroma.itertoolset import peek, repeat_copy, repeating_iterator def pick_seed(): - '''Returns a seed for a random number generator selected using - a mixture of the current time and the current process ID.''' + """Returns a seed for a random number generator selected using + a mixture of the current time and the current process ID.""" return int(time.time()) ^ (os.getpid() << 16) class Simulation(object): - def __init__(self, detector, detector_material, - seed=None, cuda_device=None, geant4_processes=4, bvh_bits=11, - use_cache=True): + def __init__(self, detector, seed=None, cuda_device=None, geant4_processes=4, bvh_bits=11, use_cache=True, nthreads_per_block=64, max_blocks=1024): self.detector = detector + self.nthreads_per_block = nthreads_per_block + self.max_blocks = max_blocks + if seed is None: self.seed = pick_seed() else: self.seed = seed - print >>sys.stderr, 'RNG seed:', self.seed + + print 'RNG seed: %i' % self.seed # We have three generators to seed: numpy.random, GEANT4, and CURAND. - # The latter two are done below + # The latter two are done below. np.random.seed(self.seed) - self.detector_material = detector_material if geant4_processes > 0: - self.photon_generator = generator.photon.G4ParallelGenerator(geant4_processes, - detector_material, - base_seed=self.seed) + self.photon_generator = generator.photon.G4ParallelGenerator(geant4_processes, detector.detector_material, base_seed=self.seed) else: self.photon_generator = None - print >>sys.stderr, 'Creating BVH with %d bits...' % (bvh_bits) - detector.build(bits=bvh_bits, use_cache=use_cache) + if not hasattr(detector, 'mesh'): + # need to build geometry + print 'Creating BVH with %i bits...' % bvh_bits + detector.build(bits=bvh_bits, use_cache=use_cache) - print >>sys.stderr, 'Initializing GPU...' - self.gpu_worker = gpu.GPU(cuda_device) + self.gpu = gpu.GPU(cuda_device) + + # geometry is loaded onto gpu by default + self.gpu_geometry = gpu.GPUGeometry(self.gpu, detector) - print >>sys.stderr, 'Loading detector onto GPU...' - self.gpu_worker.load_geometry(detector) + self.gpu_daq = gpu.GPUDaq(self.gpu_geometry, max(self.detector.pmtids)) + self.gpu_pdf = gpu.GPUPDF() + + print 'Initializing random numbers generators...' + self.rng_states = gpu.get_rng_states(self.nthreads_per_block*self.max_blocks, seed=self.seed) - print >>sys.stderr, 'Initializing random numbers generators...' - self.gpu_worker.setup_propagate(seed=self.seed) - self.gpu_worker.setup_daq(max(self.detector.pmtids)) self.pdf_config = None - def simulate(self, nevents, vertex_generator, keep_photon_start=False, keep_photon_stop=False, - run_daq=True, nreps=1): - photon_gen = repeating_iterator(self.photon_generator.generate_events(nevents, vertex_generator), - nreps) - return self.simulate_photons(nevents*nreps, - photon_gen, - keep_photon_start=keep_photon_start, keep_photon_stop=keep_photon_stop, - run_daq=run_daq) - - def simulate_photons(self, nevents, photon_generator, keep_photon_start=False, keep_photon_stop=False, - run_daq=True): - for ev in itertools.islice(photon_generator, nevents): - self.gpu_worker.load_photons(ev.photon_start) - self.gpu_worker.propagate() - self.gpu_worker.run_daq() - - ev.nphoton = len(ev.photon_start.positions) - - if not keep_photon_start: - ev.photon_start = None - - if keep_photon_stop: - ev.photon_stop = self.gpu_worker.get_photons() + def simulate(self, iterable, keep_photons_beg=False, keep_photons_end=False, run_daq=True, max_steps=10): + first_element, iterable = 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 = gpu.GPUPhotons(ev.photons_beg) + + gpu.propagate(self.gpu, gpu_photons, 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 run_daq: - ev.channels = self.gpu_worker.get_hits() + gpu_channels = self.gpu_daq.acquire(gpu_photons, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks) + ev.channels = gpu_channels.get() yield ev - def propagate_photons(self, photons, max_steps=10): - self.gpu_worker.load_photons(photons) - self.gpu_worker.propagate(max_steps=max_steps) - return self.gpu_worker.get_photons() - - def create_pdf(self, nevents, vertex_generator, tbins, trange, - qbins, qrange, nreps=1): - photon_gen = repeating_iterator(self.photon_generator.generate_events(nevents, vertex_generator), - nreps) - return self.create_pdf_from_photons(nevents*nreps, photon_gen, - tbins, trange, qbins, qrange) - - def create_pdf_from_photons(self, nevents, photon_generator, - tbins, trange, qbins, qrange): - '''Returns tuple: 1D array of channel hit counts, 3D array of (channel, time, charge) pdfs''' + def create_pdf(self, iterable, tbins, trange, qbins, qrange, nreps=1): + """Returns tuple: 1D array of channel hit counts, 3D array of + (channel, time, charge) pdfs.""" + first_element, iterable = peek(iterable) + + if isinstance(first_element, event.Event): + iterable = self.photon_generator.generate_events(iterable) + pdf_config = (tbins, trange, qbins, qrange) if pdf_config != self.pdf_config: self.pdf_config = pdf_config - self.gpu_worker.setup_pdf(max(self.detector.pmtids), tbins, trange, + self.gpu_pdf.setup_pdf(max(self.detector.pmtids), tbins, trange, qbins, qrange) else: - self.gpu_worker.clear_pdf() + self.gpu_pdf.clear_pdf() - for ev in itertools.islice(photon_generator, nevents): - self.gpu_worker.load_photons(ev.photon_start) - self.gpu_worker.propagate() - self.gpu_worker.run_daq() - self.gpu_worker.add_hits_to_pdf() - - return self.gpu_worker.get_pdfs() - - - def eval_pdf(self, event_channels, nevents, vertex_generator, - min_twidth, trange, min_qwidth, qrange, - min_bin_content=20, nreps=1, time_only=True): - photon_gen = repeating_iterator(self.photon_generator.generate_events(nevents, vertex_generator), - nreps) - return self.eval_pdf_from_photons(event_channels, nevents*nreps, photon_gen, - min_twidth, trange, min_qwidth, qrange, - min_bin_content=min_bin_content, time_only=time_only) - - def eval_pdf_from_photons(self, event_channels, nevents, photon_generator, - min_twidth, trange, min_qwidth, qrange, - min_bin_content=20, time_only=True): - '''Returns tuple: 1D array of channel hit counts, 1D array of PDF probability densities''' - self.gpu_worker.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) - - for ev in itertools.islice(photon_generator, nevents): - self.gpu_worker.load_photons(ev.photon_start) - self.gpu_worker.propagate() - self.gpu_worker.run_daq() - self.gpu_worker.accumulate_pdf_eval() + if nreps > 1: + iterable = repeating_iterator(iterable) + + for ev in iterable: + gpu_photons = gpu.GPUPhotons(ev.photons_beg) + gpu.propagate(self.gpu, gpu_photons, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks) + gpu_channels = self.gpu_daq.acquire(gpu_photons, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks) + self.gpu_pdf.add_hits_to_pdf(gpu_channels) - return self.gpu_worker.get_pdf_eval() + return self.gpu_pdf.get_pdfs() + + def eval_pdf(self, event_channels, iterable, min_twidth, trange, min_qwidth, qrange, min_bin_content=20, nreps=1, time_only=True): + """Returns tuple: 1D array of channel hit counts, 1D array of PDF + probability densities.""" + 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 = peek(iterable) + + if isinstance(first_element, event.Event): + iterable = self.photon_generator.generate_events(iterable) + + if nreps > 1: + iterable = repeating_iterator(iterable) + + for ev in iterable: + gpu_photons = gpu.GPUPhotons(ev.photons_beg) + gpu.propagate(self.gpu, gpu_photons, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks) + gpu_channels = self.gpu_daq.acquire(gpu_photons, self.rng_states, nthreads_per_block=self.nthreads_per_block, max_blocks=self.max_blocks) + self.gpu_pdf.accumulate_pdf_eval(gpu_channels) + return self.gpu_pdf.get_pdf_eval() -@profile_if_possible +@tools.profile_if_possible def main(): import optparse @@ -155,17 +154,22 @@ def main(): help='Set random number generator seed') parser.add_option('-g', type='int', dest='ngenerators', default=4, help='Number of GEANT4 generator processes') - parser.add_option('--detector', type='string', dest='detector', default='microlbne') + parser.add_option('--detector', type='string', dest='detector', + default='microlbne') parser.add_option('--nevents', type='int', dest='nevents', default=100) - parser.add_option('--particle', type='string', dest='particle', default='e-') - parser.add_option('--energy', type='float', dest='energy', default=100.0) - parser.add_option('--pos', type='string', dest='pos', default='(0,0,0)') - parser.add_option('--dir', type='string', dest='dir', default='(1,0,0)') - parser.add_option('--save-photon-start', action='store_true', - dest='save_photon_start', default=False, + parser.add_option('--particle', type='string', dest='particle', + help='particle name', default='e-') + parser.add_option('--ke', type='float', dest='ke', + help='particle kinetic energy in MeV', default=100.0) + parser.add_option('--pos', type='string', dest='pos', + help='particle vertex origin.', default='(0,0,0)') + parser.add_option('--dir', type='string', dest='dir', + help='particle vertex direction.', default='(1,0,0)') + parser.add_option('--save-photon-beg', action='store_true', + dest='save_photons_beg', default=False, help='Save initial photon vertices to disk') - parser.add_option('--save-photon-stop', action='store_true', - dest='save_photon_stop', default=False, + parser.add_option('--save-photon-end', action='store_true', + dest='save_photons_end', default=False, help='Save final photon vertices to disk') options, args = parser.parse_args() @@ -178,57 +182,52 @@ def main(): if options.nevents <= 0: sys.exit('--nevents must be greater than 0!') - position = np.array(eval(options.pos), dtype=float) - direction = np.array(eval(options.dir), dtype=float) - print >>sys.stderr, 'Loading detector %s...' % options.detector + pos = np.array(eval(options.pos), dtype=float) + dir = np.array(eval(options.dir), dtype=float) + + print 'Loading detector %s...' % options.detector + sys.stdout.flush() detector = detectors.find(options.detector) - print >>sys.stderr, 'Creating generator...' + print 'Creating particle vertex generator...' + sys.stdout.flush() if options.particle == 'pi0': - vertex_generator = generator.vertex.pi0_gun(pi0_position=constant(position), - pi0_direction=constant(direction), - pi0_total_energy=constant(options.energy)) + ev_vertex_iter = itertools.islice(generator.vertex.pi0_gun(itertools.repeat(pos), itertools.repeat(dir), itertools.repeat(options.ke)), options.nevents) else: - vertex_generator = generator.vertex.particle_gun(particle_name=constant(options.particle), - position=constant(position), - direction=constant(direction), - total_energy=constant(options.energy)) + vertex = event.Vertex(options.particle, pos, dir, None, options.ke) + ev_vertex_iter = (event.Event(i, vertex, [vertex]) for i, vertex in zip(range(options.nevents), repeat_copy(vertex))) # Initializing simulation - print >>sys.stderr, 'WARNING: ASSUMING DETECTOR IS WCSIM WATER!!' - simulation = Simulation(detector=detector, detector_material=optics.water_wcsim, - seed=options.seed, cuda_device=options.device, - geant4_processes=options.ngenerators, bvh_bits=options.nbits) + simulation = Simulation(detector=detector, seed=options.seed, cuda_device=options.device, geant4_processes=options.ngenerators, bvh_bits=options.nbits) # Create output file writer = root.RootWriter(output_filename) # Preheat generator - event_iterator = simulation.simulate(options.nevents, vertex_generator, - keep_photon_start=options.save_photon_start, - keep_photon_stop=options.save_photon_stop) + ev_iter = simulation.simulate(ev_vertex_iter, keep_photons_beg=options.save_photons_beg, keep_photons_end=options.save_photons_end) + print 'Starting simulation...' - print >>sys.stderr, 'Starting simulation...' - start_sim = time.time() nphotons = 0 + t0 = time.time() - for i, ev in enumerate(event_iterator): - assert ev.nphoton > 0, 'GEANT4 generated event with no photons!' - nphotons += ev.nphoton + for i, ev in enumerate(ev_iter): + print "\rEvent: %i" % (i+1), + sys.stdout.flush() - writer.write_event(ev) + assert ev.nphotons > 0, 'Geant4 generated event with no photons!' + nphotons += ev.nphotons - if i % 10 == 0: - print >>sys.stderr, "\rEvent:", i, + writer.write_event(ev) + print - end_sim = time.time() - print >>sys.stderr, "\rEvent:", options.nevents - 1 + elapsed = time.time() - t0 writer.close() - print >>sys.stderr, 'Done. %1.1f events/sec, %1.0f photons/sec.' % (options.nevents/(end_sim - start_sim), nphotons/(end_sim - start_sim)) + print '%f elapsed, %1.1f events/sec, %1.0f photons/sec.' % \ + (elapsed, options.nevents/elapsed, nphotons/elapsed) if __name__ == '__main__': - enable_debug_on_crash() + tools.enable_debug_on_crash() main() |