Merge branch 'master' of github.com:JakobRobnik/MicroCanonicalHMC

sampling/correlation_length.py

@@ -1,7 +1,7 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from scipy.fftpack import next_fast_len
+from scipy.fftpack import next_fast_len #type: ignore
 
 
 

sampling/sampler.py

@@ -7,12 +7,26 @@
 from . import dynamics
 from .correlation_length import ess_corr
 
+class Target():
 
+  def __init__(self, d, nlogp):
+    self.d = d
+    self.nlogp = nlogp
+    self.grad_nlogp = jax.value_and_grad(self.nlogp)
+
+  def transform(self, x):
+    return x
+
+  def prior_draw(self, key):
+    """Args: jax random key
+       Returns: one random sample from the prior"""
+
+    raise Exception("Not implemented")
 
 class Sampler:
     """the MCHMC (q = 0 Hamiltonian) sampler"""
 
-    def __init__(self, Target, L = None, eps = None,
+    def __init__(self, Target : Target, L = None, eps = None,
                  integrator = 'MN', varEwanted = 5e-4,
                  diagonal_preconditioning= False,
                  frac_tune1 = 0.1, frac_tune2 = 0.1, frac_tune3 = 0.1,

tests/test_mclmc.py

@@ -0,0 +1,42 @@
+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, Target
+
+nlogp = lambda x: 0.5*jnp.sum(jnp.square(x))
+
+class StandardGaussian(Target):
+
+  def __init__(self, d, nlogp):
+    Target.__init__(self,d,nlogp)
+
+  def transform(self, x):
+    return x[:2]
+
+  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, nlogp=nlogp)
+sampler = Sampler(target, varEwanted = 5e-4)
+
+target_simple = Target(d = 10, nlogp=nlogp)
+
+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)
+
+    Sampler(target_simple).sample(100, x_initial = jax.random.normal(shape=(10,), key=jax.random.PRNGKey(0)))

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)
 
-