1 files changed, 67 insertions, 0 deletions
diff --git a/test/test_sample_cdf.py b/test/test_sample_cdf.py
new file mode 100644
index 0000000..1510768
--- /dev/null
+++ b/test/test_sample_cdf.py
@@ -0,0 +1,67 @@
+import pycuda.driver as cuda
+from pycuda.compiler import SourceModule
+from pycuda import gpuarray
+import numpy as np
+import ROOT
+import os
+import unittest
+import chroma
+
+class TestSampling(unittest.TestCase):
+    def setUp(self):
+        self.context = chroma.gpu.create_cuda_context()
+        current_directory = os.path.split(os.path.realpath(__file__))[0]
+        from chroma.cuda import srcdir as source_directory
+        source = open(current_directory + '/test_sample_cdf.cu').read()
+        self.mod = SourceModule(source, options=['-I' + source_directory], no_extern_c=True, cache_dir=False)
+        self.test_sample_cdf = self.mod.get_function('test_sample_cdf')
+
+    def compare_sampling(self, hist, reps=10):
+        nbins = hist.GetNbinsX();
+        xaxis = hist.GetXaxis()
+        intg = hist.GetIntegral()
+        cdf_y = np.empty(nbins+1, dtype=float)
+        cdf_x = np.empty_like(cdf_y)
+
+        cdf_x[0] = xaxis.GetBinLowEdge(1)
+        cdf_y[0] = 0.0
+        for i in xrange(1,len(cdf_x)):
+            cdf_y[i] = intg[i]
+            cdf_x[i] = xaxis.GetBinUpEdge(i)
+
+        cdf_x_gpu = gpuarray.to_gpu(cdf_x.astype(np.float32))
+        cdf_y_gpu = gpuarray.to_gpu(cdf_y.astype(np.float32))
+        block =(128,1,1)
+        grid = (128, 1)
+        out_gpu = gpuarray.empty(shape=int(block[0]*grid[0]), dtype=np.float32)
+
+        out_h = ROOT.TH1D('out_h', '', hist.GetNbinsX(), 
+                          xaxis.GetXmin(),
+                          xaxis.GetXmax())
+        out_h.SetLineColor(ROOT.kGreen)
+
+        for i in xrange(reps):
+            self.test_sample_cdf(np.int32(i),
+                                 np.int32(len(cdf_x_gpu)), 
+                                 cdf_x_gpu, cdf_y_gpu, out_gpu, block=block, grid=grid)
+            out = out_gpu.get()
+            for v in out:
+                out_h.Fill(v)
+
+        prob = out_h.KolmogorovTest(hist)
+        return prob, out_h 
+
+    def test_sampling(self):
+        '''Verify that the CDF-based sampler on the GPU reproduces a binned
+        Gaussian distribution'''
+        f = ROOT.TF1('f_gaussian', 'gaus(0)', -5, 5)
+        f.SetParameters(1.0/np.sqrt(np.pi * 2), 0.0, 1.0)
+        gaussian = ROOT.TH1D('gaussian', '', 100, -5, 5)
+        gaussian.Add(f)
+
+        prob, out_h = self.compare_sampling(gaussian, reps=50)
+
+        assert prob > 0.01
+
+    def tearDown(self):
+        self.context.pop()