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

1 files changed, 24 insertions, 16 deletions

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() )