update event structure. break gpu.GPU class into separate smaller independent classes.

3 files changed, 86 insertions, 113 deletions

diff --git a/generator/g4gen.py b/generator/g4gen.py
index 9650b68..8dca086 100644
--- a/generator/g4gen.py
+++ b/generator/g4gen.py
@@ -4,7 +4,7 @@ import g4py.NISTmaterials
 import g4py.ParticleGun
 import pyublas
 import numpy as np
-import chroma.event as event
+from chroma.event import Photons, Vertex
 
 try:
     import G4chroma
@@ -20,16 +20,18 @@ except:
 
 class G4Generator(object):
     def __init__(self, material, seed=None):
-        '''Create generator to produce photons inside the specified material.
+        """Create generator to produce photons inside the specified material.
 
            material: chroma.geometry.Material object with density, 
                      composition dict and refractive_index.
 
                      composition dictionary should be 
                         { element_symbol : fraction_by_weight, ... }.
-           seed: Random number generator seed for HepRandom. If None,
-                 generator is not seeded.
-        '''
+
+           seed: int, *optional*
+               Random number generator seed for HepRandom. If None, generator
+               is not seeded.
+        """
         if seed is not None:
             HepRandom.setTheSeed(seed)
 
@@ -51,11 +53,8 @@ class G4Generator(object):
         gRunManager.SetUserAction(self.tracking_action)
         gRunManager.Initialize()
         
-        #preinitialize the process by running a simple event
-        self.generate_photons(event.Event(event_id=0, particle_name='e-', 
-                                          gen_position=(0,0,0), 
-                                          gen_direction=(1,0,0), 
-                                          gen_total_energy=1.0))
+        # preinitialize the process by running a simple event
+        self.generate_photons([Vertex('e-', (0,0,0), (1,0,0), 0, 1.0)])
         
     def create_g4material(self, material):
         g4material = G4Material('world_material', material.density * g / cm3,
@@ -96,57 +95,38 @@ class G4Generator(object):
         pol[:,2] = self.tracking_action.GetPolZ()
         
         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)
+        t0 = self.tracking_action.GetT0().astype(np.float32)
+
+        return Photons(pos, dir, pol, wavelengths, t0)
 
-    def generate_photons(self, ev):
-        '''Use GEANT4 to generate photons produced by the given particle.
+    def generate_photons(self, vertices):
+        """Use GEANT4 to generate photons produced by propagating `vertices`.
            
-               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_position,
-                                        direction=ev.gen_direction,
-                                        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())
+        Args:
+            vertices: list of event.Vertex objects
+                List of initial particle vertices.
+
+        Returns:
+            photons: event.Photons
+                Photon vertices generated by the propagation of `vertices`.
+        """
+        photons = None
+
+        for vertex in vertices:
+            self.particle_gun.SetParticleByName(vertex.particle_name)
+            mass = G4ParticleTable.GetParticleTable().FindParticle(vertex.particle_name).GetPDGMass()
+            total_energy = vertex.ke*MeV + mass
+            self.particle_gun.SetParticleEnergy(total_energy)
+            self.particle_gun.SetParticlePosition(G4ThreeVector(*vertex.pos)*m)
+            self.particle_gun.SetParticleMomentumDirection(G4ThreeVector(*vertex.dir).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)
-
-    start = time.time()
-    n = 0
-    for i in xrange(100):
-        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'
+
+            if photons is None:
+                photons = self._extract_photons_from_tracking_action()
+            else:
+                photons += self._extract_photons_from_tracking_action()
+
+        return photons
diff --git a/generator/photon.py b/generator/photon.py
index e656428..6c656b2 100644
--- a/generator/photon.py
+++ b/generator/photon.py
@@ -25,24 +25,29 @@ class G4GeneratorProcess(multiprocessing.Process):
         photon_socket = context.socket(zmq.PUSH)
         photon_socket.connect(self.photon_socket_address)
 
-        # Signal with the photon socket that we are online and ready for messages
+        # Signal with the photon socket that we are online
+        # and ready for messages.
         photon_socket.send('READY')
 
         while True:
             ev = vertex_socket.recv_pyobj()
-            ev.photon_start = gen.generate_photons(ev)
+            ev.photons_beg = gen.generate_photons(ev.vertices)
+            #print 'type(ev.photons_beg) is %s' % type(ev.photons_beg)
             photon_socket.send_pyobj(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
-    '''
+    """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.
+
+    Examples:
+        >>> 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:
@@ -51,8 +56,8 @@ def partition(num, partitions):
             yield step + (num % partitions)
 
 def vertex_sender(vertex_iterator, vertex_socket):
-    for ev in vertex_iterator:
-        vertex_socket.send_pyobj(ev)
+    for vertex in vertex_iterator:
+        vertex_socket.send_pyobj(vertex)
 
 def socket_iterator(nelements, socket):
     for i in xrange(nelements):
