Allow multiple GEANT4 generator threads to keep up with the speed of photon propagation. Defaults to 4.

1 files changed, 39 insertions, 12 deletions

diff --git a/sim.py b/sim.py
index b8382e6..4742ec9 100755
--- a/sim.py
+++ b/sim.py
@@ -33,7 +33,8 @@ def info(type, value, tb):
        pdb.pm()
 
 class GeneratorProcess(multiprocessing.Process):
-    def __init__(self, particle, energy, position, direction, nevents, material, seed=None):
+    def __init__(self, particle, energy, position, direction, nevents, material, 
+                 queue, seed=None):
         multiprocessing.Process.__init__(self)
 
         self.particle = particle
@@ -43,7 +44,7 @@ class GeneratorProcess(multiprocessing.Process):
         self.nevents = nevents
         self.material = material
         self.seed = seed
-        self.queue = multiprocessing.Queue()
+        self.queue = queue
         self.daemon = True
 
     def run(self):
@@ -63,6 +64,24 @@ class GeneratorProcess(multiprocessing.Process):
             self.queue.put(photons)
 
 
+def partition(num, partitions):
+    '''Generator that returns num//partitions, with the last item including the remainder.
+
+    Useful for partitioning a number into mostly equal parts while preserving the sum.
+
+    >>> list(partition(800, 3))
+    [266, 266, 268]
+    >>> sum(list(partition(800, 3)))
+    800
+    '''
+    step = num // partitions
+    for i in xrange(partitions):
+        if i < partitions - 1:
+            yield step
+        else:
+            yield step + (num % partitions)
+
+
 
 # Allow profile decorator to exist, but do nothing if not running under kernprof
 try:
@@ -78,6 +97,8 @@ def main():
     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,
+                      help='Number of GEANT4 generator processes')
     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-')
@@ -112,18 +133,24 @@ def main():
 
     print >>sys.stderr, 'Creating generator...'
     detector_material = optics.water_wcsim
-    generator_thread = GeneratorProcess(particle=options.particle, 
-                                        energy=options.energy,
-                                        position=position,
-                                        direction=direction,
-                                        nevents=options.nevents,
-                                        material=detector_material,
-                                        seed=options.seed)
+    queue = multiprocessing.Queue()
+    generators = [GeneratorProcess(particle=options.particle, 
+                                   energy=options.energy,
+                                   position=position,
+                                   direction=direction,
+                                   nevents=nevents,
+                                   material=detector_material,
+                                   seed=options.seed + seed_offset,
+                                   queue=queue)
+                  for seed_offset, nevents in 
+                  enumerate(partition(options.nevents, options.ngenerators))]
+    
     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 generator...'
-    generator_thread.start()
+    print >>sys.stderr, 'Starting GEANT4 generators...'
+    for generator in generators:
+        generator.start()
 
     print >>sys.stderr, 'Creating BVH for detector "%s" with %d bits...' % (options.detector, options.nbits)
     detector.build(bits=options.nbits)
@@ -149,7 +176,7 @@ def main():
     start_sim = time.time()
     nphotons = 0
     for i in xrange(options.nevents):
-        photons = generator_thread.queue.get()
+        photons = queue.get()
 
         assert len(photons['pos']) > 0, 'GEANT4 generated event with no photons!'