Add MCLMC sampler

blackjax/__init__.py

@@ -11,6 +11,7 @@
 from .mcmc.hmc import dynamic_hmc, hmc
 from .mcmc.mala import mala
 from .mcmc.marginal_latent_gaussian import mgrad_gaussian
+from .mcmc.mclmc import mclmc
 from .mcmc.nuts import nuts
 from .mcmc.periodic_orbital import orbital_hmc
 from .mcmc.random_walk import additive_step_random_walk, irmh, rmh
@@ -38,6 +39,7 @@
     "additive_step_random_walk",
     "rmh",
     "irmh",
+    "mclmc",
     "elliptical_slice",
     "ghmc",
     "sgld",  # stochastic gradient mcmc

blackjax/adaptation/step_size.py

@@ -12,12 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Step size adaptation"""
+import warnings
 from typing import Callable, NamedTuple
 
 import jax
 import jax.numpy as jnp
+from scipy.fft import next_fast_len
 
 from blackjax.mcmc.hmc import HMCState
+from blackjax.mcmc.integrators import noneuclidean_mclachlan
+from blackjax.mcmc.mclmc import IntegratorState, build_kernel, init
 from blackjax.optimizers.dual_averaging import dual_averaging
 from blackjax.types import PRNGKey
 
@@ -257,3 +261,285 @@
     rss_state = jax.lax.while_loop(do_continue, update, rss_state)
 
     return rss_state.step_size
+
+
+class MCLMCAdaptationState(NamedTuple):
+    """Tunable parameters for MCLMC"""
+
+    L: float
+    step_size: float
+
+
+def ess_corr(x):
+    """Taken from: https://blackjax-devs.github.io/blackjax/diagnostics.html
+    shape(x) = (num_samples, d)"""
+
+    input_array = jnp.array(
+        [
+            x,
+        ]
+    )
+
+    num_chains = 1  # input_array.shape[0]
+    num_samples = input_array.shape[1]
+
+    mean_across_chain = input_array.mean(axis=1, keepdims=True)
+    # Compute autocovariance estimates for every lag for the input array using FFT.
+    centered_array = input_array - mean_across_chain
+    m = next_fast_len(2 * num_samples)
+    ifft_ary = jnp.fft.rfft(centered_array, n=m, axis=1)
+    ifft_ary *= jnp.conjugate(ifft_ary)
+    autocov_value = jnp.fft.irfft(ifft_ary, n=m, axis=1)
+    autocov_value = (
+        jnp.take(autocov_value, jnp.arange(num_samples), axis=1) / num_samples
+    )
+    mean_autocov_var = autocov_value.mean(0, keepdims=True)
+    mean_var0 = (
+        jnp.take(mean_autocov_var, jnp.array([0]), axis=1)
+        * num_samples
+        / (num_samples - 1.0)
+    )
+    weighted_var = mean_var0 * (num_samples - 1.0) / num_samples
+
+    weighted_var = jax.lax.cond(
+        num_chains > 1,
+        lambda _: weighted_var + mean_across_chain.var(axis=0, ddof=1, keepdims=True),
+        lambda _: weighted_var,
+        operand=None,
+    )
+
+    # Geyer's initial positive sequence
+    num_samples_even = num_samples - num_samples % 2
+    mean_autocov_var_tp1 = jnp.take(
+        mean_autocov_var, jnp.arange(1, num_samples_even), axis=1
+    )
+    rho_hat = jnp.concatenate(
+        [
+            jnp.ones_like(mean_var0),
+            1.0 - (mean_var0 - mean_autocov_var_tp1) / weighted_var,
+        ],
+        axis=1,
+    )
+
+    rho_hat = jnp.moveaxis(rho_hat, 1, 0)
+    rho_hat_even = rho_hat[0::2]
+    rho_hat_odd = rho_hat[1::2]
+
+    mask0 = (rho_hat_even + rho_hat_odd) > 0.0
+    carry_cond = jnp.ones_like(mask0[0])
+    max_t = jnp.zeros_like(mask0[0], dtype=int)
+
+    def positive_sequence_body_fn(state, mask_t):
+        t, carry_cond, max_t = state
+        next_mask = carry_cond & mask_t
+        next_max_t = jnp.where(next_mask, jnp.ones_like(max_t) * t, max_t)
+        return (t + 1, next_mask, next_max_t), next_mask
+
+    (*_, max_t_next), mask = jax.lax.scan(
+        positive_sequence_body_fn, (0, carry_cond, max_t), mask0
+    )
+    indices = jnp.indices(max_t_next.shape)
+    indices = tuple([max_t_next + 1] + [indices[i] for i in range(max_t_next.ndim)])
+    rho_hat_odd = jnp.where(mask, rho_hat_odd, jnp.zeros_like(rho_hat_odd))
+    # improve estimation
+    mask_even = mask.at[indices].set(rho_hat_even[indices] > 0)
+    rho_hat_even = jnp.where(mask_even, rho_hat_even, jnp.zeros_like(rho_hat_even))
+
+    # Geyer's initial monotone sequence
+    def monotone_sequence_body_fn(rho_hat_sum_tm1, rho_hat_sum_t):
+        update_mask = rho_hat_sum_t > rho_hat_sum_tm1
+        next_rho_hat_sum_t = jnp.where(update_mask, rho_hat_sum_tm1, rho_hat_sum_t)
+        return next_rho_hat_sum_t, (update_mask, next_rho_hat_sum_t)
+
+    rho_hat_sum = rho_hat_even + rho_hat_odd
+    _, (update_mask, update_value) = jax.lax.scan(
+        monotone_sequence_body_fn, rho_hat_sum[0], rho_hat_sum
+    )
+
+    rho_hat_even_final = jnp.where(update_mask, update_value / 2.0, rho_hat_even)
+    rho_hat_odd_final = jnp.where(update_mask, update_value / 2.0, rho_hat_odd)
+
+    # compute effective sample size
+    ess_raw = num_chains * num_samples
+    tau_hat = (
+        -1.0
+        + 2.0 * jnp.sum(rho_hat_even_final + rho_hat_odd_final, axis=0)
+        - rho_hat_even_final[indices]
+    )
+
+    tau_hat = jnp.maximum(tau_hat, 1 / jnp.log10(ess_raw))
+    ess = ess_raw / tau_hat
+
+    neff = ess.squeeze() / num_samples
+    return 1.0 / jnp.average(1 / neff)
+
+
+def nan_reject(x, u, l, g, xx, uu, ll, gg, eps, eps_max, dK):
+    """if there are nans, let's reduce the stepsize, and not update the state. The function returns the old state in this case."""
+
+    nonans = jnp.all(jnp.isfinite(xx))
+    return nonans, *jax.tree_util.tree_map(
+        lambda new, old: jax.lax.select(nonans, jnp.nan_to_num(new), old),
+        (xx, uu, ll, gg, eps_max, dK),
+        (x, u, l, g, eps * 0.8, 0.0),
+    )
+
+
+def dynamics_adaptive(dynamics, state, L):
+    """One step of the dynamics with the adaptive stepsize"""
+
+    x, u, l, g, E, Feps, Weps, eps_max, key = state
+
+    eps = jnp.power(
+        Feps / Weps, -1.0 / 6.0
+    )  # We use the Var[E] = O(eps^6) relation here.
+    eps = (eps < eps_max) * eps + (
+        eps > eps_max
+    ) * eps_max  # if the proposed stepsize is above the stepsize where we have seen divergences
+
+    state, info = dynamics(
+        jax.random.PRNGKey(0), IntegratorState(x, u, l, g), L=L, step_size=eps
+    )
+
+    xx, uu, ll, gg = state
+    # ll, gg = -ll, -gg
+    kinetic_change = info.kinetic_change
+
+    varEwanted = 5e-4
+    sigma_xi = 1.5
+    neff = 150  # effective number of steps used to determine the stepsize in the adaptive step
+    gamma = (neff - 1.0) / (neff + 1.0)  # forgeting factor in the adaptive step
+
+    # step updating
+    success, xx, uu, ll, gg, eps_max, kinetic_change = nan_reject(
+        x, u, l, g, xx, uu, ll, gg, eps, eps_max, kinetic_change
+    )
+
+    DE = info.dE  # energy difference
+    EE = E + DE  # energy
+    # Warning: var = 0 if there were nans, but we will give it a very small weight
+    xi = (
+        (DE**2) / (xx.shape[0] * varEwanted)
+    ) + 1e-8  # 1e-8 is added to avoid divergences in log xi
+    w = jnp.exp(
+        -0.5 * jnp.square(jnp.log(xi) / (6.0 * sigma_xi))
+    )  # the weight which reduces the impact of stepsizes which are much larger on much smaller than the desired one.
+    Feps = gamma * Feps + w * (
+        xi / jnp.power(eps, 6.0)
+    )  # Kalman update the linear combinations
+    Weps = gamma * Weps + w
+
+    return xx, uu, ll, gg, EE, Feps, Weps, eps_max, key, eps * success
+
+
+def tune12(kernel, x, u, l, g, random_key, L, eps, num_steps1, num_steps2):
+    """cheap hyperparameter tuning"""
+
+    def step(state, outer_weight):
+        """one adaptive step of the dynamics"""
+        x, u, l, g, E, Feps, Weps, eps_max, key, eps = dynamics_adaptive(
+            kernel, state[0], L
+        )
+        W, F1, F2 = state[1]
+        w = outer_weight * eps
+        zero_prevention = 1 - outer_weight
+        F1 = (W * F1 + w * x) / (
+            W + w + zero_prevention
+        )  # Update <f(x)> with a Kalman filter
+        F2 = (W * F2 + w * jnp.square(x)) / (
+            W + w + zero_prevention
+        )  # Update <f(x)> with a Kalman filter
+        W += w
+
+        return ((x, u, l, g, E, Feps, Weps, eps_max, key), (W, F1, F2)), eps
+
+    # we use the last num_steps2 to compute the diagonal preconditioner
+    outer_weights = jnp.concatenate((jnp.zeros(num_steps1), jnp.ones(num_steps2)))
+
+    # initial state
+    state = (
+        (x, u, l, g, 0.0, jnp.power(eps, -6.0) * 1e-5, 1e-5, jnp.inf, random_key),
+        (0.0, jnp.zeros(len(x)), jnp.zeros(len(x))),
+    )
+    # run the steps
+    state, eps = jax.lax.scan(
+        step, init=state, xs=outer_weights, length=num_steps1 + num_steps2
+    )
+    # determine L
+    if num_steps2 != 0.0:
+        F1, F2 = state[1][1], state[1][2]
+        variances = F2 - jnp.square(F1)
+        sigma2 = jnp.average(variances)
+
+        L = jnp.sqrt(sigma2 * x.shape[0])
+
+    xx, uu, ll, gg, _, _, _, _, _ = state[0]  # the final state
+    return (
+        L,
+        eps[-1],
+        IntegratorState(xx, uu, ll, gg),
+    )  # return the tuned hyperparameters and the final state
+
+
+def tune3(kernel, state, rng_key, L, eps, num_steps):
+    """determine L by the autocorrelations (around 10 effective samples are needed for this to be accurate)"""
+
+    state, info = jax.lax.scan(
+        lambda s, k: (kernel(k, s, L, eps)), state, jax.random.split(rng_key, num_steps)
+    )
+
+    Lfactor = 0.4
+    # ESS2 = effective_sample_size(info.transformed_x)
+    # neff = ESS2.squeeze() / info.transformed_x.shape[0]
+    # ESS_alt = 1.0 / jnp.average(1 / neff)
+    ESS = ess_corr(info.transformed_x)
+    if ESS * num_steps <= 10:
+        warnings.warn(
+            "tune3 cannot be expected to work with 10 or fewer effective samples"
+        )
+
+    Lnew = Lfactor * eps / ESS
+    return Lnew, state
+
+
+def tune(
+    position,
+    logdensity_fn,
+    num_steps: int,
+    rng_key: PRNGKey,
+    params: MCLMCAdaptationState,
+) -> tuple[MCLMCAdaptationState, IntegratorState]:
+    num_tune_step_ratio_1 = 0.1
+    num_tune_step_ratio_2 = 0.1
+
+    kernel = build_kernel(
+        logdensity_fn, integrator=noneuclidean_mclachlan, transform=lambda x: x
+    )
+
+    init_key, tune1_key, tune2_key = jax.random.split(rng_key, 3)
+
+    x, u, l, g = init(position, logdensity_fn=logdensity_fn, rng_key=init_key)
+    # x, u, l, g = (
+    #     jnp.array([0.1, 0.1]),
+    #     jnp.array([-0.6755803, 0.73728645]),
+    #     -0.010000001,
+    #     -jnp.array([0.1, 0.1]),
+    # )
+
+    L, eps, state = tune12(
+        kernel,
+        x,
+        u,
+        l,
+        g,
+        tune1_key,
+        params.L,
+        params.step_size,
+        int(num_steps * num_tune_step_ratio_1),
+        int(num_steps * num_tune_step_ratio_1),
+    )
+
+    L, state = tune3(
+        kernel, state, tune2_key, L, eps, int(num_steps * num_tune_step_ratio_2)
+    )
+    return MCLMCAdaptationState(L, eps), state

blackjax/mcmc/__init__.py

@@ -4,6 +4,7 @@
     hmc,
     mala,
     marginal_latent_gaussian,
+    mclmc,
     nuts,
     periodic_orbital,
     random_walk,
@@ -18,4 +19,5 @@
     "periodic_orbital",
     "marginal_latent_gaussian",
     "random_walk",
+    "mclmc",
 ]

blackjax/mcmc/integrators.py

@@ -1,3 +1,4 @@
+# @title `integrators.py` from https://github.com/blackjax-devs/blackjax/pull/589
 # Copyright 2020- The Blackjax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");