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

1 files changed, 11 insertions, 10 deletions

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)