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

6 files changed, 115 insertions, 106 deletions

diff --git a/tests/test_fileio.py b/tests/test_fileio.py
index d21c9ff..2911976 100644
--- a/tests/test_fileio.py
+++ b/tests/test_fileio.py
@@ -5,36 +5,34 @@ import numpy as np
 
 class TestFileIO(unittest.TestCase):
     def test_file_write_and_read(self):
-        ev = event.Event(1, 'e-', (0,0,1), (1,0,0), 15)
+        ev = event.Event(1, event.Vertex('e-', (0,0,1), (1,0,0), (0,1,0), 15))
 
-        photon_start = root.make_photon_with_arrays(1)
-        photon_start.positions[0] = (1,2,3)
-        photon_start.directions[0] = (4,5,6)
-        photon_start.polarizations[0] = (7,8,9)
-        photon_start.wavelengths[0] = 400.0
-        photon_start.times[0] = 100.0
-        photon_start.histories[0] = 20
-        photon_start.last_hit_triangles[0] = 5
-        ev.photon_start = photon_start
+        photons_beg = root.make_photon_with_arrays(1)
+        photons_beg.pos[0] = (1,2,3)
+        photons_beg.dir[0] = (4,5,6)
+        photons_beg.pol[0] = (7,8,9)
+        photons_beg.wavelengths[0] = 400.0
+        photons_beg.t[0] = 100.0
+        photons_beg.last_hit_triangles[0] = 5
+        photons_beg.flags[0] = 20
+        ev.photons_beg = photons_beg
 
-        photon_stop = root.make_photon_with_arrays(1)
-        photon_stop.positions[0] = (1,2,3)
-        photon_stop.directions[0] = (4,5,6)
-        photon_stop.polarizations[0] = (7,8,9)
-        photon_stop.wavelengths[0] = 400.0
-        photon_stop.times[0] = 100.0
-        photon_stop.histories[0] = 20
-        photon_stop.last_hit_triangles[0] = 5
-        ev.photon_stop = photon_stop
+        photons_end = root.make_photon_with_arrays(1)
+        photons_end.pos[0] = (1,2,3)
+        photons_end.dir[0] = (4,5,6)
+        photons_end.pol[0] = (7,8,9)
+        photons_end.wavelengths[0] = 400.0
+        photons_end.t[0] = 100.0
+        photons_end.last_hit_triangles[0] = 5
+        photons_end.flags[0] = 20
+        ev.photons_end = photons_end
 
-        ev.nphoton = 1
-
-        ev.subtracks.append(event.Subtrack('e-', (40,30,20), (-1, -2, -3), 400, 800))
+        ev.vertices = [ev.primary_vertex]
 
         channels = event.Channels(hit=np.array([True, False]), 
                                   t=np.array([20.0, 1e9], dtype=np.float32),
                                   q=np.array([2.0, 0.0], dtype=np.float32),
-                                  histories=np.array([8, 32], dtype=np.uint32))
+                                  flags=np.array([8, 32], dtype=np.uint32))
         ev.channels = channels
 
         filename = '/tmp/chroma-filewritertest.root'
@@ -58,16 +56,18 @@ class TestFileIO(unittest.TestCase):
         # Enough screwing around, let's get the one event in the file
         newev = reader.current()
 
-
         # Now check if everything is correct in the event
-        for attribute in ['event_id', 'particle_name','gen_total_energy']:
+        for attribute in ['id']:
             self.assertEqual(getattr(ev, attribute), getattr(newev, attribute), 'compare %s' % attribute)
-        for attribute in ['gen_position', 'gen_direction']:
-            self.assertTrue(np.allclose(getattr(ev, attribute), getattr(newev, attribute)), 'compare %s' % attribute)
 
