add tests for individual functions/methods in PGBART (#64)

* add test pgbart, small refactor systematic resampling

* remove comment

Author: aloctavodia

pymc_bart/pgbart.py

@@ -232,7 +232,7 @@ def resample(self, particles, normalized_weights):
 
         Ensure particles are copied only if needed.
         """
-        new_indices = self.systematic(normalized_weights)
+        new_indices = self.systematic(normalized_weights) + 2
         seen = []
         new_particles = []
         for idx in new_indices:
@@ -253,20 +253,20 @@ def get_particle_tree(self, particles, normalized_weights):
         new_index = self.systematic(normalized_weights)[
             discrete_uniform_sampler(self.num_particles)
         ]
-        new_particle = particles[new_index - 2]
+        new_particle = particles[new_index]
         return new_particle, new_particle.tree
 
     def systematic(self, normalized_weights):
         """
         Systematic resampling.
 
-        Return indices in the range 2, ..., len(normalized_weights)+2
+        Return indices in the range 0, ..., len(normalized_weights)
 
         Note: adapted from https://github.com/nchopin/particles
         """
         lnw = len(normalized_weights)
         single_uniform = (self.uniform.random() + np.arange(lnw)) / lnw
-        return inverse_cdf(single_uniform, normalized_weights) + 2
+        return inverse_cdf(single_uniform, normalized_weights)
 
     def init_particles(self, tree_id: int) -> np.ndarray:
         """Initialize particles."""

tests/test_pgbart.py

@@ -0,0 +1,69 @@
+from unittest import TestCase
+import numpy as np
+import pymc as pm
+import pymc_bart as pmb
+from pymc_bart.pgbart import fast_mean, discrete_uniform_sampler, NormalSampler, UniformSampler
+
+
+class TestSystematic(TestCase):
+    def test_systematic(self):
+        X = np.random.normal(0, 1, size=(250, 3))
+        Y = np.random.normal(0, 1, size=250)
+        X[:, 0] = np.random.normal(Y, 0.1)
+
+        with pm.Model() as model:
+            mu = pmb.BART("mu", X, Y, m=10)
+            sigma = pm.HalfNormal("sigma", 1)
+            y = pm.Normal("y", mu, sigma, observed=Y)
+            step = pmb.PGBART([mu])
+
+        normalized_weights = np.array([0.5, 0.3, 0.2])
+        indices = step.systematic(normalized_weights)
+
+        self.assertEqual(len(indices), len(normalized_weights))
+        self.assertEqual(indices.dtype, np.int)
+        self.assertTrue(all(i >= 0 and i < len(normalized_weights) for i in indices))
+
+        normalized_weights = np.array([0, 0.25, 0.75])
+        indices = step.systematic(normalized_weights)
+        self.assertTrue(all(i >= 1 and i < len(normalized_weights) for i in indices))
+
+
+def test_fast_mean():
+    values = np.random.uniform(size=10)
+    np.testing.assert_almost_equal(fast_mean(values), np.mean(values))
+
+    values = np.random.uniform(size=(2, 10))
+    np.testing.assert_array_almost_equal(fast_mean(values), np.mean(values, 1))
+
+
+def test_discrete_uniform():
+    sample = discrete_uniform_sampler(7)
+    assert isinstance(sample, int)
+    samples = np.array([discrete_uniform_sampler(7) for i in range(1000)])
+    assert all(samples >= 0)
+    assert all(samples < 7)
+
+
+def test_normal_sampler():
+    normal = NormalSampler(2, shape=1)
+    samples = np.array([normal.random() for i in range(100000)])
+    np.testing.assert_almost_equal(samples.mean(), 0, decimal=2)
+    np.testing.assert_almost_equal(samples.std(), 2, decimal=2)
+
+    normal = NormalSampler(2, shape=2)
+    samples = np.array([normal.random() for i in range(100000)])
+    np.testing.assert_almost_equal(samples.mean(0), [0, 0], decimal=2)
+    np.testing.assert_almost_equal(samples.std(0), [2, 2], decimal=2)
+
+
+def test_uniform_sampler():
+    uniform = UniformSampler(0.5, 2, shape=1)
+    samples = np.array([uniform.random() for i in range(100000)])
+    np.testing.assert_almost_equal(samples.mean(), 1.25, decimal=2)
+    np.testing.assert_almost_equal(samples.std(), 0.43, decimal=2)
+
+    uniform = UniformSampler(0.5, 2, shape=2)
+    samples = np.array([uniform.random() for i in range(100000)])
+    np.testing.assert_almost_equal(samples.mean(0), [1.25, 1.25], decimal=2)
+    np.testing.assert_almost_equal(samples.std(0), [0.43, 0.43], decimal=2)