Rewrite test_sample_cdf to use a unittest.TestCase class so that the GPU context can be setup and torn down correctly around the test.

1 files changed, 50 insertions, 47 deletions

diff --git a/test/test_sample_cdf.py b/test/test_sample_cdf.py
index d12efa4..1510768 100644
--- a/test/test_sample_cdf.py
+++ b/test/test_sample_cdf.py
@@ -1,64 +1,67 @@
-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
+import unittest
+import chroma
 
-current_directory = os.path.split(os.path.realpath(__file__))[0]
-from chroma.cuda import srcdir as source_directory
+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')
 
-source = open(current_directory + '/test_sample_cdf.cu').read()
+    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)
 
-mod = SourceModule(source, options=['-I' + source_directory], no_extern_c=True, cache_dir=False)
+        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)
 
-test_sample_cdf = mod.get_function('test_sample_cdf')
+        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)
 
-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)
+        out_h = ROOT.TH1D('out_h', '', hist.GetNbinsX(), 
+                          xaxis.GetXmin(),
+                          xaxis.GetXmax())
+        out_h.SetLineColor(ROOT.kGreen)
 
-    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)
+        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)
 
-    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)
+        prob = out_h.KolmogorovTest(hist)
+        return prob, out_h 
 
-    out_h = ROOT.TH1D('out_h', '', hist.GetNbinsX(), 
-                      xaxis.GetXmin(),
-                      xaxis.GetXmax())
-    out_h.SetLineColor(ROOT.kGreen)
+    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)
 
-    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 = self.compare_sampling(gaussian, reps=50)
 
-    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
+        assert prob > 0.01
 
+    def tearDown(self):
+        self.context.pop()