merge

13 files changed, 375 insertions, 254 deletions

diff --git a/camera.py b/camera.py
index d7c2038..fa38060 100755
--- a/camera.py
+++ b/camera.py
@@ -457,7 +457,7 @@ class EventViewer(Camera):
         hit = hit.astype(np.bool)
 
         # Important: Compute range only with HIT channels
-        solid_colors = map_to_color(t, range=(t[hit].min(),t[hit].mean()))
+        solid_colors = map_to_color(q, range=(q[hit].min(),q[hit].max()))
         self.gpu.color_solids(hit, solid_colors)
         self.update()
 
diff --git a/event.py b/event.py
new file mode 100644
index 0000000..bf7c007
--- /dev/null
+++ b/event.py
@@ -0,0 +1,51 @@
+import numpy as np
+
+class Photons(object):
+    def __init__(self, positions, directions, polarizations, times, wavelengths,
+                 last_hit_triangles=None, histories=None):
+        self.positions = positions
+        self.directions = directions
+        self.polarizations = polarizations
+        self.times = times
+        self.wavelengths = wavelengths
+        self.last_hit_triangles = last_hit_triangles
+        self.histories = histories
+
+def concatenate_photons(photons):
+    '''Merge a list of Photons objects into one long list of photons'''
+    return Photons(positions=np.concatenate([p.positions for p in photons]),
+                   directions=np.concatenate([p.directions for p in photons]),
+                   polarizations=np.concatenate([p.polarizations for p in photons]),
+                   times=np.concatenate([p.times for p in photons]),
+                   wavelengths=np.concatenate([p.wavelengths for p in photons]))
+
+class Channels(object):
+    def __init__(self, hit, t, q, histories=None):
+        self.hit = hit
+        self.t = t
+        self.q = q
+        self.histories=histories
+
+    def hit_channels(self):
+        return self.hit.nonzero(), self.t[self.hit], self.q[self.hit]
+
+class Subtrack(object):
+    def __init__(self, particle_name, position, direction, start_time, total_energy):
+        self.particle_name = particle_name
+        self.position = position
+        self.direction = direction
+        self.start_time = start_time
+        self.total_energy = total_energy
+
+class Event(object):
+    def __init__(self, event_id, particle_name=None, gen_position=None, gen_direction=None, gen_total_energy=None, ):
+        self.event_id = event_id
+        self.particle_name = particle_name
+        self.gen_position = gen_position
+        self.gen_direction = gen_direction
+        self.gen_total_energy = gen_total_energy
+
+        self.subtracks = []
+        self.photon_start = None
+        self.photon_stop = None
+        self.channels = None
diff --git a/fileio/root.C b/fileio/root.C
index 9a959f5..8f7d0a9 100644
--- a/fileio/root.C
+++ b/fileio/root.C
@@ -18,6 +18,7 @@ struct Track {
   double t;
   TVector3 pos;
   TVector3 dir;
+  double start_time;
   double total_energy;
 };
 
diff --git a/fileio/root.py b/fileio/root.py
index 004484f..5bebda0 100644
--- a/fileio/root.py
+++ b/fileio/root.py
@@ -3,7 +3,8 @@ import os.path
 
 ROOT.gROOT.ProcessLine('.L '+os.path.join(os.path.dirname(__file__), 'root.C+g'))
 
-from ROOT import Event
+import ROOT
+import event
 
 class RootWriter(object):
     def __init__(self, filename):
@@ -14,42 +15,43 @@ class RootWriter(object):
         self.ev = ROOT.Event()
         self.T.Branch('ev', self.ev)
 
-    def set_generated_particle(self, name, position, direction, total_e):
-        self.ev.mc.particle = name
-        self.ev.mc.gen_pos.SetXYZ(*position)
-        self.ev.mc.gen_dir.SetXYZ(*direction)
-        self.ev.mc.gen_total_e = total_e
+    def write_event(self, pyev):
+        '''Write an event.Event object to the ROOT tree as a ROOT.Event object.'''
+        self.ev.event_id = pyev.event_id
+
+        self.ev.mc.particle = pyev.particle_name
+        self.ev.mc.gen_pos.SetXYZ(*pyev.gen_position)
+        self.ev.mc.gen_dir.SetXYZ(*pyev.gen_direction)
+        self.ev.mc.gen_total_energy = pyev.gen_total_energy
 
