diff options
Diffstat (limited to 'sim.py')
| -rwxr-xr-x | sim.py | 133 |
1 files changed, 88 insertions, 45 deletions
@@ -2,6 +2,7 @@ import sys import optparse import time +import multiprocessing import detectors import optics @@ -27,8 +28,54 @@ def info(type, value, tb): sys.excepthook = info -if __name__ == '__main__': - +class GeneratorProcess(multiprocessing.Process): + def __init__(self, particle, energy, position, direction, nevents, material): + multiprocessing.Process.__init__(self) + + self.particle = particle + self.energy = energy + self.position = position + self.direction = direction + self.nevents = nevents + self.material = material + self.queue = multiprocessing.Queue() + + def run(self): + print >>sys.stderr, 'Starting generator thread...' + generator = g4gen.G4Generator(self.material) + + for i in xrange(self.nevents): + photons = generator.generate_photons(particle_name=self.particle, + total_energy=self.energy, + position=self.position, + direction=self.direction) + self.queue.put(photons) + + +def write_event(T, ev, event_id, hits, photon_start=None, photon_stop=None): + ev.event_id = event_id + if photon_start is not None: + photons = photon_start + root.fill_photons(ev, True, + len(photons['pos']), + np.ravel(photons['pos']), + np.ravel(photons['dir']), + np.ravel(photons['pol']), + photons['wavelength'], photons['t0']) + if photon_stop is not None: + photons = photon_stop + root.fill_photons(ev, False, + len(photons['pos']), + np.ravel(photons['pos']), + np.ravel(photons['dir']), + np.ravel(photons['pol']), + photons['wavelength'], photons['t0'], + photons['histories'], photons['last_hit_triangles']) + + root.fill_hits(ev, len(hits), hits) + T.Fill() + +def main(): parser = optparse.OptionParser('%prog') parser.add_option('-b', type='int', dest='nbits', default=10) parser.add_option('-j', type='int', dest='device', default=None) @@ -39,7 +86,13 @@ if __name__ == '__main__': 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='(0,1,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, + help='Save initial photon vertices to disk') + parser.add_option('--save-photon-stop', action='store_true', + dest='save_photon_stop', default=False, + help='Save final photon vertices to disk') options, args = parser.parse_args() if len(args) != 1: @@ -59,11 +112,20 @@ if __name__ == '__main__': print >>sys.stderr, 'Creating detector...' detector.build(bits=options.nbits) - print >>sys.stderr, 'Initializing generator...' + print >>sys.stderr, 'Creating generator...' detector_material = optics.water - generator = g4gen.G4Generator(detector_material) + generator_thread = GeneratorProcess(particle=options.particle, + energy=options.energy, + position=position, + direction=direction, + nevents=options.nevents, + material=detector_material) print >>sys.stderr, 'WARNING: ASSUMING DETECTOR IS WATER!!' + # Do this now so we can get ahead of the photon propagation + print >>sys.stderr, 'Starting GEANT4 generator...' + generator_thread.start() + print >>sys.stderr, 'Initializing GPU...' gpu_worker = gpu.GPU(options.device) @@ -74,61 +136,39 @@ if __name__ == '__main__': gpu_worker.setup_propagate() gpu_worker.setup_daq(max(detector.pmtids)) - print >>sys.stderr, 'Loading tables in GEANT4...' - # Do easy event to force tables to load, throw away photons - generator.generate_photons(particle_name='e-', - total_energy=1.5, - position=(0,0,0), - direction=(1,0,0)) - - # Create output file f = ROOT.TFile(output_filename, 'RECREATE') ev, T = root.make_tree('T') print >>sys.stderr, 'Starting simulation...' start_sim = time.time() - #### Do the generation and writing in this offset order to ensure - #### that propagation on the GPU overlaps with CPU work to create - #### and save photons. - #### WARNING: THIS MAKES THE CODE LOOK A LITTLE CRAZY. I AM SORRY - - # Do first event - photons = generator.generate_photons(particle_name=options.particle, - total_energy=options.energy, - position=position, - direction=direction) - nphotons = len(photons['pos']) - gpu_worker.load_photons(**photons) - gpu_worker.propagate() - gpu_worker.run_daq() - - for i in xrange(1, options.nevents-1): - # photons for next event while previous event propagates on GPU - photons = generator.generate_photons(particle_name=options.particle, - total_energy=options.energy, - position=position, - direction=direction) + nphotons = 0 + for i in xrange(options.nevents): + photons = generator_thread.queue.get() + + assert len(photons['pos']) > 0, 'GEANT4 generated event with no photons!' + nphotons += len(photons['pos']) - - # this will stop and wait for event to finish - hit_times = gpu_worker.get_hits() - # turn around next event gpu_worker.load_photons(**photons) gpu_worker.propagate() gpu_worker.run_daq() + hit_times = gpu_worker.get_hits() - # write out this event - root.fill_event(T, ev, i-1, position, len(hit_times), hit_times) - + if options.save_photon_start: + photon_start = photons + else: + photon_start = None + if options.save_photon_stop: + photon_stop = gpu_worker.get_photons() + else: + photon_stop = None + + write_event(T, ev, i, hit_times, + photon_start=photon_start, photon_stop=photon_stop) if i % 10 == 0: print >>sys.stderr, "\rEvent:", i, - # Get results for last event and write it out - hit_times = gpu_worker.get_hits() - root.fill_event(T, ev, options.nevents - 1, position, len(hit_times), hit_times) - end_sim = time.time() print >>sys.stderr, "\rEvent:", options.nevents - 1 @@ -136,3 +176,6 @@ if __name__ == '__main__': f.Close() print >>sys.stderr, 'Done. %1.1f events/sec, %1.0f photons/sec.' % (options.nevents/(end_sim - start_sim), nphotons/(end_sim - start_sim)) + +if __name__ == '__main__': + main() |