@@ -66,8 +71,7 @@ class G4ParallelGenerator(object):
         base_address = 'ipc:///tmp/chroma_'+str(uuid.uuid4())
         self.vertex_address = base_address + '.vertex'
         self.photon_address = base_address + '.photon'
-        self.processes = [ G4GeneratorProcess(i, material, self.vertex_address, self.photon_address, seed=base_seed + i)
-                           for i in xrange(nprocesses) ]
+        self.processes = [ G4GeneratorProcess(i, material, self.vertex_address, self.photon_address, seed=base_seed + i) for i in xrange(nprocesses) ]
 
         for p in self.processes:
             p.start()
@@ -78,18 +82,16 @@ class G4ParallelGenerator(object):
         self.photon_socket = self.zmq_context.socket(zmq.PULL)
         self.photon_socket.bind(self.photon_address)
 
-        # Verify everyone is running and connected to avoid sending all the events to one client
+        # Verify everyone is running and connected to avoid
+        # sending all the events to one client.
         for i in xrange(nprocesses):
             msg = self.photon_socket.recv()
             assert msg == 'READY'
 
-    def generate_events(self, nevents, vertex_iterator):
-        # Doing this to avoid a deadlock caused by putting to one queue while getting from another
-        limited_iterator = itertools.islice(vertex_iterator, nevents)
-        sender_thread = threading.Thread(target=vertex_sender, args=(limited_iterator,
-                                                                     self.vertex_socket))
+    def generate_events(self, vertex_iterator):
+        # Doing this to avoid a deadlock caused by putting to one queue
+        # while getting from another.
+        vertex_list = list(vertex_iterator)
+        sender_thread = threading.Thread(target=vertex_sender, args=(vertex_list, self.vertex_socket))
         sender_thread.start()
-        return socket_iterator(nevents, self.photon_socket)
-
-        
-            
+        return socket_iterator(len(vertex_list), self.photon_socket)
diff --git a/generator/vertex.py b/generator/vertex.py
index 450c0e4..5521d6b 100644
--- a/generator/vertex.py
+++ b/generator/vertex.py
@@ -2,7 +2,7 @@ import numpy as np
 from itertools import izip, count
 
 from chroma.pi0 import pi0_decay
-from chroma.event import Event, Subtrack
+from chroma import event
 from chroma.sample import uniform_sphere
 from chroma.itertoolset import repeat_func
 
@@ -40,43 +40,34 @@ def flat(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 \
-            izip(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 particle_gun(particle_name_iter, pos_iter, dir_iter, ke_iter, start_id=0):
+    for i, particle_name, pos, dir, ke in izip(count(start_id), particle_name_iter, pos_iter, dir_iter, ke_iter):
+        dir /= np.linalg.norm(dir)
+        vertex = event.Vertex(particle_name, pos, dir, None, ke)
+        ev_vertex = event.Event(i, vertex, [vertex])
+        yield ev_vertex
 
-def pi0_gun(pi0_position, pi0_direction, pi0_total_energy, start_id=0):
-    for i, ev_position, ev_direction, ev_total_energy in \
-            izip(count(start_id), pi0_position, pi0_direction, pi0_total_energy):
+def pi0_gun(pos_iter, dir_iter, ke_iter, start_id=0):
+    for i, pos, dir, ke in izip(count(start_id), pos_iter, dir_iter, ke_iter):
+        dir /= np.linalg.norm(dir)
+        primary_vertex = event.Vertex('pi0', pos, dir, None, ke)
 
-        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_decay(energy=ev_total_energy,
-                                                                       direction=ev_direction,
-                                                                       theta=theta_rest,
-                                                                       phi=phi_rest)
+        (gamma1_e, gamma1_dir), (gamma2_e, gamma2_dir) = \
+            pi0_decay(ke+134.9766, dir, theta_rest, phi_rest)
+
+        gamma1_vertex = event.Vertex('gamma', pos, gamma1_dir, None, gamma1_e)
+        gamma2_vertex = event.Vertex('gamma', pos, gamma2_dir, None, gamma2_e)
+
+        ev_vertex = event.Event(i, primary_vertex, [gamma1_vertex, gamma2_vertex])
 
-        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
+        yield ev_vertex
 
 
-def constant_particle_gun(particle_name, position, direction, total_energy,
-                          start_id=0):
+def constant_particle_gun(particle_name, pos, dir, ke, start_id=0):
     '''Convenience wrapper around particle gun that assumes all
     arguments are constants, rather than generators.'''
-    return particle_gun(constant(particle_name), constant(position),
-                        constant(direction), constant(total_energy),
-                        start_id=start_id)
+    return particle_gun(constant(particle_name), constant(pos), constant(dir), constant(ke), start_id=start_id)