-    def write_event(self, event_id, hits, photon_start=None, photon_stop=None, subtracks=None):
-        self.ev.event_id = event_id
-        if photon_start is not None:
-            photons = photon_start
+        if pyev.photon_start is not None:
+            photons = pyev.photon_start
             ROOT.fill_photons(self.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:
+                              len(photons.positions), 
+                              np.ravel(photons.positions),
+                              np.ravel(photons.directions),
+                              np.ravel(photons.polarizations),
+                              photons.wavelengths, photons.times)
+        if pyev.photon_stop is not None:
             photons = photon_stop
-            ROOT.fill_photons(self.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_photons(self.ev, True,
+                              len(photons.positions), 
+                              np.ravel(photons.positions),
+                              np.ravel(photons.directions),
+                              np.ravel(photons.polarizations),
+                              photons.wavelengths, photons.times)
 
         self.ev.mc.subtrack.resize(0)
-        if subtracks is not None:
-            self.ev.mc.subtrack.resize(len(subtracks))
-            for i, subtrack in enumerate(subtracks):
-                self.ev.mc.subtrack[i].name = subtrack['name']
-                self.ev.mc.subtrack[i].pos.SetXYZ(*subtrack['pos'])
-                self.ev.mc.subtrack[i].dir.SetXYZ(*subtrack['dir'])
-                self.ev.mc.subtrack[i].total_energy = subtrack['total_e']
-
-        ROOT.fill_hits(self.ev, len(hits['t']), hits['t'], hits['q'], hits['history'])
+        if pyev.subtracks is not None:
+            self.ev.mc.subtrack.resize(len(pyev.subtracks))
+            for i, subtrack in enumerate(pyev.subtracks):
+                self.ev.mc.subtrack[i].name = subtrack.particle_name
+                self.ev.mc.subtrack[i].pos.SetXYZ(*subtrack.position)
+                self.ev.mc.subtrack[i].dir.SetXYZ(*subtrack.direction)
+                self.ev.mc.subtrack[i].start_time = subtrack.start_time
+                self.ev.mc.subtrack[i].total_energy = subtrack.total_energy
+
+        ROOT.fill_hits(self.ev, len(pyev.hits.t), pyev.hits.t, pyev.hits.q, pyev.hits.histories)
         self.T.Fill()
 
     def close(self):
diff --git a/G4chroma.cc b/generator/G4chroma.cc
index 1aba133..1aba133 100644
--- a/G4chroma.cc
+++ b/generator/G4chroma.cc
diff --git a/G4chroma.hh b/generator/G4chroma.hh
index 4f085aa..4f085aa 100644
--- a/G4chroma.hh
+++ b/generator/G4chroma.hh
diff --git a/generator/__init__.py b/generator/__init__.py
new file mode 100644
index 0000000..c3179d0
--- /dev/null
+++ b/generator/__init__.py
@@ -0,0 +1,5 @@
+from vertex import *
+from event import *
+from photon import *
+
+
diff --git a/g4gen.py b/generator/g4gen.py
index 6766d42..d4e79a1 100644
--- a/g4gen.py
+++ b/generator/g4gen.py
@@ -4,7 +4,7 @@ import g4py.NISTmaterials
 import g4py.ParticleGun
 import pyublas
 import numpy as np
-import pi0
+import event
 
 try:
     import G4chroma
@@ -71,58 +71,7 @@ class G4Generator(object):
         g4material.SetMaterialPropertiesTable(prop_table)
         return g4material
 
-    def generate_pi0(self, total_energy, position, direction):
-        '''Use GEANT4 to generate photons produced by a pi0 decay.
-
-        pi0 event will be generated by running two gammas with the appropriate
-        energy and angular distribution.
-
-               total_energy: Total energy of pi0 (incl rest mass) in MeV
-               position: 3-tuple of position of particle in global coordinates
-               direction: 3-tuple direction vector.  
-                          Does not have to be normalized.
-        '''
-        direction = direction / np.linalg.norm(direction)
-        
-        cos_theta_rest = np.random.random_sample() * 2 - 1
-        theta_rest = np.arccos(cos_theta_rest)
-        phi_rest = np.random.random_sample() * 2 * np.pi
-
-        (gamma1_e, gamma1_dir), (gamma2_e, gamma2_dir) = pi0.pi0_decay(energy=total_energy,
-                                                                       direction=direction,
-                                                                       theta=theta_rest,
-                                                                       phi=phi_rest)
-
-        gamma1 = self.generate_photons('gamma', gamma1_e, position, gamma1_dir)
-        gamma2 = self.generate_photons('gamma', gamma2_e, position, gamma2_dir)
-
-        photons = {'subtracks' : [ {'name':'gamma', 'pos': position, 't0' : 0.0,
-                                    'dir' : gamma1_dir, 'total_e' : gamma1_e },
-                                   {'name':'gamma', 'pos': position, 't0' : 0.0,
-                                    'dir' : gamma2_dir, 'total_e' : gamma2_e } ]
-                    }
-        for key in gamma1.keys():
-            photons[key] = np.concatenate((gamma1[key], gamma2[key]))
-                                                                       
-        return photons
-
-    def generate_photons(self, particle_name, total_energy, position, direction):
-        '''Use GEANT4 to generate photons produced by the given particle.
-           
-               particle_name: GEANT4 name of particle.  'e-', 'mu-', etc
-               total_energy: Total energy of particle (incl rest mass) in MeV
-               position: 3-tuple of position of particle in global coordinates
-               direction: 3-tuple direction vector.  
-                          Does not have to be normalized.
-        '''
-        self.particle_gun.SetParticleByName(particle_name)
-        self.particle_gun.SetParticleEnergy(total_energy * MeV)
-        self.particle_gun.SetParticlePosition(G4ThreeVector(*position))
-        self.particle_gun.SetParticleMomentumDirection(G4ThreeVector(*direction).unit())
-
-        self.tracking_action.Clear()
-
-        gRunManager.BeamOn(1)
+    def _extract_photons_from_tracking_action(self):
         n = self.tracking_action.GetNumPhotons()        
         pos = np.zeros(shape=(n,3), dtype=np.float32)
         pos[:,0] = self.tracking_action.GetX()
@@ -139,27 +88,60 @@ class G4Generator(object):
         pol[:,1] = self.tracking_action.GetPolY()
         pol[:,2] = self.tracking_action.GetPolZ()
         
-        wavelength = self.tracking_action.GetWavelength().astype(np.float32)
-        t0 = self.tracking_action.GetT0().astype(np.float32)
-        return { 'pos' : pos, 
-                 'dir' : dir, 
-                 'pol' : pol,
-                 'wavelength' : wavelength, 
-                 't0' : t0 }
+        wavelengths = self.tracking_action.GetWavelength().astype(np.float32)
+        times = self.tracking_action.GetT0().astype(np.float32)
+
+        return event.Photons(positions=pos, directions=dir, polarizations=pol, times=times, wavelengths=wavelengths)
+
+    def generate_photons(self, ev):
+        '''Use GEANT4 to generate photons produced by the given particle.
+           
+               ev: a generator.event.Event object with the particle
+                   properties set.  If it contains subtracks, those
+                   will be used to create the photon vertices rather
+                   than the main particle.
+
+               Returns an instance of event.Photons containing the
+               generated photon vertices for the primary particle or
+               all the subtracks, if present.
+        '''
+        photons = []
+        if ev.subtracks:
+            subtracks = ev.subtracks
+        else:
+            # Create temporary subtrack for single primary particle
+            subtracks = [event.Subtrack(particle_name=ev.particle_name,
+                                        position=ev.gen_pos,
+                                        direction=ev.gen_dir,
+                                        start_time=0.0,
+                                        total_energy=ev.gen_total_energy)]
+
+        for subtrack in subtracks:
+            self.particle_gun.SetParticleByName(subtrack.particle_name)
+            self.particle_gun.SetParticleEnergy(subtrack.total_energy * MeV)
+            self.particle_gun.SetParticlePosition(G4ThreeVector(*subtrack.position)*m)
+            self.particle_gun.SetParticleMomentumDirection(G4ThreeVector(*subtrack.direction).unit())
+            
+            self.tracking_action.Clear()
+            gRunManager.BeamOn(1)
+            photons.append(self._extract_photons_from_tracking_action())
+
+        # Merge all photon lists into one big list
+        return event.concatenate_photons(photons)
 
 if __name__ == '__main__':
     import time
     import optics
     gen = G4Generator(optics.water)
     # prime things
-    gen.generate_photons('e-', 1, (0,0,0), (1,0,0))
+    gen.generate_photons(event.Event('e-', (0,0,0), (1,0,0), 1.0))
 
-    start = time.time()   
+    start = time.time()
     n = 0
     for i in xrange(100):
-        photons = gen.generate_photons('mu-', 700, (0,0,0), (1,0,0))
-        n += len(photons['t0'])
-        print photons['pos'][0].min(), photons['pos'][0].max()
+        photons = gen.generate_photons(event.Event('mu-', (0,0,0), (1,0,0), 1.0))
+        n += len(photons.times)
+        print photons.positions[0].min(), photons.positions[0].max()
     stop = time.time()
     print stop - start, 'sec'
     print n / (stop-start), 'photons/sec'
diff --git a/generator/photon.py b/generator/photon.py
new file mode 100644
index 0000000..dcfc292
--- /dev/null
+++ b/generator/photon.py
@@ -0,0 +1,74 @@
+import g4gen
+import multiprocessing
+import threading
+import numpy as np
+import itertools
+
+class G4GeneratorProcess(multiprocessing.Process):
+    def __init__(self, idnum, material, input_queue, output_queue, seed=None):
+        multiprocessing.Process.__init__(self)
+
+        self.idnum = idnum
+        self.material = material
+        self.input_queue = input_queue
+        self.output_queue = output_queue
+        self.seed = seed
+        self.daemon = True
+
+    def run(self):
+        gen = g4gen.G4Generator(self.material, seed=self.seed)
+        while True:
+            ev = self.input_queue.get()
+            ev.photon_start = gen.generate_photons(ev)
+            self.output_queue.put(ev)
+
+class G4VertexGeneratorProcess(multiprocessing.Process):
+    def __init__(self, vertex_iterator, queue):
+        multiprocessing.Process.__init__(self)
+        self.vertex_iterator = vertex_iterator
+        self.queue = queue
+
+    def run(self):
+        for ev in self.vertex_iterator:
+            self.queue.put(ev)
+
+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)
+
+def queue_iterator(nelements, queue):
+    for i in xrange(nelements):
+        yield queue.get()
+
+class G4ParallelGenerator(object):
+    def __init__(self, nprocesses, material, base_seed=None):
+        self.material = material
+        self.vertex_queue = multiprocessing.Queue()
+        self.photon_queue = multiprocessing.Queue()
+        self.processes = [ G4GeneratorProcess(i, material, self.vertex_queue, self.photon_queue, seed=base_seed + i)
+                           for i in xrange(nprocesses) ]
+        for p in self.processes:
+            p.start()
+
+    def generate_events(self, nevents, vertex_iterator):
+        # Doing this to avoid a deadlock caused by putting to one queue while getting from another
+        for ev in itertools.islice(vertex_iterator, nevents):
+            self.vertex_queue.put(ev)
+
+        return queue_iterator(nevents, self.photon_queue)
+
+        
+            
diff --git a/generator/vertex.py b/generator/vertex.py
new file mode 100644
index 0000000..1fa61cb
--- /dev/null
+++ b/generator/vertex.py
@@ -0,0 +1,71 @@
+import itertools
+import numpy as np
+from math import sin, cos, sqrt
+import copy
+
+import pi0
+from event import Event, Subtrack
+
+# generator parts for use with gun()
+
+def constant(obj):
+    while True:
+        yield obj
+
+def isotropic():
+    while True:
+        cos_theta = np.random.uniform(-1.0, 1.0)
+        sin_theta = sqrt(1.0 - cos_theta**2)
+        phi = np.random.uniform(0.0, 2*np.pi)
+        direction = np.array([sin_theta * cos(phi),
+                              sin_theta * sin(phi),
+                              cos_theta])
+        yield direction
+        
+def line_segment(point1, point2):
+    while True:
+        frac = np.random.uniform(0.0, 1.0)
+        yield frac * point1 + (1.0 - frac) * point2
+
+def fill_shell(center, radius):
+    for direction in isotropic():
+        r = radius * np.random.uniform(0.0, 1.0)**(1.0/3.0)
+        yield r * direction
+
+def flat(e_lo, e_hi):
+    while True:
+        yield np.random.uniform(e_lo, e_hi)
+
+# vertex generators
+
+def particle_gun(particle_name, position, direction, total_energy, start_id=0):
+    for i, ev_particle, ev_position, ev_direction, ev_total_energy in \
+            itertools.izip(itertools.count(start_id), 
+                           particle_name, position, direction, total_energy):
+        ev = Event(event_id=i, particle_name=ev_particle, 
+                   gen_position=ev_position, gen_direction=ev_direction, gen_total_energy=ev_total_energy)
+        yield ev
+
+def pi0_gun(pi0_position, pi0_direction, pi0_total_energy, start_id=0):
+    for i, ev_position, ev_direction, ev_total_energy in \
+            itertools.izip(itertools.count(start_id), pi0_position, pi0_direction, pi0_total_energy):
+
+        ev = Event(event_id=i, particle_name='pi0',
+                   gen_position=ev_position, gen_direction=ev_direction, gen_total_energy=ev_total_energy)
+
+        ev_direction = ev_direction / np.linalg.norm(ev_direction)
+        
+        cos_theta_rest = np.random.random_sample() * 2 - 1
+        theta_rest = np.arccos(cos_theta_rest)
+        phi_rest = np.random.random_sample() * 2 * np.pi
+
+        (gamma1_e, gamma1_dir), (gamma2_e, gamma2_dir) = pi0.pi0_decay(energy=ev_total_energy,
+                                                                       direction=ev_direction,
+                                                                       theta=theta_rest,
+                                                                       phi=phi_rest)
+
+        ev.subtracks = [Subtrack(particle_name='gamma', position=ev_position, direction=gamma1_dir,
+                                 start_time=0.0, total_energy=gamma1_e),
+                        Subtrack(particle_name='gamma', position=ev_position, direction=gamma2_dir,
+                                 start_time=0.0, total_energy=gamma2_e)]
+        yield ev
diff --git a/gpu.py b/gpu.py
index 7e45c3d..327e692 100644
--- a/gpu.py
+++ b/gpu.py
@@ -13,6 +13,7 @@ import src
 from geometry import standard_wavelengths
 from color import map_to_color
 import sys
