merge

1 files changed, 33 insertions, 116 deletions

diff --git a/sim.py b/sim.py
index 4742ec9..a7dc220 100755
--- a/sim.py
+++ b/sim.py
@@ -3,93 +3,22 @@ import sys
 import optparse
 import time
 import os
-import multiprocessing
+import numpy as np
 
 import detectors
 import optics
+import generator
+from generator import constant
 import gpu
-import g4gen
 from fileio import root
-import numpy as np
-import math
-import ROOT
+from tools import profile_if_possible, enable_debug_on_crash
 
 def pick_seed():
     '''Returns a seed for a random number generator selected using
     a mixture of the current time and the current process ID.'''
     return int(time.time()) ^ (os.getpid() << 16)
 
-def info(type, value, tb):
-    if hasattr(sys, 'ps1') or not sys.stderr.isatty():
-       # we are in interactive mode or we don't have a tty-like
-       # device, so we call the default hook
-       sys.__excepthook__(type, value, tb)
-    else:
-       import traceback, pdb
-       # we are NOT in interactive mode, print the exception...
-       traceback.print_exception(type, value, tb)
-       print
-       # ...then start the debugger in post-mortem mode.
-       pdb.pm()
-
-class GeneratorProcess(multiprocessing.Process):
-    def __init__(self, particle, energy, position, direction, nevents, material, 
-                 queue, seed=None):
-        multiprocessing.Process.__init__(self)
-
-        self.particle = particle
-        self.energy = energy
-        self.position = position
-        self.direction = direction
-        self.nevents = nevents
-        self.material = material
-        self.seed = seed
-        self.queue = queue
-        self.daemon = True
-
-    def run(self):
-        print >>sys.stderr, 'Starting generator thread...'
-        generator = g4gen.G4Generator(self.material, seed=self.seed)
-
-        for i in xrange(self.nevents):
-            if self.particle == 'pi0':
-                photons = generator.generate_pi0(total_energy=self.energy,
-                                                 position=self.position,
-                                                 direction=self.direction)
-            else:
-                photons = generator.generate_photons(particle_name=self.particle, 
-                                                     total_energy=self.energy,
-                                                     position=self.position,
-                                                     direction=self.direction)
-            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:
-    profile = profile
-except NameError:
-    profile = lambda x: x
-
-@profile
+@profile_if_possible
 def main():
     parser = optparse.OptionParser('%prog')
     parser.add_option('-b', type='int', dest='nbits', default=10)
@@ -130,27 +59,28 @@ def main():
         options.seed = pick_seed()
 
     print >>sys.stderr, 'RNG seed:', options.seed
+    # We have three generators to seed: numpy.random, GEANT4, and CURAND.
+    # The latter two are done below
+    np.random.seed(options.seed)
 
     print >>sys.stderr, 'Creating generator...'
+    if options.particle == 'pi0':
+        vertex_generator = generator.vertex.pi0_gun(pi0_position=constant(position),
+                                                    pi0_direction=constant(direction),
+                                                    pi0_total_energy=constant(options.energy))
+    else:
+        vertex_generator = generator.vertex.particle_gun(particle_name=constant(options.particle),
+                                                         position=constant(position),
+                                                         direction=constant(direction),
+                                                         total_energy=constant(options.energy))
     detector_material = optics.water_wcsim
-    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))]
-    
+    photon_generator = generator.photon.G4ParallelGenerator(options.ngenerators, detector_material,
+                                                           base_seed=options.seed)
     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 generators...'
-    for generator in generators:
-        generator.start()
+    event_iterator = photon_generator.generate_events(options.nevents, vertex_generator)
 
     print >>sys.stderr, 'Creating BVH for detector "%s" with %d bits...' % (options.detector, options.nbits)
     detector.build(bits=options.nbits)
@@ -168,41 +98,28 @@ def main():
     # Create output file
     writer = root.RootWriter(output_filename)
 
-    # Set generator info
-    writer.set_generated_particle(name=options.particle, position=position,
-                                  direction=direction, total_e=options.energy)
-
     print >>sys.stderr, 'Starting simulation...'
     start_sim = time.time()
     nphotons = 0
-    for i in xrange(options.nevents):
-        photons = queue.get()
 
-        assert len(photons['pos']) > 0, 'GEANT4 generated event with no photons!'
+    for i, ev in enumerate(event_iterator):
+        photons = ev.photon_start
+        assert len(photons.positions) > 0, 'GEANT4 generated event with no photons!'
 
-        nphotons += len(photons['pos'])
+        nphotons += len(photons.positions)
+
+        gpu_worker.load_photons(photons)
 
-        gpu_worker.load_photons(pos=photons['pos'], dir=photons['dir'], pol=photons['pol'],
-                                t0=photons['t0'], wavelength=photons['wavelength'])
         gpu_worker.propagate()
         gpu_worker.run_daq()
-        hits = gpu_worker.get_hits()
 
-        if options.save_photon_start:
-            photon_start = photons
-        else:
-            photon_start = None
+        ev.hits = gpu_worker.get_hits()        
+        if not options.save_photon_start:
+            ev.photon_start = None
         if options.save_photon_stop:
-            photon_stop = gpu_worker.get_photons()
-        else:
-            photon_stop = None
-        
-        if 'subtracks' in photons:
-            subtracks = photons['subtracks']
-        else:
-            subtracks = None
-        writer.write_event(i, hits, photon_start=photon_start, photon_stop=photon_stop,
-                           subtracks=subtracks)
+            ev.photon_stop = gpu_worker.get_photons()
+        writer.write_event(ev)
+
         if i % 10 == 0:
             print >>sys.stderr, "\rEvent:", i,
 
@@ -214,5 +131,5 @@ def main():
     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__':
-    sys.excepthook = info
+    enable_debug_on_crash()
     main()