Add the MALA sampling algorithm

Author: jeremiecoullon

blackjax/__init__.py

@@ -3,6 +3,7 @@
 from .kernels import (
     adaptive_tempered_smc,
     hmc,
+    mala,
     nuts,
     rmh,
     tempered_smc,
@@ -13,6 +14,7 @@
 
 __all__ = [
     "hmc",  # mcmc
+    "mala",
     "nuts",
     "rmh",
     "window_adaptation",  # mcmc adaptation

blackjax/kernels.py

@@ -11,6 +11,7 @@
 __all__ = [
     "adaptive_tempered_smc",
     "hmc",
+    "mala",
     "nuts",
     "rmh",
     "tempered_smc",
@@ -222,6 +223,77 @@ def step_fn(rng_key: PRNGKey, state):
         return SamplingAlgorithm(init_fn, step_fn)
 
 
+class mala:
+    """Implements the (basic) user interface for the MALA kernel.
+
+    The general mala kernel (:meth:`blackjax.mcmc.mala.kernel`, alias `blackjax.mala.kernel`) can be
+    cumbersome to manipulate. Since most users only need to specify the kernel
+    parameters at initialization time, we provide a helper function that
+    specializes the general kernel.
+
+    We also add the general kernel and state generator as an attribute to this class so
+    users only need to pass `blackjax.mala` to SMC, adaptation, etc. algorithms.
+
+    Examples
+    --------
+
+    A new MALA kernel can be initialized and used with the following code:
+
+    .. code::
+
+        mala = blackjax.mala(logprob_fn, step_size)
+        state = mala.init(position)
+        new_state, info = mala.step(rng_key, state)
+
+    Kernels are not jit-compiled by default so you will need to do it manually:
+
+    .. code::
+
+       step = jax.jit(mala.step)
+       new_state, info = step(rng_key, state)
+
+    Should you need to you can always use the base kernel directly:
+
+    .. code::
+
+       kernel = blackjax.mala.kernel(logprob_fn)
+       state = blackjax.mala.init(position, logprob_fn)
+       state, info = kernel(rng_key, state, logprob_fn, step_size)
+
+    Parameters
+    ----------
+    logprob_fn
+        The logprobability density function we wish to draw samples from. This
+        is minus the potential function.
+    step_size
+        The value to use for the step size in the symplectic integrator.
+
+    Returns
+    -------
+    A ``SamplingAlgorithm``.
+
+    """
+
+    init = staticmethod(mcmc.mala.init)
+    kernel = staticmethod(mcmc.mala.kernel)
+
+    def __new__(  # type: ignore[misc]
+        cls,
+        logprob_fn: Callable,
+        step_size: float,
+    ) -> SamplingAlgorithm:
+
+        step = cls.kernel()
+
+        def init_fn(position: PyTree):
+            return cls.init(position, logprob_fn)
+
+        def step_fn(rng_key: PRNGKey, state):
+            return step(rng_key, state, logprob_fn, step_size)
+
+        return SamplingAlgorithm(init_fn, step_fn)
+
+
 class nuts:
     """Implements the (basic) user interface for the nuts kernel.
 

blackjax/mcmc/__init__.py

@@ -1,3 +1,3 @@
-from . import hmc, nuts, rmh
+from . import hmc, mala, nuts, rmh
 
-__all__ = ["hmc", "nuts", "rmh"]
+__all__ = ["hmc", "mala", "nuts", "rmh"]

blackjax/mcmc/diffusion.py

@@ -0,0 +1,39 @@
+"""Solvers for Langevin diffusions."""
+from typing import NamedTuple
+
+import jax
+import jax.numpy as jnp
+
+from blackjax.types import PRNGKey, PyTree
+
+__all__ = ["overdamped_langevin"]
+
+
+class DiffusionState(NamedTuple):
+    position: PyTree
+    logprob_grad: PyTree
+
+
+def generate_gaussian_noise(rng_key: PRNGKey, position):
+    position_flat, unravel_fn = jax.flatten_util.ravel_pytree(position)
+    noise_flat = jax.random.normal(rng_key, shape=jnp.shape(position_flat))
+    return unravel_fn(noise_flat)
+
+
+def overdamped_langevin(logprob_grad_fn):
+    """Euler solver for overdamped Langevin diffusion."""
+
+    def one_step(rng_key, state: DiffusionState, step_size: float, batch: tuple = ()):
+        position, logprob_grad = state
+        noise = generate_gaussian_noise(rng_key, position)
+        position = jax.tree_util.tree_multimap(
+            lambda p, g, n: p + step_size * g + jnp.sqrt(2 * step_size) * n,
+            position,
+            logprob_grad,
+            noise,
+        )
+
+        logprob_grad = logprob_grad_fn(position, *batch)
+        return DiffusionState(position, logprob_grad)
+
+    return one_step

blackjax/mcmc/mala.py

@@ -0,0 +1,103 @@
+"""Public API for Metropolis Adjusted Langevin kernels."""
+from typing import Callable, NamedTuple, Tuple
+
+import jax
+import jax.numpy as jnp
+
+from blackjax.mcmc.diffusion import overdamped_langevin
+from blackjax.types import PRNGKey, PyTree
+
+__all__ = ["MALAState", "MALAInfo", "init", "kernel"]
+
+
+class MALAState(NamedTuple):
+    """State of the MALA algorithm.
+
+    The MALA algorithm takes one position of the chain and returns another
+    position. In order to make computations more efficient, we also store
+    the current log-probability density as well as the current gradient of the
+    log-probability density.
+
+    """
+
+    position: PyTree
+    logprob_grad: PyTree
+
+
+class MALAInfo(NamedTuple):
+    """Additional information on the MALA transition.
+
+    This additional information can be used for debugging or computing
+    diagnostics.
+
+    acceptance_probability
+        The acceptance probability of the transition.
+    is_accepted
+        Whether the proposed position was accepted or the original position
+        was returned.
+
+    """
+
+    acceptance_probability: float
+    is_accepted: bool
+
+
+def init(position: PyTree, logprob_fn: Callable) -> MALAState:
+    grad_fn = jax.grad(logprob_fn)
+    logprob_grad = grad_fn(position)
+    return MALAState(position, logprob_grad)
+
+
+def kernel():
+    """Build a MALA kernel.
+
+    Returns
+    -------
+    A kernel that takes a rng_key and a Pytree that contains the current state
+    of the chain and that returns a new state of the chain along with
+    information about the transition.
+
+    """
+
+    def transition_probability(state, new_state, step_size):
+        """Transition probability to go from `state` to `new_state`"""
+        theta = new_state.position - state.position - step_size * state.logprob_grad
+        return -0.25 * (1.0 / step_size) * jnp.dot(theta, theta)
+
+    def one_step(
+        rng_key: PRNGKey, state: MALAState, logprob_fn: Callable, step_size: float
+    ) -> Tuple[MALAState, MALAInfo]:
+        """Generate a new sample with the MALA kernel.
+
+        TODO expand the docstring.
+
+        """
+        grad_fn = jax.grad(logprob_fn)
+        integrator = overdamped_langevin(grad_fn)
+
+        key_integrator, key_rmh = jax.random.split(rng_key)
+
+        new_state = integrator(key_integrator, state, step_size)
+
+        delta = (
+            logprob_fn(new_state.position)
+            - logprob_fn(state.position)
+            + transition_probability(new_state, state, step_size)
+            - transition_probability(state, new_state, step_size)
+        )
+        delta = jnp.where(jnp.isnan(delta), -jnp.inf, delta)
+        p_accept = jnp.clip(jnp.exp(delta), a_max=1)
+
+        do_accept = jax.random.bernoulli(key_rmh, p_accept)
+
+        new_state = MALAState(*new_state)
+        info = MALAInfo(p_accept, do_accept)
+
+        return jax.lax.cond(
+            do_accept,
+            lambda _: (new_state, info),
+            lambda _: (state, info),
+            operand=None,
+        )
+
+    return one_step

docs/sampling.rst

@@ -8,6 +8,7 @@ Sampling
 
   hmc
   nuts
+  mala
   rmh
   tempered_smc
   adaptive_tempered_smc
@@ -50,6 +51,11 @@ HMC
 
 .. autoclass:: blackjax.hmc
 
+MALA
+~~~~
+
+.. autoclass:: blackjax.mala
+
 NUTS
 ~~~~
 

tests/test_sampling.py

@@ -121,6 +121,13 @@ def test_linear_regression(self, case, is_mass_matrix_diagonal):
         "num_sampling_steps": 20_000,
         "burnin": 5_000,
     },
+    {
+        "algorithm": blackjax.mala,
+        "initial_position": 1.0,
+        "parameters": {"step_size": 1e-1},
+        "num_sampling_steps": 20_000,
+        "burnin": 2_000,
+    },
 ]