-        for attribute in ['positions', 'directions', 'wavelengths', 'polarizations', 'times',
-                          'histories', 'last_hit_triangles']:    
-            self.assertTrue(np.allclose(getattr(ev.photon_start, attribute), 
-                                        getattr(newev.photon_start, attribute)), 'compare %s' % attribute)
-            self.assertTrue(np.allclose(getattr(ev.photon_stop, attribute), 
-                                        getattr(newev.photon_stop, attribute)), 'compare %s' % attribute)
+        for attribute in ['pos', 'dir', 'pol', 'ke', 't0']:
+            self.assertTrue(np.allclose(getattr(ev.primary_vertex, attribute), getattr(newev.primary_vertex, attribute)), 'compare %s' % attribute)
+
+            for i in range(len(ev.vertices)):
+                self.assertTrue(np.allclose(getattr(ev.vertices[i], attribute), getattr(newev.vertices[i], attribute)), 'compare %s' % attribute)
+
+        for attribute in ['pos', 'dir', 'pol', 'wavelengths', 't', 'last_hit_triangles', 'flags']:    
+            self.assertTrue(np.allclose(getattr(ev.photons_beg, attribute), 
+                                        getattr(newev.photons_beg, attribute)), 'compare %s' % attribute)
+            self.assertTrue(np.allclose(getattr(ev.photons_end, attribute), 
+                                        getattr(newev.photons_end, attribute)), 'compare %s' % attribute)
diff --git a/tests/test_generator_photon.py b/tests/test_generator_photon.py
index 9684126..13501fe 100644
--- a/tests/test_generator_photon.py
+++ b/tests/test_generator_photon.py
@@ -1,20 +1,21 @@
 import unittest
+import itertools
 
-import chroma.generator.photon
+from chroma import generator
 from chroma.generator.vertex import constant_particle_gun
-from chroma.optics import water_wcsim
+from chroma import optics
 
 class TestG4ParallelGenerator(unittest.TestCase):
     def test_center(self):
         '''Generate Cherenkov light at the center of the world volume'''
-        gen = chroma.generator.photon.G4ParallelGenerator(1, water_wcsim)
-        vertex = constant_particle_gun('e-', (0,0,0), (1,0,0), 100)
-        for event in gen.generate_events(10, vertex):
-            self.assertGreater(len(event.photon_start.positions), 0)
+        gen = generator.photon.G4ParallelGenerator(1, optics.water_wcsim)
+        vertex = itertools.islice(constant_particle_gun('e-', (0,0,0), (1,0,0), 100), 10)
+        for event in gen.generate_events(vertex):
+            self.assertGreater(len(event.photons_beg.pos), 0)
 
     def test_off_center(self):
         '''Generate Cherenkov light at (1 m, 0 m, 0 m)'''
-        gen = chroma.generator.photon.G4ParallelGenerator(1, water_wcsim)
-        vertex = constant_particle_gun('e-', (1,0,0), (1,0,0), 100)
-        for event in gen.generate_events(10, vertex):
-            self.assertGreater(len(event.photon_start.positions), 0)
+        gen = generator.photon.G4ParallelGenerator(1, optics.water_wcsim)
+        vertex = itertools.islice(constant_particle_gun('e-', (1,0,0), (1,0,0), 100), 10)
+        for event in gen.generate_events(vertex):
+            self.assertGreater(len(event.photons_beg.pos), 0)
diff --git a/tests/test_generator_vertex.py b/tests/test_generator_vertex.py
index d773579..cec363d 100644
--- a/tests/test_generator_vertex.py
+++ b/tests/test_generator_vertex.py
@@ -8,16 +8,16 @@ class TestParticleGun(unittest.TestCase):
         '''Generate electron vertices at the center of the world volume.'''
         vertex = chroma.generator.vertex.constant_particle_gun('e-', (0,0,0), (1,0,0), 100)
         for ev in itertools.islice(vertex, 100):
-            self.assertEquals(ev.particle_name, 'e-')
-            self.assertTrue(np.allclose(ev.gen_position, [0,0,0]))
-            self.assertTrue(np.allclose(ev.gen_direction, [1,0,0]))
-            self.assertTrue(np.allclose(ev.gen_total_energy, 100))
+            self.assertEquals(ev.primary_vertex.particle_name, 'e-')
+            self.assertTrue(np.allclose(ev.primary_vertex.pos, [0,0,0]))
+            self.assertTrue(np.allclose(ev.primary_vertex.dir, [1,0,0]))
+            self.assertTrue(np.allclose(ev.primary_vertex.ke, 100))
 
     def test_off_center(self):
         '''Generate electron vertices at (1,0,0) in the world volume.'''
         vertex = chroma.generator.vertex.constant_particle_gun('e-', (1,0,0), (1,0,0), 100)
         for ev in itertools.islice(vertex, 100):
