From 9cdaa78716037257d6cad20e961db045c303e640 Mon Sep 17 00:00:00 2001 From: Stan Seibert Date: Fri, 12 Aug 2011 18:04:06 -0400 Subject: Allow multiple GEANT4 generator threads to keep up with the speed of photon propagation. Defaults to 4. --- sim.py | 51 +++++++++++++++++++++++++++++++++++++++------------ 1 file changed, 39 insertions(+), 12 deletions(-) (limited to 'sim.py') 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!' -- cgit