Comparing MCMC algorithms with MCMCInference¶

In addition to the NumPyro inference methods, dynestyx provides the MCMCInference interface for advanced MCMC methods using NumPyro and BlackJAX.

The interface mirrors that of Filter: you pass in a config, which includes all relevant hyperparameters. This is passed to an MCMCInference object, which accepts

• the MCMC config dataclass (NUTSConfig, HMCConfig, AdaptiveMetropolisConfig, SGLDConfig, MALAConfig)
• a model function with signature model(obs_times, obs_values, ctrl_times=None, ctrl_values=None, predict_times=None, ...)

For inference, you still call inference.run(rng_key, obs_times, obs_values) (no predict_times required when you only filter on observations).

For simulation / Predictive, pass the time grid as predict_times=.... Simulator outputs use the f_ prefix and include a leading n_simulations axis (and under Predictive, a leading num_samples axis), e.g. f_observations has shape (num_samples, n_sim, T, obs_dim).

Then inference can be run using

posterior_samples = inference.run(rng_key, obs_times, obs_values)

This works with the various different inference contexts, e.g., Filter or DiscreteTimeSimulator.

In [1]:

import os
import time
from pathlib import Path

os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=4"

import arviz as az
import dynestyx as dsx
import jax.numpy as jnp
import jax.random as jr
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import numpyro
import numpyro.distributions as dist
from numpyro.infer import Predictive

from dynestyx import DiscreteTimeSimulator
from dynestyx.inference.configs.mcmc import (
    AdaptiveMetropolisConfig,
    HMCConfig,
    MALAConfig,
    NUTSConfig,
)
from dynestyx.inference.filters import Filter, KFConfig
from dynestyx.inference.mcmc import MCMCInference
from dynestyx.models.lti_dynamics import LTI_discrete

STYLE_COLORS = {
    "warm_red": "#E64B35",
    "orange": "#E69F00",
    "teal": "#009E73",
    "muted_blue": "#56B4E9",
    "purple": "#8C79B8",
    "olive": "#8DA35B",
    "black": "#222222",
}

mpl.rcParams.update(
    {
        "figure.dpi": 160,
        "savefig.dpi": 300,
        "savefig.bbox": "tight",
        "figure.facecolor": "white",
        "axes.facecolor": "white",
        "font.family": "DejaVu Sans",
        "mathtext.fontset": "dejavusans",
        "font.size": 9,
        "axes.titlesize": 10,
        "axes.labelsize": 9,
        "xtick.labelsize": 8,
        "ytick.labelsize": 8,
        "legend.fontsize": 8,
        "axes.linewidth": 0.9,
        "xtick.direction": "out",
        "ytick.direction": "out",
        "legend.frameon": False,
    }
)


def despine_curve_axis(ax):
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.tick_params(direction="out")

import os import time from pathlib import Path os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=4" import arviz as az import dynestyx as dsx import jax.numpy as jnp import jax.random as jr import matplotlib as mpl import matplotlib.pyplot as plt import numpy as np import numpyro import numpyro.distributions as dist from numpyro.infer import Predictive from dynestyx import DiscreteTimeSimulator from dynestyx.inference.configs.mcmc import ( AdaptiveMetropolisConfig, HMCConfig, MALAConfig, NUTSConfig, ) from dynestyx.inference.filters import Filter, KFConfig from dynestyx.inference.mcmc import MCMCInference from dynestyx.models.lti_dynamics import LTI_discrete STYLE_COLORS = { "warm_red": "#E64B35", "orange": "#E69F00", "teal": "#009E73", "muted_blue": "#56B4E9", "purple": "#8C79B8", "olive": "#8DA35B", "black": "#222222", } mpl.rcParams.update( { "figure.dpi": 160, "savefig.dpi": 300, "savefig.bbox": "tight", "figure.facecolor": "white", "axes.facecolor": "white", "font.family": "DejaVu Sans", "mathtext.fontset": "dejavusans", "font.size": 9, "axes.titlesize": 10, "axes.labelsize": 9, "xtick.labelsize": 8, "ytick.labelsize": 8, "legend.fontsize": 8, "axes.linewidth": 0.9, "xtick.direction": "out", "ytick.direction": "out", "legend.frameon": False, } ) def despine_curve_axis(ax): ax.spines["top"].set_visible(False) ax.spines["right"].set_visible(False) ax.tick_params(direction="out")