-            self.assertEquals(ev.particle_name, 'e-')
-            self.assertTrue(np.allclose(ev.gen_position, [1,0,0]))
-            self.assertTrue(np.allclose(ev.gen_direction, [1,0,0]))
-            self.assertTrue(np.allclose(ev.gen_total_energy, 100))
+            self.assertEquals(ev.primary_vertex.particle_name, 'e-')
+            self.assertTrue(np.allclose(ev.primary_vertex.pos, [1,0,0]))
+            self.assertTrue(np.allclose(ev.primary_vertex.dir, [1,0,0]))
+            self.assertTrue(np.allclose(ev.primary_vertex.ke, 100))
diff --git a/tests/test_pdf.py b/tests/test_pdf.py
index 9a08b53..791f119 100644
--- a/tests/test_pdf.py
+++ b/tests/test_pdf.py
@@ -1,11 +1,12 @@
 import unittest
 import numpy as np
+import itertools
 
 import chroma.detectors
 from chroma.generator.photon import G4ParallelGenerator
 from chroma.generator.vertex import constant_particle_gun
 from chroma.optics import water_wcsim
-from chroma.gpu import GPU
+from chroma import gpu
 from chroma.sim import Simulation
 
 class TestPDF(unittest.TestCase):
@@ -18,21 +19,26 @@ class TestPDF(unittest.TestCase):
         '''Create a hit count and (q,t) PDF for 10 MeV events in MicroLBNE'''
 
         g4generator = G4ParallelGenerator(1, water_wcsim)
-        gpu = GPU(0)
-        gpu.load_geometry(self.detector)
-        gpu.setup_propagate()
-        gpu.setup_daq(max(self.detector.pmtids))
-        gpu.setup_pdf(max(self.detector.pmtids), 100, (-0.5, 999.5), 10, (-0.5, 9.5))
+        gpu_instance = gpu.GPU(0)
+        gpu_geometry = gpu.GPUGeometry(gpu_instance, self.detector)
+
+        nthreads_per_block, max_blocks = 64, 1024
+
+        rng_states = gpu.get_rng_states(nthreads_per_block*max_blocks)
+
+        gpu_daq = gpu.GPUDaq(gpu_geometry, max(self.detector.pmtids))
+        gpu_pdf = gpu.GPUPDF()
+        gpu_pdf.setup_pdf(max(self.detector.pmtids), 100, (-0.5, 999.5), 10, (-0.5, 9.5))
         
-        gpu.clear_pdf()
+        gpu_pdf.clear_pdf()
 
-        for ev in g4generator.generate_events(10, self.vertex_gen):
-            gpu.load_photons(ev.photon_start)
-            gpu.propagate()
-            gpu.run_daq()
-            gpu.add_hits_to_pdf()
+        for ev in g4generator.generate_events(itertools.islice(self.vertex_gen, 10)):
+            gpu_photons = gpu.GPUPhotons(ev.photons_beg)
+            gpu.propagate(gpu_instance, gpu_photons, rng_states, nthreads_per_block, max_blocks)
+            gpu_channels = gpu_daq.acquire(gpu_photons, rng_states, nthreads_per_block, max_blocks)
+            gpu_pdf.add_hits_to_pdf(gpu_channels)
 
-        hitcount, pdf = gpu.get_pdfs()
+        hitcount, pdf = gpu_pdf.get_pdfs()
         self.assertTrue( (hitcount > 0).any() )
         self.assertTrue( (pdf > 0).any() )
 
@@ -41,9 +47,11 @@ class TestPDF(unittest.TestCase):
             self.assertEqual(nhits, pdf[i].sum())
 
     def testSimPDF(self):
