Move CUDA source inside chroma package, rename tests directory to test

1 files changed, 64 insertions, 0 deletions

diff --git a/test/test_sample_cdf.py b/test/test_sample_cdf.py
new file mode 100644
index 0000000..2baa2ac
--- /dev/null
+++ b/test/test_sample_cdf.py
@@ -0,0 +1,64 @@
+import pycuda.autoinit
+import pycuda.driver as cuda
+from pycuda.compiler import SourceModule
+from pycuda import gpuarray
+import numpy as np
+import ROOT
+import os
+
+current_directory = os.path.split(os.path.realpath(__file__))[0]
+source_directory = current_directory + '/../src'
+
+source = open(current_directory + '/test_sample_cdf.cu').read()
+
+mod = SourceModule(source, options=['-I' + source_directory], no_extern_c=True, cache_dir=False)
+
+test_sample_cdf = mod.get_function('test_sample_cdf')
+
+def compare_sampling(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):
+        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():
+    '''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 = compare_sampling(gaussian, reps=50)
+
+    assert prob > 0.01
+