Tests

tests/test_mclmc.py

@@ -0,0 +1,45 @@
+import sys  
+sys.path.insert(0, '../../')
+sys.path.insert(0, './')
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import matplotlib.pyplot as plt
+
+from sampling.sampler import Sampler
+
+nlogp = lambda x: 0.5*jnp.sum(jnp.square(x))
+value_grad = jax.value_and_grad(nlogp)
+
+class StandardGaussian():
+
+  def __init__(self, d):
+    self.d = d
+
+  def grad_nlogp(self, x):
+    """should return nlogp and gradient of nlogp"""
+    return value_grad(x)
+
+  def transform(self, x):
+    return x[:2]
+    #return x
+
+  def prior_draw(self, key):
+    """Args: jax random key
+       Returns: one random sample from the prior"""
+
+    return jax.random.normal(key, shape = (self.d, ), dtype = 'float64') * 4
+
+target = StandardGaussian(d = 10)
+sampler = Sampler(target, varEwanted = 5e-4)
+
+
+def test_mclmc():
+    samples1 = sampler.sample(100, 1, random_key=jax.random.PRNGKey(0))
+    samples2 = sampler.sample(100, 1, random_key=jax.random.PRNGKey(0))
+    samples3 = sampler.sample(100, 1, random_key=jax.random.PRNGKey(1))
+    assert jnp.array_equal(samples1,samples2), "sampler should be pure"
+    assert not jnp.array_equal(samples1,samples3), "this suggests that seed is not being used"
+    # run with multiple chains
+    sampler.sample(100, 3)

tests/test_momentum_update.py

@@ -7,9 +7,9 @@
 
 import jax.numpy as jnp
 
-def update_momentum_unstable(d, eps):
+def update_momentum_unstable(d):
 
-    def update(u, g):
+    def update(eps, u, g):
         g_norm = jnp.linalg.norm(g)
         e = - g / g_norm
         delta = eps * g_norm / (d-1)
@@ -26,11 +26,10 @@ def test_momentum_update():
     u = jax.random.uniform(key=jax.random.PRNGKey(0),shape=(d,))
     u = u / jnp.linalg.norm(u)
     g = jax.random.uniform(key=jax.random.PRNGKey(1),shape=(d,))
-    update_stable = update_momentum(d, eps)
-    update_unstable = update_momentum_unstable(d, eps)
-    update1 = update_stable(u, g)[0]
-    update2 = update_unstable(u, g)
+    update_stable = update_momentum(d, sequential=True)
+    update_unstable = update_momentum_unstable(d)
+    update1 = update_stable(eps, u, g)[0]
+    update2 = update_unstable(eps, u, g)
     print(update1, update2)
     assert jnp.allclose(update1,update2)
 
-