Expose interface to PDF creation on the GPU in the Simulation class.

1 files changed, 28 insertions, 0 deletions

diff --git a/sim.py b/sim.py
index 3909ec4..a056c50 100755
--- a/sim.py
+++ b/sim.py
@@ -55,6 +55,7 @@ class Simulation(object):
         print >>sys.stderr, 'Initializing random numbers generators...'
         self.gpu_worker.setup_propagate(seed=self.seed)
         self.gpu_worker.setup_daq(max(self.detector.pmtids))
+        self.pdf_config = None
 
     def simulate(self, nevents, vertex_generator, keep_photon_start=False, keep_photon_stop=False,
                  run_daq=True, nreps=1):
@@ -91,6 +92,33 @@ class Simulation(object):
         self.gpu_worker.propagate(max_steps=max_steps)
         return self.gpu_worker.get_photons()
 
+    def create_pdf(self, nevents, vertex_generator, tbins, trange,
+                   qbins, qrange, nreps=1):
+        photon_gen = repeating_iterator(self.photon_generator.generate_events(nevents, vertex_generator),
+                                        nreps)
+        return self.create_pdf_from_photons(nevents*nreps, photon_gen,
+                                            tbins, trange, qbins, qrange)
+
+    def create_pdf_from_photons(self, nevents, photon_generator,
+                                tbins, trange, qbins, qrange):
+        '''Returns tuple: 1D array of channel hit counts, 3D array of (channel, time, charge) pdfs'''
+        pdf_config = (tbins, trange, qbins, qrange)
+        if pdf_config != self.pdf_config:
+            self.pdf_config = pdf_config
+            self.gpu_worker.setup_pdf(max(self.detector.pmtids), tbins, trange,
+                                      qbins, qrange)
+        else:
+            self.gpu_worker.clear_pdf()
+
+        for ev in itertools.islice(photon_generator, nevents):
+            self.gpu_worker.load_photons(ev.photon_start)
+            self.gpu_worker.propagate()
+            self.gpu_worker.run_daq()
+            self.gpu_worker.add_hits_to_pdf()
+        
+        return self.gpu_worker.get_pdfs()
+        
+
 @profile_if_possible
 def main():
     import optparse