merge heads

1 files changed, 88 insertions, 45 deletions

diff --git a/sim.py b/sim.py
index 7520caf..13e53c6 100755
--- a/sim.py
+++ b/sim.py
@@ -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()