1 files changed, 67 insertions, 0 deletions
diff --git a/test/test_pdf.py b/test/test_pdf.py
new file mode 100644
index 0000000..571cbd4
--- /dev/null
+++ b/test/test_pdf.py
@@ -0,0 +1,67 @@
+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 import gpu
+from chroma.sim import Simulation
+
+class TestPDF(unittest.TestCase):
+    def setUp(self):
+        self.detector = chroma.detectors.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)
+
+        context = gpu.create_cuda_context()
+
+        gpu_geometry = gpu.GPUGeometry(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_pdf.clear_pdf()
+
+        for ev in g4generator.generate_events(itertools.islice(self.vertex_gen, 10)):
+            gpu_photons = gpu.GPUPhotons(ev.photons_beg)
+            gpu_photons.propagate(gpu_geometry, 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_pdf.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())
+
+        context.pop()
+
+    def testSimPDF(self):
+        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() )
+
+        # Consistency checks
+        for i, nhits in enumerate(hitcount):
+            self.assertEqual(nhits, pdf[i].sum())