+import event
 
 cuda.init()
 
@@ -232,40 +233,31 @@ class GPU(object):
                                  grid=(self.max_blocks,1))
         #self.context.synchronize()
 
-    def load_photons(self, pos, dir, pol, wavelength, t0, 
-                     histories=None, last_hit_triangles=None):
-        '''Load N photons onto the GPU.
-
-           pos: numpy.array(shape=(N, 3)) of photon starting positions (meters)
-           dir: numpy.array(shape=(N, 3)) of photon starting directions (unit vectors)
-           pol: numpy.array(shape=(N, 3)) of photon polarization directions (unit vectors)
-           wavelength: numpy.array(shape=N) of photon wavelengths (nm)
-           t0: numpy.array(shape=N) of photon start times (s)
-           
-           Optional args will be loaded with defaults on GPU if not set:
-           histories: Bitmask of interactions that have occurred over history of photon
-           last_hit_triangles: The triangle ID number that the photon last interacted with,
-                               if any.  -1 if no triangle was hit in the last step
+    def load_photons(self, photons):
+        '''Load photons onto the GPU from an event.Photons object.
+
+           If photon.histories or photon.last_hit_triangles are set to none,
+           they will be initialized to 0 and -1 on the GPU, respectively.
         '''
-        self.nphotons = len(pos)
-        assert len(dir) == self.nphotons
-        assert len(pol) == self.nphotons
-        assert len(wavelength) == self.nphotons
-        assert len(t0) == self.nphotons
-
-        self.positions_gpu = gpuarray.to_gpu(pos.astype(np.float32).view(gpuarray.vec.float3))
-        self.directions_gpu = gpuarray.to_gpu(dir.astype(np.float32).view(gpuarray.vec.float3))
-        self.polarizations_gpu = gpuarray.to_gpu(pol.astype(np.float32).view(gpuarray.vec.float3))
-        self.wavelengths_gpu = gpuarray.to_gpu(wavelength.astype(np.float32))
-        self.times_gpu = gpuarray.to_gpu(t0.astype(np.float32))
-
-        if histories is not None:
-            self.histories_gpu = gpuarray.to_gpu(histories.astype(np.uint32))
+        self.nphotons = len(photons.positions)
+        assert len(photons.directions) == self.nphotons
+        assert len(photons.polarizations) == self.nphotons
+        assert len(photons.wavelengths) == self.nphotons
+        assert len(photons.times) == self.nphotons
+
+        self.positions_gpu = gpuarray.to_gpu(photons.positions.astype(np.float32).view(gpuarray.vec.float3))
+        self.directions_gpu = gpuarray.to_gpu(photons.directions.astype(np.float32).view(gpuarray.vec.float3))
+        self.polarizations_gpu = gpuarray.to_gpu(photons.polarizations.astype(np.float32).view(gpuarray.vec.float3))
+        self.wavelengths_gpu = gpuarray.to_gpu(photons.wavelengths.astype(np.float32))
+        self.times_gpu = gpuarray.to_gpu(photons.times.astype(np.float32))
+
+        if photons.histories is not None:
+            self.histories_gpu = gpuarray.to_gpu(photons.histories.astype(np.uint32))
         else:
             self.histories_gpu = gpuarray.zeros(self.nphotons, dtype=np.uint32)
 
-        if last_hit_triangles is not None:
-            self.last_hit_triangles_gpu = gpuarray.to_gpu(last_hit_triangles.astype(np.int32))
+        if photons.last_hit_triangles is not None:
+            self.last_hit_triangles_gpu = gpuarray.to_gpu(photons.last_hit_triangles.astype(np.int32))
         else:
             self.last_hit_triangles_gpu = gpuarray.empty(self.nphotons, dtype=np.int32)
             self.last_hit_triangles_gpu.fill(-1)
@@ -328,13 +320,13 @@ class GPU(object):
 
         Contents of dictionary have the same names as the parameters to load_photons().
         '''
