Benchmark and unit test for PDF creation

2 files changed, 113 insertions, 4 deletions

diff --git a/benchmark.py b/benchmark.py
index 7773437..7988399 100755
--- a/benchmark.py
+++ b/benchmark.py
@@ -9,6 +9,12 @@ from chroma.gpu import GPU, to_float3
 from chroma.camera import get_rays
 from chroma.event import Photons
 from chroma.sample import uniform_sphere
+from chroma.generator.photon import G4ParallelGenerator
+from chroma.optics import water_wcsim
+from chroma.generator.vertex import constant_particle_gun
+from chroma.tools import profile_if_possible
+# Generator processes need to fork BEFORE the GPU context is setup
+g4generator = G4ParallelGenerator(4, water_wcsim)
 
 def progress(seq):
     "Print progress while iterating over `seq`."
@@ -122,6 +128,55 @@ def propagate(gpu, number=10, nphotons=500000):
 
     return nphotons/ufloat((np.mean(run_times),np.std(run_times)))
 
+@profile_if_possible
+def pdf(gpu, max_pmt_id, npdfs=10, nevents=100, nreps=1):
+    """
+    Return the mean and standard deviation of the number of 100 MeV
+    events per second that can be histogrammed a ufloat for the
+    geometry loaded onto `gpu`.
+
+    Args:
+        - gpu, chroma.gpu.GPU
+            The GPU object with a geometry already loaded.
+        - max_pmt_id, int
+            The channel number of the highest PMT
+        - npdfs, int
+            The number of pdf generations to average.
+        - nevents, int
+            The number of 100 MeV events to generate for each PDF.
+        - nreps, int
+            The number of times to propagate each event and add to PDF
+    """
+    gpu.setup_propagate()
+
+    run_times = []
+
+    gpu.setup_propagate()
+    gpu.setup_daq(max_pmt_id)
+    gpu.setup_pdf(max_pmt_id, 100, (-0.5, 999.5), 10, (-0.5, 9.5))
+    vertex_gen = constant_particle_gun('e-', (0,0,0), (1,0,0), 100)
+
+    for i in progress(range(npdfs)):
+        t0 = time.time()
+        gpu.clear_pdf()
+
+        for ev in g4generator.generate_events(nevents, vertex_gen):
+            for j in xrange(nreps):
+                gpu.load_photons(ev.photon_start)
+                gpu.propagate()
+                gpu.run_daq()
+                gpu.add_hits_to_pdf()
+
+        hitcount, pdf = gpu.get_pdfs()
+
+        elapsed = time.time() - t0
+
+        if i > 0:
+            # first kernel call incurs some driver overhead
+            run_times.append(elapsed)
+
+    return nevents*nreps/ufloat((np.mean(run_times),np.std(run_times)))
+
 if __name__ == '__main__':
     from chroma.detectors import build_lbne_200kton, build_minilbne
 
@@ -130,7 +185,8 @@ if __name__ == '__main__':
 
     gpu = GPU()
     gpu.load_geometry(lbne, print_usage=False)
-
-    print '%s track steps/s' % ray_trace(gpu)
-    print '%s loaded photons/s' % load_photons(gpu)
-    print '%s propagated photons/s' % propagate(gpu)
+    
+    #print '%s track steps/s' % ray_trace(gpu)
+    #print '%s loaded photons/s' % load_photons(gpu)
+    #print '%s propagated photons/s' % propagate(gpu)
+    print '%s histogrammed 100 MeV events/s' % pdf(gpu, max(lbne.pmtids))
diff --git a/tests/test_pdf.py b/tests/test_pdf.py
new file mode 100644
index 0000000..9a08b53
--- /dev/null
+++ b/tests/test_pdf.py
@@ -0,0 +1,53 @@
+import unittest
+import numpy as np
+
+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.sim import Simulation
+
+class TestPDF(unittest.TestCase):
+    def setUp(self):
+        self.detector = chroma.detectors.find("microlbne")
+        self.detector.build()
+        self.vertex_gen = constant_particle_gun('e-', (0,0,0), (1,0,0), 10)
+
+    def testGPUPDF(self):
+        '''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.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()
+
+        hitcount, pdf = gpu.get_pdfs()
+        self.assertTrue( (hitcount > 0).any() )
+        self.assertTrue( (pdf > 0).any() )
+
+        # Consistency checks
+        for i, nhits in enumerate(hitcount):
+            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))
+
+        self.assertTrue( (hitcount > 0).any() )
+        self.assertTrue( (pdf > 0).any() )
+
+        # Consistency checks
+        for i, nhits in enumerate(hitcount):
+            self.assertEqual(nhits, pdf[i].sum())