merge

3 files changed, 104 insertions, 0 deletions

diff --git a/src/random.h b/src/random.h
index b084c19..6abc415 100644
--- a/src/random.h
+++ b/src/random.h
@@ -19,6 +19,30 @@ __device__ float3 uniform_sphere(curandState *s)
 	return make_float3(c*cosf(theta), c*sinf(theta), u);
 }
 
+// Draw a random sample given a cumulative distribution function
+// Assumptions: ncdf >= 2, cdf_y[0] is 0.0, and cdf_y[ncdf-1] is 1.0
+__device__ float sample_cdf(curandState *rng, int ncdf, 
+			    float *cdf_x, float *cdf_y)
+{
+  float u = curand_uniform(rng);
+
+  // Find u in cdf_y with binary search
+  // list must contain at least 2 elements: 0.0 and 1.0
+  int lower=0;
+  int upper=ncdf-1;
+  while(lower < upper-1) {
+    int half = (lower+upper) / 2;
+    if (u < cdf_y[half])
+      upper = half;
+    else
+      lower = half;
+  }
+  
+  float frac = (u - cdf_y[lower]) / (cdf_y[upper] - cdf_y[lower]);
+  return cdf_x[lower] * frac + cdf_x[upper] * (1.0f - frac);
+}
+
+
 extern "C"
 {
 
diff --git a/tests/test_sample_cdf.cu b/tests/test_sample_cdf.cu
new file mode 100644
index 0000000..1401772
--- /dev/null
+++ b/tests/test_sample_cdf.cu
@@ -0,0 +1,16 @@
+// -*-c++-*-
+#include "random.h"
+
+extern "C" {
+
+__global__ void test_sample_cdf(int offset, int ncdf, 
+				float *cdf_x, float *cdf_y, float *out)
+{
+  int id = blockDim.x * blockIdx.x + threadIdx.x;
+  curandState s;
+  curand_init(0, id, offset, &s);
+
+  out[id] = sample_cdf(&s, ncdf, cdf_x, cdf_y);
+}
+
+}
diff --git a/tests/test_sample_cdf.py b/tests/test_sample_cdf.py
new file mode 100644
index 0000000..2baa2ac
--- /dev/null
+++ b/tests/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
+