Switch sim.py to spawn a separate process for GEANT4 and start it generating

photons while the detector geometry is being built.

1 files changed, 46 insertions, 42 deletions

diff --git a/sim.py b/sim.py
index b8deef6..44fb1dd 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,31 @@ def info(type, value, tb):
 
 sys.excepthook = info
 
-if __name__ == '__main__':
+class GeneratorThread(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 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)
@@ -59,11 +83,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 = GeneratorThread(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 +107,29 @@ 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()
-
-        # write out this event
+        hit_times = gpu_worker.get_hits()
         root.fill_event(T, ev, i-1, position, len(hit_times), hit_times)
         
         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
 
@@ -138,3 +139,6 @@ if __name__ == '__main__':
     print >>sys.stderr, 'Done. %1.1f events/sec, %1.0f photons/sec.' % (options.nevents/(end_sim - start_sim), nphotons/(end_sim - start_sim))
 
     gpu_worker.shutdown()
+
+if __name__ == '__main__':
+    main()