Enable creation of charge and time PDF on GPU

2 files changed, 75 insertions, 1 deletions

diff --git a/gpu.py b/gpu.py
index ff8db8c..4d2197b 100644
--- a/gpu.py
+++ b/gpu.py
@@ -115,7 +115,7 @@ class GPU(object):
         self.daq_module = SourceModule(chroma.src.daq, options=cuda_options, no_extern_c=True)
         self.daq_funcs = CUDAFuncs(self.daq_module, 
                               ['reset_earliest_time_int', 'run_daq',
-                               'convert_sortable_int_to_float'])
+                               'convert_sortable_int_to_float', 'bin_hits'])
 
     def print_device_usage(self):
         print 'device usage:'
@@ -375,5 +375,52 @@ class GPU(object):
                         q=self.channel_q_gpu.get().astype(np.float32),
                         histories=self.channel_history_gpu.get())
 
+    def setup_pdf(self, max_pmt_id, tbins, trange, qbins, qrange):
+        '''Setup GPU arrays to hold PDF information.
+
+           max_pmt_id: int, largest PMT id #
+           tbins: number of time bins
+           trange: tuple of (min, max) time in PDF
+           qbins: number of charge bins
+           qrange: tuple of (min, max) charge in PDF
+        '''
+        self.events_in_histogram = 0
+        self.hitcount_gpu = gpuarray.zeros(max_pmt_id+1, dtype=np.uint32)
+        self.pdf_gpu = gpuarray.zeros(shape=(max_pmt_id+1, tbins, qbins), 
+                                      dtype=np.uint32)
+        self.tbins = tbins
+        self.trange = trange
+        self.qbins = qbins
+        self.qrange = qrange
+
+    def clear_pdf(self):
+        '''Rezero the PDF counters.'''
+        self.hitcount_gpu.fill(0)
+        self.pdf_gpu.fill(0)
+
+    def add_hits_to_pdf(self):
+        '''Put the most recent results of run_daq() into the PDF.'''
+
+        self.daq_funcs.bin_hits(np.int32(len(self.hitcount_gpu)),
+                                self.channel_q_gpu,
+                                self.earliest_time_gpu,
+                                self.hitcount_gpu,
+                                np.int32(self.tbins),
+                                np.float32(self.trange[0]),
+                                np.float32(self.trange[1]),
+                                np.int32(self.qbins),
+                                np.float32(self.qrange[0]),
+                                np.float32(self.qrange[1]),
+                                self.pdf_gpu,
+                                block=(self.nthreads_per_block,1,1), 
+                                grid=(len(self.earliest_time_gpu)//self.nthreads_per_block+1,1))
+
+
+        self.events_in_histogram += 1
+
+    def get_pdfs(self):
+        '''Returns the 1D hitcount array and the 3D [channel, time, charge] histogram'''
+        return self.hitcount_gpu.get(), self.pdf_gpu.get()
+
     def __del__(self):
         self.context.pop()
diff --git a/src/daq.cu b/src/daq.cu
index 2b95560..a34fbbe 100644
--- a/src/daq.cu
+++ b/src/daq.cu
@@ -80,4 +80,31 @@ __global__ void convert_sortable_int_to_float(int n,
 		float_output[id] = sortable_int_to_float(sortable_ints[id]);
 }
 
+__global__ void bin_hits(int nchannels,
+			 unsigned int *channel_q, float *channel_time,
+			 unsigned int *hitcount, 
+			 int tbins, float tmin, float tmax,
+			 int qbins, float qmin, float qmax,
+			 unsigned int *pdf)
+{
+	int id = threadIdx.x + blockDim.x * blockIdx.x;
+
+	if (id >= nchannels)
+	  return;
+
+	unsigned int q = channel_q[id];
+	float t = channel_time[id];
+	
+	if (t < 1e8 && t >= tmin && t < tmax && q >= qmin && q < qmax) {
+		hitcount[id] += 1;
+		
+		int tbin = (t - tmin) / (tmax - tmin) * tbins;
+		int qbin = (q - qmin) / (qmax - qmin) * qbins;
+	    
+		// row major order (channel, t, q)
+		int bin = id * (tbins * qbins) + tbin * qbins + qbin;
+		pdf[bin] += 1;
+	}
+}
+
 } // extern "C"