Add MH variants and MALA (#18)

* Add MH variants and MALA

examples/bivariate_gaussian_snl.py

@@ -100,6 +100,7 @@ def run():
             params=params,
             observable=y_observed,
             data=data,
+            sampler="imh",
         )
         params, info = snl.fit(
             jr.fold_in(jr.PRNGKey(23), i), data=data, optimizer=optimizer

examples/slcp_ssnl.py

@@ -5,6 +5,7 @@
 
 import argparse
 from functools import partial
+from timeit import default_timer as timer
 
 import distrax
 import haiku as hk
@@ -189,6 +190,7 @@ def run(use_surjectors):
     optimizer = optax.adam(1e-3)
 
     data, params = None, {}
+    start = timer()
     for i in range(5):
         data, _ = snl.simulate_data_and_possibly_append(
             jr.fold_in(jr.PRNGKey(12), i),
@@ -200,6 +202,8 @@ def run(use_surjectors):
         params, info = snl.fit(
             jr.fold_in(jr.PRNGKey(23), i), data=data, optimizer=optimizer
         )
+    end = timer()
+    print(end - start)
 
     sample_key, rng_key = jr.split(jr.PRNGKey(123))
     snl_samples, _ = snl.sample_posterior(sample_key, params, y_observed)

pyproject.toml

@@ -20,7 +20,7 @@ classifiers = [
 ]
 requires-python = ">=3.8"
 dependencies = [
-    "blackjax-nightly>=0.9.6.post127",
+    "blackjax-nightly>=1.0.0.post17",
     "distrax>=0.1.2",
     "dm-haiku>=0.0.9",
     "optax>=0.1.3",

sbijax/generator.py

@@ -1,6 +1,7 @@
 from collections import namedtuple
 
 import chex
+import numpy as np
 from jax import lax
 from jax import numpy as jnp
 from jax import random as jr
@@ -62,7 +63,7 @@ def as_batch_iterator(
 
     def get_batch(idx, idxs=idxs):
         start_idx = idx * batch_size
-        step_size = jnp.minimum(n - start_idx, batch_size)
+        step_size = np.minimum(n - start_idx, batch_size)
         ret_idx = lax.dynamic_slice_in_dim(idxs, idx * batch_size, step_size)
         batch = {
             name: lax.index_take(array, (ret_idx,), axes=(0,))

sbijax/mcmc/irmh.py

@@ -0,0 +1,73 @@
+import blackjax as bj
+import distrax
+import jax
+from jax import random as jr
+
+
+# pylint: disable=too-many-arguments,unused-argument
+def sample_with_imh(
+    rng_key, lp, prior, *, n_chains=4, n_samples=2_000, n_warmup=1_000, **kwargs
+):
+    """
+    Sample from a distribution using the indepdendent Metropolis-Hastings
+    sampler.
+
+    Parameters
+    ----------
+    rng_seq: hk.PRNGSequence
+        a hk.PRNGSequence
+    lp: Callable
+        the logdensity you wish to sample from
+    prior: Callable
+        a function that returns a prior sample
+    n_chains: int
+        number of chains to sample
+    n_samples: int
+        number of samples per chain
+    n_warmup: int
+        number of samples to discard
+
+    Returns
+    -------
+    jnp.ndarrau
+        a JAX array of dimension n_samples \times n_chains \times len_theta
+    """
+
+    def _inference_loop(rng_key, kernel, initial_state, n_samples):
+        @jax.jit
+        def _step(states, rng_key):
+            keys = jax.random.split(rng_key, n_chains)
+            states, _ = jax.vmap(kernel)(keys, states)
+            return states, states
+
+        sampling_keys = jax.random.split(rng_key, n_samples)
+        _, states = jax.lax.scan(_step, initial_state, sampling_keys)
+        return states
+
+    init_key, rng_key = jr.split(rng_key)
+    initial_states, kernel = _mh_init(init_key, n_chains, prior, lp)
+
+    states = _inference_loop(init_key, kernel, initial_states, n_samples)
+    _ = states.position["theta"].block_until_ready()
+    thetas = states.position["theta"][n_warmup:, :, :]
+
+    return thetas
+
+
+def _irmh_proposal_distribution(shape):
+    def fn(rng_key):
+        return {"theta": jax.random.normal(rng_key, shape=(shape,))}
+
+    return fn
+
+
+# pylint: disable=missing-function-docstring,no-member
+def _mh_init(rng_key, n_chains, prior: distrax.Distribution, lp):
+    init_key, rng_key = jr.split(rng_key)
+    initial_positions = prior(seed=init_key, sample_shape=(n_chains,))
+    kernel = bj.irmh(
+        lp, _irmh_proposal_distribution(initial_positions.shape[1])
+    )
+    initial_positions = {"theta": initial_positions}
+    initial_state = jax.vmap(kernel.init)(initial_positions)
+    return initial_state, kernel.step

sbijax/mcmc/mala.py

@@ -0,0 +1,63 @@
+import blackjax as bj
+import distrax
+import jax
+from jax import random as jr
+
+
+# pylint: disable=too-many-arguments,unused-argument
+def sample_with_mala(
+    rng_key, lp, prior, *, n_chains=4, n_samples=2_000, n_warmup=1_000, **kwargs
+):
+    """
+    Sample from a distribution using the MALA sampler.
+
+    Parameters
+    ----------
+    rng_seq: hk.PRNGSequence
+        a hk.PRNGSequence
+    lp: Callable
+        the logdensity you wish to sample from
+    prior: Callable
+        a function that returns a prior sample
+    n_chains: int
+        number of chains to sample
+    n_samples: int
+        number of samples per chain
+    n_warmup: int
+        number of samples to discard
+
+    Returns
+    -------
+    jnp.ndarrau
+        a JAX array of dimension n_samples \times n_chains \times len_theta
+    """
+
+    def _inference_loop(rng_key, kernel, initial_state, n_samples):
+        @jax.jit
+        def _step(states, rng_key):
+            keys = jax.random.split(rng_key, n_chains)
+            states, _ = jax.vmap(kernel)(keys, states)
+            return states, states
+
+        sampling_keys = jax.random.split(rng_key, n_samples)
+        _, states = jax.lax.scan(_step, initial_state, sampling_keys)
+        return states
+
+    init_key, rng_key = jr.split(rng_key)
+    initial_states, kernel = _mala_init(init_key, n_chains, prior, lp)
+
+    states = _inference_loop(init_key, kernel, initial_states, n_samples)
+    _ = states.position["theta"].block_until_ready()
+    thetas = states.position["theta"][n_warmup:, :, :]
+
+    return thetas
+
+
+# pylint: disable=missing-function-docstring,no-member
+def _mala_init(rng_key, n_chains, prior: distrax.Distribution, lp):
+    init_key, rng_key = jr.split(rng_key)
+    initial_positions = prior(seed=init_key, sample_shape=(n_chains,))
+    kernel = bj.mala(lp, 1.0)
+    initial_positions = {"theta": initial_positions}
+    initial_state = jax.vmap(kernel.init)(initial_positions)
+    return initial_state, kernel.step

sbijax/mcmc/rmh.py

@@ -0,0 +1,69 @@
+import blackjax as bj
+import distrax
+import jax
+from jax import numpy as jnp
+from jax import random as jr
+
+
+# pylint: disable=too-many-arguments,unused-argument
+def sample_with_rmh(
+    rng_key, lp, prior, *, n_chains=4, n_samples=2_000, n_warmup=1_000, **kwargs
+):
+    """
+    Sample from a distribution using the Rosenbluth-Metropolis-Hastings sampler.
+
+    Parameters
+    ----------
+    rng_seq: hk.PRNGSequence
+        a hk.PRNGSequence
+    lp: Callable
+        the logdensity you wish to sample from
+    prior: Callable
+        a function that returns a prior sample
+    n_chains: int
+        number of chains to sample
+    n_samples: int
+        number of samples per chain
+    n_warmup: int
+        number of samples to discard
+
+    Returns
+    -------
+    jnp.ndarrau
+        a JAX array of dimension n_samples \times n_chains \times len_theta
+    """
+
+    def _inference_loop(rng_key, kernel, initial_state, n_samples):
+        @jax.jit
+        def _step(states, rng_key):
+            keys = jax.random.split(rng_key, n_chains)
+            states, _ = jax.vmap(kernel)(keys, states)
+            return states, states
+
+        sampling_keys = jax.random.split(rng_key, n_samples)
+        _, states = jax.lax.scan(_step, initial_state, sampling_keys)
+        return states
+
+    init_key, rng_key = jr.split(rng_key)
+    initial_states, kernel = _mh_init(init_key, n_chains, prior, lp)
+
+    states = _inference_loop(init_key, kernel, initial_states, n_samples)
+    _ = states.position["theta"].block_until_ready()
+    thetas = states.position["theta"][n_warmup:, :, :]
+
+    return thetas
+
+
+# pylint: disable=missing-function-docstring,no-member
+def _mh_init(rng_key, n_chains, prior: distrax.Distribution, lp):
+    init_key, rng_key = jr.split(rng_key)
+    initial_positions = prior(seed=init_key, sample_shape=(n_chains,))
+    kernel = bj.rmh(
+        lp,
+        bj.mcmc.random_walk.normal(
+            jnp.full_like(initial_positions.shape[1], 0.25)
+        ),
+    )
+    initial_positions = {"theta": initial_positions}
+    initial_state = jax.vmap(kernel.init)(initial_positions)
+    return initial_state, kernel.step

sbijax/snl.py

@@ -12,6 +12,9 @@
 from sbijax.mcmc import mcmc_diagnostics, sample_with_nuts, sample_with_slice
 
 # pylint: disable=too-many-arguments,unused-argument
+from sbijax.mcmc.irmh import sample_with_imh
+from sbijax.mcmc.mala import sample_with_mala
+from sbijax.mcmc.rmh import sample_with_rmh
 from sbijax.nn.early_stopping import EarlyStopping
 
 
@@ -283,6 +286,27 @@ def lp__(theta):
                 return jax.vmap(_joint_logdensity_fn)(theta)
 
             sampling_fn = sample_with_slice
+        elif "sampler" in kwargs and kwargs["sampler"] == "rmh":
+            kwargs.pop("sampler", None)
+
+            def lp__(theta):
+                return _joint_logdensity_fn(**theta)
+
+            sampling_fn = sample_with_rmh
+        elif "sampler" in kwargs and kwargs["sampler"] == "imh":
+            kwargs.pop("sampler", None)
+
+            def lp__(theta):
+                return _joint_logdensity_fn(**theta)
+
+            sampling_fn = sample_with_imh
+        elif "sampler" in kwargs and kwargs["sampler"] == "mala":
+            kwargs.pop("sampler", None)
+
+            def lp__(theta):
+                return _joint_logdensity_fn(**theta)
+
+            sampling_fn = sample_with_mala
         else:
 
             def lp__(theta):