-        return { 'pos' : self.positions_gpu.get().view(np.float32).reshape(self.positions_gpu.size, 3),
-                 'dir' : self.directions_gpu.get().view(np.float32).reshape(self.directions_gpu.size, 3),
-                 'pol' : self.polarizations_gpu.get().view(np.float32).reshape(self.polarizations_gpu.size, 3),
-                 'wavelength' : self.wavelengths_gpu.get(),
-                 't0' : self.times_gpu.get(),
-                 'histories' : self.histories_gpu.get(),
-                 'last_hit_triangles' : self.last_hit_triangles_gpu.get()}
+        return event.Photons(positions=self.positions_gpu.get().view(np.float32).reshape(self.positions_gpu.size, 3),
+                             directions=self.directions_gpu.get().view(np.float32).reshape(self.directions_gpu.size, 3),
+                             polarizations=self.polarizations_gpu.get().view(np.float32).reshape(self.polarizations_gpu.size, 3),
+                             wavelengths=self.wavelengths_gpu.get(),
+                             times=self.times_gpu.get(),
+                             histories=self.histories_gpu.get(),
+                             last_hit_triangles=self.last_hit_triangles_gpu.get())
     
     def setup_daq(self, max_pmt_id, pmt_rms=1.2e-9):
         self.earliest_time_gpu = gpuarray.GPUArray(shape=(max_pmt_id+1,), dtype=np.float32)
@@ -378,9 +370,11 @@ class GPU(object):
         
 
     def get_hits(self):
-        return { 't': self.earliest_time_gpu.get(), 
-                 'q': self.channel_q_gpu.get().astype(np.float32),
-                 'history': self.channel_history_gpu.get()}
+        t = self.earliest_time_gpu.get()
+        # For now, assume all channels with small enough hit time were hit
+        return event.Channels(hit=t<1e8, t=t, 
+                        q=self.channel_q_gpu.get().astype(np.float32),
+                        histories=self.channel_history_gpu.get())
 
     def __del__(self):
         self.context.pop()
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()
diff --git a/tools.py b/tools.py
index 6835228..ba36f6d 100644
--- a/tools.py
+++ b/tools.py
@@ -1,6 +1,30 @@
 import numpy as np
 import time
 import datetime
+import sys
+
+def debugger_hook(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()
+
+def enable_debug_on_crash():
+    '''Start the PDB console when an uncaught exception propagates to the top.'''
+    sys.excepthook = debugger_hook
+
+# Allow profile decorator to exist, but do nothing if not running under kernprof
+try:
+    profile_if_possible = profile
+except NameError:
+    profile_if_possible = lambda x: x
 
 def timeit(func):
     def f(*args, **kwargs):