-        sim = Simulation(self.detector, water_wcsim)
-        hitcount, pdf = sim.create_pdf(100, self.vertex_gen,
-                                        100, (-0.5, 999.5), 10, (-0.5, 9.5))
+        sim = Simulation(self.detector)
+
+        vertex_iter = itertools.islice(self.vertex_gen, 10)
+
+        hitcount, pdf = sim.create_pdf(vertex_iter, 100, (-0.5, 999.5), 10, (-0.5, 9.5))
 
         self.assertTrue( (hitcount > 0).any() )
         self.assertTrue( (pdf > 0).any() )
diff --git a/tests/test_propagation.py b/tests/test_propagation.py
index 331242b..bf886bd 100644
--- a/tests/test_propagation.py
+++ b/tests/test_propagation.py
@@ -11,43 +11,43 @@ class TestPropagation(unittest.TestCase):
     def testAbort(self):
         '''Photons that hit a triangle at normal incidence should not abort.
 
-        Photons that hit a triangle at exactly normal incidence can sometimes produce a dot product
-        that is outside the range allowed by acos().  Trigger these with axis aligned photons in a box
+        Photons that hit a triangle at exactly normal incidence can sometimes
+        produce a dot product that is outside the range allowed by acos().
+        Trigger these with axis aligned photons in a box.
         '''
 
         # Setup geometry
-        cube = Geometry()
+        cube = Geometry(vacuum)
         cube.add_solid(Solid(box(100,100,100), vacuum, vacuum))
         cube.pmtids = [0]
-        sim = Simulation(cube, vacuum, bvh_bits=4, geant4_processes=0,
-                              use_cache=False)
+        sim = Simulation(cube, bvh_bits=4, geant4_processes=0, use_cache=False)
 
         # Create initial photons
         nphotons = 10000
-        positions = np.tile([0,0,0], (nphotons,1)).astype(np.float32)
-        directions = np.tile([0,0,1], (nphotons,1)).astype(np.float32)
-        polarizations = np.zeros_like(positions)
+        pos = np.tile([0,0,0], (nphotons,1)).astype(np.float32)
+        dir = np.tile([0,0,1], (nphotons,1)).astype(np.float32)
+        pol = np.zeros_like(pos)
         phi = np.random.uniform(0, 2*np.pi, nphotons).astype(np.float32)
-        polarizations[:,0] = np.cos(phi)
-        polarizations[:,1] = np.sin(phi)
-        times = np.zeros(nphotons, dtype=np.float32)
+        pol[:,0] = np.cos(phi)
+        pol[:,1] = np.sin(phi)
+        t = np.zeros(nphotons, dtype=np.float32)
         wavelengths = np.empty(nphotons, np.float32)
         wavelengths.fill(400.0)
 
-        photons = Photons(positions=positions, directions=directions, polarizations=polarizations,
-                          times=times, wavelengths=wavelengths)
+        photons = Photons(pos=pos, dir=dir, pol=pol, t=t, wavelengths=wavelengths)
 
         # First make one step to check for strangeness
-        photon_stop = sim.propagate_photons(photons, max_steps=1)
-        self.assertFalse(np.isnan(photon_stop.positions).any())
-        self.assertFalse(np.isnan(photon_stop.directions).any())
-        self.assertFalse(np.isnan(photon_stop.polarizations).any())
-        self.assertFalse(np.isnan(photon_stop.times).any())
-        self.assertFalse(np.isnan(photon_stop.wavelengths).any())
+        photons_end = sim.simulate([photons], keep_photons_end=True, max_steps=1).next().photons_end
+
+        self.assertFalse(np.isnan(photons_end.pos).any())
+        self.assertFalse(np.isnan(photons_end.dir).any())
+        self.assertFalse(np.isnan(photons_end.pol).any())
+        self.assertFalse(np.isnan(photons_end.t).any())
+        self.assertFalse(np.isnan(photons_end.wavelengths).any())
 
         # Now let it run the usual ten steps
-        photon_stop = sim.propagate_photons(photons, max_steps=10)
-        aborted = (photon_stop.histories & (1 << 31)) > 0
+        photons_end = sim.simulate([photons], keep_photons_end=True, max_steps=10).next().photons_end
+        aborted = (photons_end.flags & (1 << 31)) > 0
         print 'aborted photons: %1.1f' % (float(np.count_nonzero(aborted)) / nphotons)
         self.assertFalse(aborted.any())
         
