diff options
Diffstat (limited to 'sim.py')
| -rwxr-xr-x | sim.py | 129 |
1 files changed, 85 insertions, 44 deletions
@@ -9,8 +9,11 @@ import detectors import optics import generator from generator import constant +import itertools +import threading import gpu from fileio import root + from tools import profile_if_possible, enable_debug_on_crash def pick_seed(): @@ -18,12 +21,77 @@ def pick_seed(): 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): + self.detector = detector + + if seed is None: + self.seed = pick_seed() + else: + self.seed = seed + print >>sys.stderr, 'RNG seed:', self.seed + # We have three generators to seed: numpy.random, GEANT4, and CURAND. + # 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) + else: + self.photon_generator = None + + print >>sys.stderr, 'Creating BVH with %d bits...' % (bvh_bits) + detector.build(bits=bvh_bits) + + print >>sys.stderr, 'Initializing GPU...' + self.gpu_worker = gpu.GPU(cuda_device) + + print >>sys.stderr, 'Loading detector onto GPU...' + self.gpu_worker.load_geometry(detector) + + print >>sys.stderr, 'Initializing random numbers generators...' + self.gpu_worker.setup_propagate(seed=self.seed) + self.gpu_worker.setup_daq(max(self.detector.pmtids)) + + def simulate(self, nevents, vertex_generator, keep_photon_start=False, keep_photon_stop=False, + run_daq=True): + return self.simulate_photons(nevents, self.photon_generator.generate_events(nevents, vertex_generator), + 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() + + if run_daq: + ev.hits = self.gpu_worker.get_hits() + + yield ev + + def propagate_photons(self, photons, max_steps=10): + self.gpu_worker.load_photons(photons) + self.gpu_worker.propagate() + return self.gpu_worker.get_photons() + @profile_if_possible def main(): parser = optparse.OptionParser('%prog') - parser.add_option('-b', type='int', dest='nbits', default=10) + parser.add_option('-b', type='int', dest='nbits', default=11) parser.add_option('-j', type='int', dest='device', default=None) - parser.add_option('-n', type='int', dest='nblocks', default=64) parser.add_option('-s', type='int', dest='seed', default=None, help='Set random number generator seed') parser.add_option('-g', type='int', dest='ngenerators', default=4, @@ -54,14 +122,8 @@ def main(): position = np.array(eval(options.pos), dtype=float) direction = np.array(eval(options.dir), dtype=float) + print >>sys.stderr, 'Loading detector %s...' % options.detector detector = detectors.find(options.detector) - if options.seed is None: - options.seed = pick_seed() - - print >>sys.stderr, 'RNG seed:', options.seed - # We have three generators to seed: numpy.random, GEANT4, and CURAND. - # The latter two are done below - np.random.seed(options.seed) print >>sys.stderr, 'Creating generator...' if options.particle == 'pi0': @@ -73,51 +135,30 @@ def main(): position=constant(position), direction=constant(direction), total_energy=constant(options.energy)) - detector_material = optics.water_wcsim - photon_generator = generator.photon.G4ParallelGenerator(options.ngenerators, detector_material, - base_seed=options.seed) - print >>sys.stderr, 'WARNING: ASSUMING DETECTOR IS WCSIM WATER!!' - - # Do this now so we can get ahead of the photon propagation - print >>sys.stderr, 'Starting GEANT4 generators...' - event_iterator = photon_generator.generate_events(options.nevents, vertex_generator) - - print >>sys.stderr, 'Creating BVH for detector "%s" with %d bits...' % (options.detector, options.nbits) - detector.build(bits=options.nbits) - - print >>sys.stderr, 'Initializing GPU...' - gpu_worker = gpu.GPU(options.device) - - print >>sys.stderr, 'Loading detector onto GPU...' - gpu_worker.load_geometry(detector) - - print >>sys.stderr, 'Initializing random numbers generators...' - gpu_worker.setup_propagate(seed=options.seed) - gpu_worker.setup_daq(max(detector.pmtids)) + + # 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) # 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) + + print >>sys.stderr, 'Starting simulation...' start_sim = time.time() nphotons = 0 for i, ev in enumerate(event_iterator): - photons = ev.photon_start - assert len(photons.positions) > 0, 'GEANT4 generated event with no photons!' - - nphotons += len(photons.positions) - - gpu_worker.load_photons(photons) - - gpu_worker.propagate() - gpu_worker.run_daq() + assert ev.nphoton > 0, 'GEANT4 generated event with no photons!' + nphotons += ev.nphoton - ev.hits = gpu_worker.get_hits() - if not options.save_photon_start: - ev.photon_start = None - if options.save_photon_stop: - ev.photon_stop = gpu_worker.get_photons() writer.write_event(ev) if i % 10 == 0: |