/Users/danwaxman/Documents/dynestyx/.venv/lib/python3.12/site-packages/arviz/__init__.py:50: FutureWarning: 
ArviZ is undergoing a major refactor to improve flexibility and extensibility while maintaining a user-friendly interface.
Some upcoming changes may be backward incompatible.
For details and migration guidance, visit: https://python.arviz.org/en/latest/user_guide/migration_guide.html
  warn(

In [2]:

def discrete_time_lti_simplified_model(
    obs_times=None,
    obs_values=None,
    ctrl_times=None,
    ctrl_values=None,
    predict_times=None,
):
    alpha = numpyro.sample("alpha", dist.Uniform(-0.7, 0.7))

    A = jnp.array([[alpha, 0.1], [0.1, 0.8]])
    Q = 0.1 * jnp.eye(2)
    H = jnp.array([[1.0, 0.0]])
    R = jnp.array([[0.5**2]])
    B = jnp.array([[0.1], [0.0]])
    D = jnp.array([[0.01]])

    dynamics = LTI_discrete(A=A, Q=Q, H=H, R=R, B=B, D=D)

    dsx.sample(
        "f",
        dynamics,
        obs_times=obs_times,
        obs_values=obs_values,
        ctrl_times=ctrl_times,
        ctrl_values=ctrl_values,
        predict_times=predict_times,
    )


def make_data(seed=0):
    # Time grid for forward simulation (Predictive API uses predict_times).
    predict_times = jnp.arange(start=0.0, stop=30.0, step=0.05)
    obs_times = predict_times
    true_params = {"alpha": jnp.array(0.35)}

    predictive = Predictive(
        discrete_time_lti_simplified_model,
        params=true_params,
        num_samples=1,
        exclude_deterministic=False,
    )

    with DiscreteTimeSimulator():
        synthetic = predictive(jr.PRNGKey(seed), predict_times=predict_times)

    # f_observations: (num_samples, n_sim, T, obs_dim); take the single trajectory [0, 0].
    obs_values = synthetic["f_observations"][0, 0]
    return obs_times, obs_values, true_params

def discrete_time_lti_simplified_model( obs_times=None, obs_values=None, ctrl_times=None, ctrl_values=None, predict_times=None, ): alpha = numpyro.sample("alpha", dist.Uniform(-0.7, 0.7)) A = jnp.array([[alpha, 0.1], [0.1, 0.8]]) Q = 0.1 * jnp.eye(2) H = jnp.array([[1.0, 0.0]]) R = jnp.array([[0.5**2]]) B = jnp.array([[0.1], [0.0]]) D = jnp.array([[0.01]]) dynamics = LTI_discrete(A=A, Q=Q, H=H, R=R, B=B, D=D) dsx.sample( "f", dynamics, obs_times=obs_times, obs_values=obs_values, ctrl_times=ctrl_times, ctrl_values=ctrl_values, predict_times=predict_times, ) def make_data(seed=0): # Time grid for forward simulation (Predictive API uses predict_times). predict_times = jnp.arange(start=0.0, stop=30.0, step=0.05) obs_times = predict_times true_params = {"alpha": jnp.array(0.35)} predictive = Predictive( discrete_time_lti_simplified_model, params=true_params, num_samples=1, exclude_deterministic=False, ) with DiscreteTimeSimulator(): synthetic = predictive(jr.PRNGKey(seed), predict_times=predict_times) # f_observations: (num_samples, n_sim, T, obs_dim); take the single trajectory [0, 0]. obs_values = synthetic["f_observations"][0, 0] return obs_times, obs_values, true_params

In [3]:

obs_times, obs_values, true_params = make_data(seed=0)

fig, ax = plt.subplots(figsize=(7, 3))
ax.plot(np.asarray(obs_times), np.asarray(obs_values).squeeze(-1), marker="o", lw=1)
ax.set_title("Synthetic observations")
ax.set_xlabel("time")
ax.set_ylabel("y")
plt.show()

print("True alpha:", float(true_params["alpha"]))

obs_times, obs_values, true_params = make_data(seed=0) fig, ax = plt.subplots(figsize=(7, 3)) ax.plot(np.asarray(obs_times), np.asarray(obs_values).squeeze(-1), marker="o", lw=1) ax.set_title("Synthetic observations") ax.set_xlabel("time") ax.set_ylabel("y") plt.show() print("True alpha:", float(true_params["alpha"]))

True alpha: 0.3499999940395355

Benchmarking¶

We will run an informal benchmark for different MCMC algorithms, including

• wall-clock time (seconds)
• ESS (bulk) for alpha
• ESS/sec = ESS / seconds

In [4]:

def _as_chain_draw(x):
    x = np.asarray(x)
    if x.ndim == 1:
        return x[None, :]
    if x.ndim == 2:
        return x
    return x.reshape(x.shape[0], x.shape[1], -1)[..., 0]


def run_one(name, mcmc_config, seed):
    with Filter(filter_config=KFConfig()):
        inference = MCMCInference(
            mcmc_config=mcmc_config,
            model=discrete_time_lti_simplified_model,
        )
        t0 = time.perf_counter()
        posterior = inference.run(jr.PRNGKey(seed), obs_times, obs_values)
        posterior["alpha"].block_until_ready()
        elapsed = time.perf_counter() - t0
        diagnostics = inference.get_diagnostics()

    alpha_chain_draw = _as_chain_draw(posterior["alpha"])
    idata = az.from_dict(posterior={"alpha": alpha_chain_draw})
    ess_bulk = float(
        az.ess(idata, var_names=["alpha"], method="bulk")["alpha"].values
    )
    acceptance = diagnostics.get("mean_acceptance_rate")

    return {
        "name": name,
        "elapsed_sec": elapsed,
        "ess_bulk": ess_bulk,
        "ess_per_sec": ess_bulk / elapsed,
        "mean_acceptance": (
            float(jnp.mean(acceptance)) if acceptance is not None else np.nan
        ),
        "alpha_samples": alpha_chain_draw.reshape(-1),
    }

def _as_chain_draw(x): x = np.asarray(x) if x.ndim == 1: return x[None, :] if x.ndim == 2: return x return x.reshape(x.shape[0], x.shape[1], -1)[..., 0] def run_one(name, mcmc_config, seed): with Filter(filter_config=KFConfig()): inference = MCMCInference( mcmc_config=mcmc_config, model=discrete_time_lti_simplified_model, ) t0 = time.perf_counter() posterior = inference.run(jr.PRNGKey(seed), obs_times, obs_values) posterior["alpha"].block_until_ready() elapsed = time.perf_counter() - t0 diagnostics = inference.get_diagnostics() alpha_chain_draw = _as_chain_draw(posterior["alpha"]) idata = az.from_dict(posterior={"alpha": alpha_chain_draw}) ess_bulk = float( az.ess(idata, var_names=["alpha"], method="bulk")["alpha"].values ) acceptance = diagnostics.get("mean_acceptance_rate") return { "name": name, "elapsed_sec": elapsed, "ess_bulk": ess_bulk, "ess_per_sec": ess_bulk / elapsed, "mean_acceptance": ( float(jnp.mean(acceptance)) if acceptance is not None else np.nan ), "alpha_samples": alpha_chain_draw.reshape(-1), }

In [5]:

num_samples = 500
num_warmup = 500
num_chains = 4

samplers = [
    (
        "NUTS (BlackJAX)",
        NUTSConfig(
            num_samples=num_samples,
            num_warmup=num_warmup,
            num_chains=num_chains,
            mcmc_source="blackjax",
        ),
    ),
    (
        "NUTS (NumPyro)",
        NUTSConfig(
            num_samples=num_samples,
            num_warmup=num_warmup,
            num_chains=num_chains,
            mcmc_source="numpyro",
        ),
    ),
    (
        "Adaptive Metropolis",
        AdaptiveMetropolisConfig(
            num_samples=num_samples,
            num_warmup=num_warmup,
            num_chains=num_chains,
            initial_proposal_scale=0.005,
        ),
    ),
    (
        "HMC (BlackJAX)",
        HMCConfig(
            num_samples=num_samples,
            num_warmup=num_warmup,
            num_chains=num_chains,
            mcmc_source="blackjax",
            step_size=5e-3,
            num_steps=8,
        ),
    ),
    (
        "HMC (NumPyro)",
        HMCConfig(
            num_samples=num_samples,
            num_warmup=num_warmup,
            num_chains=num_chains,
            mcmc_source="numpyro",
            step_size=5e-3,
            num_steps=8,
        ),
    ),
    (
        "MALA (BlackJAX)",
        MALAConfig(
            num_samples=num_samples,
            num_warmup=num_warmup,
            num_chains=num_chains,
            mcmc_source="blackjax",
            step_size=5e-3,
        ),
    ),
]

results = []
for i, (name, cfg) in enumerate(samplers):
    print(f"Running {name}...")
    results.append(run_one(name, cfg, seed=100 + i))

results_sorted = sorted(results, key=lambda r: r["ess_per_sec"], reverse=True)
for result in results_sorted:
    acceptance = result["mean_acceptance"]
    acceptance_text = f"{acceptance:.3f}" if np.isfinite(acceptance) else "n/a"
    print(
        f"{result['name']:<36} "
        f"time={result['elapsed_sec']:.2f}s  "
        f"accept={acceptance_text:>5}  "
        f"ESS={result['ess_bulk']:.1f}  "
        f"ESS/sec={result['ess_per_sec']:.2f}"
    )

num_samples = 500 num_warmup = 500 num_chains = 4 samplers = [ ( "NUTS (BlackJAX)", NUTSConfig( num_samples=num_samples, num_warmup=num_warmup, num_chains=num_chains, mcmc_source="blackjax", ), ), ( "NUTS (NumPyro)", NUTSConfig( num_samples=num_samples, num_warmup=num_warmup, num_chains=num_chains, mcmc_source="numpyro", ), ), ( "Adaptive Metropolis", AdaptiveMetropolisConfig( num_samples=num_samples, num_warmup=num_warmup, num_chains=num_chains, initial_proposal_scale=0.005, ), ), ( "HMC (BlackJAX)", HMCConfig( num_samples=num_samples, num_warmup=num_warmup, num_chains=num_chains, mcmc_source="blackjax", step_size=5e-3, num_steps=8, ), ), ( "HMC (NumPyro)", HMCConfig( num_samples=num_samples, num_warmup=num_warmup, num_chains=num_chains, mcmc_source="numpyro", step_size=5e-3, num_steps=8, ), ), ( "MALA (BlackJAX)", MALAConfig( num_samples=num_samples, num_warmup=num_warmup, num_chains=num_chains, mcmc_source="blackjax", step_size=5e-3, ), ), ] results = [] for i, (name, cfg) in enumerate(samplers): print(f"Running {name}...") results.append(run_one(name, cfg, seed=100 + i)) results_sorted = sorted(results, key=lambda r: r["ess_per_sec"], reverse=True) for result in results_sorted: acceptance = result["mean_acceptance"] acceptance_text = f"{acceptance:.3f}" if np.isfinite(acceptance) else "n/a" print( f"{result['name']:<36} " f"time={result['elapsed_sec']:.2f}s " f"accept={acceptance_text:>5} " f"ESS={result['ess_bulk']:.1f} " f"ESS/sec={result['ess_per_sec']:.2f}" )

Running NUTS (BlackJAX)...
Running NUTS (NumPyro)...

  0%|          | 0/1000 [00:00<?, ?it/s]

  0%|          | 0/1000 [00:00<?, ?it/s]

  0%|          | 0/1000 [00:00<?, ?it/s]

  0%|          | 0/1000 [00:00<?, ?it/s]

Running Adaptive Metropolis...
Running HMC (BlackJAX)...
Running HMC (NumPyro)...

/Users/danwaxman/Documents/dynestyx/dynestyx/inference/mcmc.py:133: UserWarning: If both `num_steps` and `trajectory_length` are specified step size can't be adapted
  HMC(

  0%|          | 0/1000 [00:00<?, ?it/s]

  0%|          | 0/1000 [00:00<?, ?it/s]

  0%|          | 0/1000 [00:00<?, ?it/s]

  0%|          | 0/1000 [00:00<?, ?it/s]

Running MALA (BlackJAX)...
Adaptive Metropolis                  time=4.27s  accept=0.448  ESS=523.0  ESS/sec=122.53
NUTS (NumPyro)                       time=10.75s  accept=0.913  ESS=831.9  ESS/sec=77.37
NUTS (BlackJAX)                      time=15.35s  accept=0.932  ESS=725.5  ESS/sec=47.27
MALA (BlackJAX)                      time=11.74s  accept=  n/a  ESS=33.0  ESS/sec=2.81
HMC (BlackJAX)                       time=40.36s  accept=  n/a  ESS=29.9  ESS/sec=0.74
HMC (NumPyro)                        time=25.09s  accept=  n/a  ESS=5.4  ESS/sec=0.22

Finally, besides ESS, we should make sure each method provides accurate results!

In [6]:

method_colors = {
    "NUTS (BlackJAX)": STYLE_COLORS["warm_red"],
    "NUTS (NumPyro)": STYLE_COLORS["orange"],
    "Adaptive Metropolis": STYLE_COLORS["teal"],
    "HMC (BlackJAX)": STYLE_COLORS["muted_blue"],
    "HMC (NumPyro)": STYLE_COLORS["purple"],
    "MALA (BlackJAX)": STYLE_COLORS["olive"],
}

fig, ax = plt.subplots(figsize=(9, 4.8))
for result in results:
    ax.hist(
        result["alpha_samples"],
        bins=35,
        density=True,
        histtype="step",
        color=method_colors[result["name"]],
        linewidth=1.7,
        label=result["name"],
    )

ax.axvline(
    float(true_params["alpha"]),
    color=STYLE_COLORS["black"],
    linestyle="--",
    linewidth=1.4,
    label="truth",
)
ax.set(xlabel=r"$\alpha$", ylabel="density")
despine_curve_axis(ax)
ax.legend(
    ncols=4,
    loc="upper center",
    bbox_to_anchor=(0.5, -0.16),
)
fig.subplots_adjust(bottom=0.24)
output_dir = Path("../../outputs")
output_dir.mkdir(parents=True, exist_ok=True)
fig.savefig(output_dir / "mcmc_algorithm_comparison.pdf")
fig.savefig(output_dir / "mcmc_algorithm_comparison.png", dpi=300)
plt.show()

method_colors = { "NUTS (BlackJAX)": STYLE_COLORS["warm_red"], "NUTS (NumPyro)": STYLE_COLORS["orange"], "Adaptive Metropolis": STYLE_COLORS["teal"], "HMC (BlackJAX)": STYLE_COLORS["muted_blue"], "HMC (NumPyro)": STYLE_COLORS["purple"], "MALA (BlackJAX)": STYLE_COLORS["olive"], } fig, ax = plt.subplots(figsize=(9, 4.8)) for result in results: ax.hist( result["alpha_samples"], bins=35, density=True, histtype="step", color=method_colors[result["name"]], linewidth=1.7, label=result["name"], ) ax.axvline( float(true_params["alpha"]), color=STYLE_COLORS["black"], linestyle="--", linewidth=1.4, label="truth", ) ax.set(xlabel=r"$\alpha$", ylabel="density") despine_curve_axis(ax) ax.legend( ncols=4, loc="upper center", bbox_to_anchor=(0.5, -0.16), ) fig.subplots_adjust(bottom=0.24) output_dir = Path("../../outputs") output_dir.mkdir(parents=True, exist_ok=True) fig.savefig(output_dir / "mcmc_algorithm_comparison.pdf") fig.savefig(output_dir / "mcmc_algorithm_comparison.png", dpi=300) plt.show()