diff --git a/tests/test_rayleigh.py b/tests/test_rayleigh.py
index 4394ada..7e5fc92 100644
--- a/tests/test_rayleigh.py
+++ b/tests/test_rayleigh.py
@@ -13,43 +13,43 @@ ROOT.gROOT.SetBatch(1)
 
 class TestRayleigh(unittest.TestCase):
     def setUp(self):
-        self.cube = Geometry()
+        self.cube = Geometry(water_wcsim)
         self.cube.add_solid(Solid(box(100,100,100), water_wcsim, water_wcsim))
         self.cube.pmtids = [0]
-        self.sim = Simulation(self.cube, water_wcsim, bvh_bits=4, geant4_processes=0,
+        self.sim = Simulation(self.cube, bvh_bits=4, geant4_processes=0,
                               use_cache=False)
         nphotons = 100000
-        positions = np.tile([0,0,0], (nphotons,1)).astype(np.float32)
-        directions = np.tile([0,0,1], (nphotons,1)).astype(np.float32)
-        polarizations = np.zeros_like(positions)
+        pos = np.tile([0,0,0], (nphotons,1)).astype(np.float32)
+        dir = np.tile([0,0,1], (nphotons,1)).astype(np.float32)
+        pol = np.zeros_like(pos)
         phi = np.random.uniform(0, 2*np.pi, nphotons).astype(np.float32)
-        polarizations[:,0] = np.cos(phi)
-        polarizations[:,1] = np.sin(phi)
-        times = np.zeros(nphotons, dtype=np.float32)
+        pol[:,0] = np.cos(phi)
+        pol[:,1] = np.sin(phi)
+        t = np.zeros(nphotons, dtype=np.float32)
         wavelengths = np.empty(nphotons, np.float32)
         wavelengths.fill(400.0)
 
-        self.photons = Photons(positions=positions, directions=directions, polarizations=polarizations,
-                               times=times, wavelengths=wavelengths)
+        self.photons = Photons(pos=pos, dir=dir, pol=pol, t=t, wavelengths=wavelengths)
 
     def testAngularDistributionPolarized(self):
         # Fully polarized photons
-        self.photons.polarizations[:] = [1.0, 0.0, 0.0]
+        self.photons.pol[:] = [1.0, 0.0, 0.0]
 
-        photon_stop = self.sim.propagate_photons(self.photons, max_steps=1)
-        aborted = (photon_stop.histories & (1 << 31)) > 0
+        photons_end = self.sim.simulate([self.photons], keep_photons_end=True, max_steps=1).next().photons_end
+        aborted = (photons_end.flags & (1 << 31)) > 0
         self.assertFalse(aborted.any())
 
-        # Compute the dot product between initial and final directions
-        rayleigh_scatters = (photon_stop.histories & (1 << 4)) > 0
-        cos_scatter = (self.photons.directions[rayleigh_scatters] * photon_stop.directions[rayleigh_scatters]).sum(axis=1)
+        # Compute the dot product between initial and final dir
+        rayleigh_scatters = (photons_end.flags & (1 << 4)) > 0
+        cos_scatter = (self.photons.dir[rayleigh_scatters] * photons_end.dir[rayleigh_scatters]).sum(axis=1)
         theta_scatter = np.arccos(cos_scatter)
         h = histogram.Histogram(bins=100, range=(0, np.pi))
         h.fill(theta_scatter)
         h = rootify(h)
 
-        # The functional form for polarized light should be (1 + \cos^2 \theta)\sin \theta
-        # according to GEANT4 physics reference manual
+        # The functional form for polarized light should be
+        # (1 + \cos^2 \theta)\sin \theta according to GEANT4 physics
+        # reference manual.
         f = ROOT.TF1("pol_func", "[0]*(1+cos(x)**2)*sin(x)", 0, np.pi)
         h.Fit(f)
         self.assertGreater(f.GetProb(), 1e-3)