Show code cell source

Hide code cell source

import matplotlib.pyplot as plt
%matplotlib inline
import matplotlib_inline
matplotlib_inline.backend_inline.set_matplotlib_formats('svg')
import seaborn as sns

# Uncomment the next two lines if running for the book
import warnings
warnings.filterwarnings("ignore")

from functools import partial
import jax
from jax import lax, tree, vmap, jit
import jax.random as jr
import jax.numpy as jnp
from jax.flatten_util import ravel_pytree
import equinox as eqx
import numpyro
import numpyro.distributions as dist
import numpyro.handlers
from numpyro.infer.util import initialize_model
import optax
import numpy as np
import pandas as pd
# This example uses ArviZ's InferenceData API from the 0.22 release.
import sys
from importlib import metadata
try:
    arviz_version = metadata.version("arviz")
except metadata.PackageNotFoundError:
    arviz_version = None
if arviz_version != "0.22.0":
    if "google.colab" not in sys.modules:
        raise ImportError("This notebook requires arviz==0.22.0; install it in the notebook environment.")
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "arviz==0.22.0"])
import arviz as az

jax.config.update("jax_enable_x64", True)
key = jr.PRNGKey(0)
pprint = eqx.tree_pprint

Population Uncertainty

Cars on a bumpy road

There is a bump on the road that causes cars to oscillate after hitting it. The nature of the oscillation depends on the car’s mass and suspension system. You’ve installed a camera on the highway that can capture snapshots of each car’s vertical displacement. You capture 20 snapshots per car before they drive out of the camera’s view. Suppose you want to infer the cars’ suspension dynamics parameters (with uncertainty).

First, we need a forward model for the vertical displacement \(x\) of a car. We’ll model this as a damped harmonic oscillator

\[\begin{split} \begin{aligned} &\ddot{x} + 2\zeta\omega\dot{x} + \omega^2 x = 0 \\ &x(0) = x_0 \\ &\dot{x}(0) = 0 \end{aligned} \end{split}\]

where \(\zeta\) is the damping ratio and \(\omega\) is the natural frequency. Let \(x(t; x_0, \zeta, \omega)\) be the vertical position of a car at time \(t\), which is obtained by solving the above ODE.

This ODE happens to have an analytic solution, which we plot next for three damping regimes.

Show code cell source

Hide code cell source

@jit
def damped_harmonic_oscillator(t, x0, v0, zeta, omega):
    """
    Computes the displacement x(t) of a damped harmonic oscillator.
    
    Parameters
    ----------
        t: Time (scalar or array).
        x0: Initial displacement.
        v0: Initial velocity.
        omega: Natural frequency (rad/s).
        zeta: Damping ratio.

    Returns:
        x: Displacement x(t) at time t.
    """
    kwargs = dict(x0=x0, v0=v0, omega=omega)
    index = jnp.where(zeta < 1.0, 0, jnp.where(jnp.isclose(zeta, 1.0), 1, 2))
    x = lax.switch(
        index,
        [partial(underdamped_solution, zeta=zeta, **kwargs), partial(critically_damped_solution, **kwargs), partial(overdamped_solution, zeta=zeta, **kwargs)],
        t
    )
    return x

def underdamped_solution(t, x0, v0, zeta, omega):
    omega_d = omega * jnp.sqrt(1 - zeta**2)  # Damped natural frequency
    A = x0
    B = (v0 + zeta * omega * x0) / omega_d  # From initial velocity
    x = jnp.exp(-zeta * omega * t) * (A * jnp.cos(omega_d * t) + B * jnp.sin(omega_d * t))
    return x

def critically_damped_solution(t, x0, v0, omega):
    A = x0
    B = v0 + omega * x0
    x = (A + B * t) * jnp.exp(-omega * t)
    return x

def overdamped_solution(t, x0, v0, zeta, omega):
    r1 = -omega * (zeta - jnp.sqrt(zeta**2 - 1))
    r2 = -omega * (zeta + jnp.sqrt(zeta**2 - 1))
    A = (v0 - r2 * x0) / (r1 - r2)
    B = (r1 * x0 - v0) / (r1 - r2)
    x = A * jnp.exp(r1 * t) + B * jnp.exp(r2 * t)
    return x

# Parameters
x0 = 1.0       # Initial displacement
v0 = 0.0       # Initial velocity
omega = 3.0  # Natural frequency (rad/s)
zetas = [0.1, 1.0, 10.0]  # Damping ratio

# Time array
t = jnp.linspace(0, 10, 500)

# Compute displacement and velocity
xs = [damped_harmonic_oscillator(t, x0, v0, zeta_i, omega) for zeta_i in zetas]
x_underdamped, x_critically_damped, x_overdamped = xs

fig, ax = plt.subplots(figsize=FIGURE_SIZES["half_standard"])
ax.plot(t, x_underdamped, lw=1.5, color='0.10', linestyle='-', label='underdamped')
ax.plot(t, x_critically_damped, lw=1.5, color='0.35', linestyle='--', label='critically damped')
ax.plot(t, x_overdamped, lw=1.5, color='0.60', linestyle=':', label='overdamped')
ax.axhline(0, color="black", lw=1, ls='--', zorder=-10)
ax.set_ylabel("Displacement")
ax.set_xlabel("Time")
ax.legend()
finalize_axes(keep_box=False)

array([<Axes: xlabel='Time', ylabel='Displacement'>], dtype=object)

Hierarchical suspension model

Population distribution

Before seeing any road data, the population distribution describes uncertainty about the damping ratio \(\zeta\) and natural frequency \(\omega\) of a randomly selected car. We write it as

\[\begin{split} \underbrace{p(\zeta, \omega)}_{\substack{\text{population} \\ \text{distribution}}} = \int \underbrace{p(\zeta, \omega | \theta_\text{pop})}_{\substack{\text{conditional} \\ \text{prior}}} \underbrace{p(\theta_\text{pop})}_{\substack{\text{prior on} \\ \text{population} \\ \text{parameters}}} d\theta_\text{pop} \end{split}\]

where \(\theta_\text{pop}\) contains the population parameters. We choose the conditional prior

\[\begin{split} \begin{aligned} \log \zeta \mid \mu_\zeta, \tau_\zeta &\sim \text{Normal}(\mu_\zeta, \tau_\zeta^2) \\ \log \omega \mid \mu_\omega, \tau_\omega &\sim \text{Normal}(\mu_\omega, \tau_\omega^2) \end{aligned} \end{split}\]

where \(\operatorname{Normal}(m,s^2)\) denotes a normal distribution with mean \(m\) and standard deviation \(s\). Thus \(\mu_\zeta\) and \(\tau_\zeta\) are the population mean and standard deviation of \(\log \zeta\), with analogous definitions for \(\omega\). The population parameters are

\[ \theta_\text{pop}=(\mu_\zeta, \tau_\zeta, \mu_\omega, \tau_\omega). \]

We assign the hyperpriors

\[\begin{split} \begin{aligned} \mu_\zeta &\sim \text{Normal}(-2, 1^2) \\ \tau_\zeta &\sim \text{Exponential}(\text{rate}=10) \\ \mu_\omega &\sim \text{Normal}(0, 0.5^2) \\ \tau_\omega &\sim \text{Exponential}(\text{rate}=10) \end{aligned} \end{split}\]

This is the basic hierarchical pattern from the preceding section with \(\phi=\theta_\text{pop}\) and local parameter \(\theta_i=(\zeta_i,\omega_i)\). The following directed acyclic graph shows how the four shared population parameters govern the suspension parameters \((\zeta,\omega)\) of a single car:

Connecting the population distribution to the data

The example uses \(N_\mathrm{cars}=100\) cars and \(N_\mathrm{obs}=20\) displacement measurements per car. We assume that the initial displacement \(x_0\) is shared by all cars and lies between 0 and 5 centimeters:

\[ x_0 \sim \text{Uniform}([0, 5]) \]

The observed displacement of car \(i\) at time \(t_{ij}\) is

\[ y_{ij}| t_{ij}, \zeta_i, \omega_i, x_0 \sim \text{Normal}\Big(x(t_{ij}; x_0, \zeta_i, \omega_i), \sigma^2 \Big) \]

where the measurement noise \(\sigma\) is known.

The full graphical model expands the data node of the basic pattern into repeated measurements. The outer plate contains one pair \((\zeta_i,\omega_i)\) for each car \(i=1,\ldots,N_\mathrm{cars}\), and the nested inner plate contains that car’s measurements \(y_{ij}\) for \(j=1,\ldots,N_\mathrm{obs}\). The population parameters and \(x_0\) lie outside the plates because they are shared. The known inputs \(t_{ij}\) and the known measurement standard deviation \(\sigma\) are omitted from the diagram.

We can write down the posterior as

\[\begin{split} \begin{aligned} &p(\underbrace{\theta_\text{pop}, \boldsymbol{\zeta}, \boldsymbol{\omega}, x_0}_\text{unknowns} | \underbrace{\mathbf{t}, \mathbf{y}}_\text{data}) \\ &\propto \prod_{i=1}^{N_\mathrm{cars}} \Bigg\{ \prod_{j=1}^{N_\mathrm{obs}} \underbrace{p(y_{ij}|t_{ij}, \zeta_i, \omega_i, x_0, \sigma)}_\text{likelihood of an observation} \Bigg\} \\ &\qquad {}\times \prod_{i=1}^{N_\mathrm{cars}} \underbrace{p(\zeta_i|\mu_\zeta, \tau_\zeta) p(\omega_i|\mu_\omega, \tau_\omega)}_\text{conditional priors} \\ &\qquad {}\times \underbrace{p(\mu_\zeta) p(\tau_\zeta) p(\mu_\omega) p(\tau_\omega)}_{\substack{\text{prior on} \\ \text{population parameters}}} \underbrace{p(x_0)}_{\substack{\text{initial condition} \\ \text{prior}}}, \end{aligned} \end{split}\]

where \(\boldsymbol{\zeta}=(\zeta_1,\ldots,\zeta_{N_\mathrm{cars}})\) and \(\boldsymbol{\omega}=(\omega_1,\ldots,\omega_{N_\mathrm{cars}})\) collect the physical parameters of all cars in the data set. Finally, we’ll transform all random variables to a single random vector \(\xi\) which lives in unconstrained space \(\mathbb{R}^d\).

Building the model with NumPyro

We use NumPyro to construct both the log probability density \(\log p(\xi|\mathbf{t}, \mathbf{y})\) and the transformation \(\xi \mapsto (\theta_\text{pop}, \boldsymbol{\zeta}, \boldsymbol{\omega}, x_0)\). The nested numpyro.plate contexts implement the same repeated structure as the two plates in the graphical model: the outer context indexes cars, and the inner context indexes observations within each car. The companion notebook writes the model with NumPyro objects, generates a synthetic data set from a deliberately misspecified population, and samples the posterior with NUTS.

N_TIMES = 20
N_INDIVIDUALS = 100
MEASUREMENT_NOISE = 0.1
PARAMETERIZATION = 'centered'
times = jnp.linspace(0, 4, N_TIMES)

if PARAMETERIZATION == 'centered':

    def model(obs, gamma, prior_only=False):
        # Population parameters
        mu_zeta = numpyro.sample("mu_zeta", dist.Normal(-2.0, 1.0))
        tau_zeta = numpyro.sample("tau_zeta", dist.Exponential(10.0))
        mu_omega = numpyro.sample("mu_omega", dist.Normal(0.0, 0.5))
        tau_omega = numpyro.sample("tau_omega", dist.Exponential(10.0))

        # Initial condition
        x0 = numpyro.sample("x0", dist.Uniform(0, 5))

        # Physical parameters
        with numpyro.plate("individuals", N_INDIVIDUALS):
            log_zeta = numpyro.sample("log_zeta", dist.Normal(mu_zeta, tau_zeta))
            log_omega = numpyro.sample("log_omega", dist.Normal(mu_omega, tau_omega))
            zeta = jnp.exp(log_zeta)
            omega = jnp.exp(log_omega)

            if not prior_only:
                # Solve the ODE
                solver = lambda zeta, omega: damped_harmonic_oscillator(t=times, x0=x0, v0=0.0, zeta=zeta, omega=omega)
                x = vmap(solver, out_axes=-1)(zeta, omega)

                # Observations
                with numpyro.plate("observations", N_TIMES):
                    with numpyro.handlers.scale(scale=gamma):
                        y = numpyro.sample("y", dist.Normal(x, MEASUREMENT_NOISE), obs=obs)
        
        return locals()  # Returns a dict of all locally-defined variables

if PARAMETERIZATION == 'noncentered':

    def model(obs, gamma, prior_only=False):
        # Population parameters
        mu_zeta = numpyro.sample("mu_zeta", dist.Normal(-2.0, 1.0))
        tau_zeta = numpyro.sample("tau_zeta", dist.Exponential(10.0))
        mu_omega = numpyro.sample("mu_omega", dist.Normal(0.0, 0.5))
        tau_omega = numpyro.sample("tau_omega", dist.Exponential(10.0))

        # Initial condition
        x0 = numpyro.sample("x0", dist.Uniform(0, 5))

        # Physical parameters
        with numpyro.plate("individuals", N_INDIVIDUALS):
            log_zeta_noncentered = numpyro.sample("log_zeta_noncentered", dist.Normal())
            log_omega_noncentered = numpyro.sample("log_omega_noncentered", dist.Normal())
            log_zeta = mu_zeta + tau_zeta*log_zeta_noncentered
            log_omega = mu_omega + tau_omega*log_omega_noncentered
            zeta = jnp.exp(log_zeta)
            omega = jnp.exp(log_omega)

            if not prior_only:
                # Solve the ODE
                solver = lambda zeta, omega: damped_harmonic_oscillator(t=times, x0=x0, v0=0.0, zeta=zeta, omega=omega)
                x = vmap(solver, out_axes=-1)(zeta, omega)

                # Observations
                with numpyro.plate("observations", N_TIMES):
                    with numpyro.handlers.scale(scale=gamma):
                        y = numpyro.sample("y", dist.Normal(x, MEASUREMENT_NOISE), obs=obs)
        
        return locals()  # Returns a dict of all locally-defined variables

You can check that there are no syntax errors by sampling the model:

Show code cell source

Hide code cell source

# NOTE: This code cell is not needed - it is just useful for checking syntax/value errors in the model definition above.
#       Basically, we want this cell NOT to throw an error.

# The following `with` block applies an "effect handler". 
# It tells NumPyro to perform a task behind the scenes.
# Here, we set the NumPyro random seed to 1.
# Only statements inside the `with` block will be affected.
# NumPyro raises an error if we call `model` without applying the seed effect handler.
with numpyro.handlers.seed(rng_seed=1):

    # Sample the hierarchical model defined above by calling `model` with its arguments.
    dummy_y_obs = jnp.ones((N_TIMES, N_INDIVIDUALS))
    samples = model(dummy_y_obs, 1.0)

# Pretty-print the output
eqx.tree_pprint(samples)

Show code cell output

Hide code cell output

{
  'obs': f64[20,100],
  'gamma': 1.0,
  'prior_only': False,
  'mu_zeta': f64[],
  'tau_zeta': f64[],
  'mu_omega': f64[],
  'tau_omega': f64[],
  'log_zeta': f64[100],
  'log_omega': f64[100],
  'zeta': f64[100],
  'omega': f64[100],
  'solver': <function model.<locals>.<lambda>>,
  'x': f64[20,100],
  'y': f64[20,100],
  'x0': f64[]
}

The model is syntactically valid. We now generate a synthetic data set. The synthetic values of \((\log\zeta,\log\omega)\) follow a correlated, mildly nonlinear joint distribution, whereas the fitted hierarchy assumes conditional independence given \(\theta_\text{pop}\). This deliberate misspecification makes the example a test of approximation rather than an exact model-recovery exercise.

Show code cell source

Hide code cell source

# This cell generates a synthetic dataset for this toy problem.

def model_ground_truth(obs, gamma, prior_only=False):
    # Initial condition
    x0 = 3.5

    # Physical parameters
    with numpyro.plate("individuals", N_INDIVIDUALS):
        # Simulate samples from some "ground truth" population distribution
        p = numpyro.sample("p", dist.MultivariateNormal(jnp.array([-1.5, 1.0]), jnp.array([[0.1, 0.07], [0.07, 0.1]])))
        f = lambda x: x + 0.3*jnp.cos(2*(x - 1.0))
        zeta = jnp.exp(p[..., 0])
        omega = jnp.exp(f(p[..., 1]))

        if not prior_only:
            # Solve the ODE
            solver = lambda zeta, omega: damped_harmonic_oscillator(t=times, x0=x0, v0=0.0, zeta=zeta, omega=omega)
            x = vmap(solver, out_axes=-1)(zeta, omega)

            # Observations
            with numpyro.plate("observations", N_TIMES):
                with numpyro.handlers.scale(scale=gamma):
                    y = numpyro.sample("y", dist.Normal(x, MEASUREMENT_NOISE), obs=obs)
    
    return locals()  # Returns a dict of all locally-defined variables

# These will override the `numpyro.sample` statements in `model`
key, subkey = jr.split(key)
simulated_ground_truth = numpyro.infer.Predictive(model_ground_truth, num_samples=1)(subkey, None, 1.0)
y_obs = simulated_ground_truth['y'].squeeze(0)

fig, ax = plt.subplots(figsize=FIGURE_SIZES["half_standard"])
for i in range(N_INDIVIDUALS):
    ax.scatter(times, y_obs[:, i], 8, alpha=0.35, color='0.20', linewidths=0)
ax.set_xlabel("Time (s)")
ax.set_ylabel("Position (cm)")
finalize_axes(keep_box=False)

Show code cell output

Hide code cell output

array([<Axes: xlabel='Time (s)', ylabel='Position (cm)'>], dtype=object)

Next, we obtain the probability density and transformation functions from NumPyro using the BlackJAX interoperability pattern (Cabezas et al., 2024):

Show code cell source

Hide code cell source

model_default_args = (y_obs, 1.0, False)

key, subkey = jr.split(key)
(
    init_params,  # We don't need this
    potential_fn_gen, 
    postprocess_fn_gen, 
    model_trace  # We also don't need this
) = initialize_model(
    subkey,
    model,
    model_args=model_default_args,  # Dummy arguments
    dynamic_args=True,
)

# Get the probability density.
# This is p(ξ|y,t)
joint_log_prob_tempered = lambda x, gamma: -potential_fn_gen(y_obs, gamma, False)(x)
joint_log_prob = lambda x: joint_log_prob_tempered(x, 1.0)

# Get the transformation function.
# This is ξ ↦ (θ_pop, ζ, ω, x0)
constrain = lambda x: postprocess_fn_gen(y_obs, 1.0)(x)

# And get the inverse transformation function.
# This is (θ_pop, ζ, ω, x0) ↦ ξ 
unconstrain = jit(lambda x: numpyro.infer.util.unconstrain_fn(model, model_default_args, {}, x))

We now have the NumPyro quantities needed for sampling with BlackJAX. To demonstrate, here is how to evaluate \(p(\xi|\mathbf{t}, \mathbf{y})\) at some point \(\xi\):

# Create a dummy ξ
xi = {
    'mu_zeta': jnp.ones(()),
    'tau_zeta': jnp.ones(()),
    'mu_omega': jnp.ones(()),
    'tau_omega': jnp.ones(()),
    'log_zeta': jnp.ones((N_INDIVIDUALS,)),
    'log_omega': jnp.ones((N_INDIVIDUALS,)),
    'x0': jnp.ones(()),
}

joint_log_prob(xi)

Array(-334230.56120081, dtype=float64)

And here is how to transform xi to the original parameter ranges:

constrain(xi)

{'mu_zeta': Array(1., dtype=float64),
 'tau_zeta': Array(2.71828183, dtype=float64),
 'mu_omega': Array(1., dtype=float64),
 'tau_omega': Array(2.71828183, dtype=float64),
 'log_zeta': Array([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],      dtype=float64),
 'log_omega': Array([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],      dtype=float64),
 'x0': Array(3.65529289, dtype=float64)}

And unconstrain takes us back to unconstrained space:

eqx.tree_equal( unconstrain(constrain(xi)), xi )

Array(True, dtype=bool)

If the structure of xi is unclear, init_params.z gives the default unconstrained structure generated by NumPyro.

Sampling the hierarchical model posterior

We set up NUTS with BlackJAX for this problem. First, let’s pick starting points for each sampling chain by sampling from the prior:

NUM_CHAINS = 3

# Here is how to sample from the prior (in unconstrained space)
@partial(jit, static_argnums=1)
def sample_prior_xi(key, num_samples):
    s = numpyro.infer.Predictive(model, num_samples=num_samples)(key, *model_default_args)
    xi = vmap(unconstrain)(s)
    xi = {k: v for k, v in xi.items() if k in init_params.z.keys()}
    return xi

initial_xis = sample_prior_xi(key, 3)

# Print the shapes of `initial_xis`
eqx.tree_pprint(initial_xis)

{
  'log_omega': f64[3,100],
  'log_zeta': f64[3,100],
  'mu_omega': f64[3],
  'mu_zeta': f64[3],
  'tau_omega': f64[3],
  'tau_zeta': f64[3],
  'x0': f64[3]
}

The inference loop follows the BlackJAX change-of-variables pattern (Cabezas et al., 2024):

import blackjax

# @eqx.filter_jit
def inference_loop_multiple_chains(
    key, 
    initial_states, 
    sampler_params, 
    log_prob_fn, 
    num_samples, 
    num_chains, 
    likelihood_scale_schedule
):
    kernel = blackjax.nuts.build_kernel()

    @eqx.debug.assert_max_traces(max_traces=1)
    def step_fn(key, state, gamma, **params):
        return kernel(key, state, lambda x: log_prob_fn(x, gamma), **params)

    def one_step(states, fixed):
        key, gamma = fixed
        keys = jr.split(key, num_chains)
        states, infos = jax.vmap(partial(step_fn, gamma=gamma, **sampler_params))(keys, states)
        return states, (states, infos)

    keys = jr.split(key, num_samples)
    gammas = likelihood_scale_schedule(jnp.arange(num_samples))
    fixed = (keys, gammas)
    _, (states, infos) = lax.scan(one_step, initial_states, fixed)

    return (states, infos)

The loop permits likelihood tempering through the scale \(\gamma\). In this example we keep \(\gamma=1\) throughout, so both warmup and sampling target the full posterior; no annealing is applied.

Show code cell source

Hide code cell source

def full_posterior_schedule(steps):
    """Return unit likelihood scale at every transition."""
    return jnp.ones_like(steps, dtype=jnp.float64)

likelihood_scale_schedule = full_posterior_schedule

Finally, let’s run MCMC:

Show code cell source

Hide code cell source

# NUTS parameters
nuts_params = {
    'step_size': 0.001,
    'inverse_mass_matrix': jnp.ones(len(ravel_pytree(xi)[0]))
}

# Initialize the NUTS sampler states
nuts = blackjax.nuts(joint_log_prob, **nuts_params)
initial_states = vmap(nuts.init)(initial_xis)

# Split the key for warmup and sampling
key, warmup_key, sample_key = jr.split(key, 3)

# Warmup
num_warmup = 1000
warmup_states, warmup_infos = inference_loop_multiple_chains(
    warmup_key, initial_states, nuts_params, joint_log_prob_tempered, num_warmup, NUM_CHAINS, likelihood_scale_schedule
)

# Sample
num_samples = 1000
last_warmup_states = tree.map(lambda x: x[-1], warmup_states)
states, infos = inference_loop_multiple_chains(
    sample_key, last_warmup_states, nuts_params, joint_log_prob_tempered, num_samples, NUM_CHAINS, lambda x: jnp.ones_like(x)
)

# Put the samples in a dictionary of arrays whose leading dimensions are NUM_CHAINS and NUM_INDIVIDUALS.
xi_samples_all_chains = {k: v.swapaxes(0, 1) for k, v in states.position.items()}

The MCMC chains are stored in xi_samples_all_chains:

eqx.tree_pprint(xi_samples_all_chains)

{
  'log_omega': f64[3,1000,100],
  'log_zeta': f64[3,1000,100],
  'mu_omega': f64[3,1000],
  'mu_zeta': f64[3,1000],
  'tau_omega': f64[3,1000],
  'tau_zeta': f64[3,1000],
  'x0': f64[3,1000]
}

Here are the posterior sample histograms, trace plots, and R-hat convergence metric:

Show code cell source

Hide code cell source

# Visualize the posterior samples
samples_dataset_all_chains = az.from_dict(posterior=xi_samples_all_chains)
az.plot_trace(samples_dataset_all_chains, backend_kwargs={'tight_layout': True})

# Visualize rhat
def plot_rhats(samples):
    rhats = az.rhat(samples)
    rhats = np.hstack([rhats[k] for k in xi.keys()])
    fig, ax = plt.subplots(figsize=FIGURE_SIZES["half_standard"])
    ax.scatter(range(rhats.shape[0]), rhats, 9, facecolors='white', edgecolors='black', linewidths=0.6)
    ax.axhline(1.0, color="black", lw=1, ls='--', zorder=-10)
    ax.set_ylim(0, max(ax.get_ylim()[1], 1.5))
    ax.set_xlabel("Parameter index")
    ax.set_ylabel("R-hat")
    finalize_axes(keep_box=False)
    return ax

plot_rhats(samples_dataset_all_chains);

Show code cell output

Hide code cell output

Remove any chains that look like they didn’t converge:

Show code cell source

Hide code cell source

def remove_bad_chains(samples, bad_chain_ind, num_chains):
    """Splits samples into good and bad chains.
    
    Parameters
    ----------
    samples: dict
        Dictionary of samples.
    bad_chain_ind: list
        Indices of bad chains.
    num_chains: int
        Number of chains.

    Returns
    -------
    samples_good: dict
        Dictionary of samples from good chains.
    samples_bad: dict
        Dictionary of samples from bad chains.
    """
    is_good_chain = jnp.ones(num_chains, dtype=bool)
    if len(bad_chain_ind) > 0:
        is_good_chain = is_good_chain.at[jnp.array(bad_chain_ind)].set(False)
    is_bad_chain = ~is_good_chain
    return tree.map(lambda x: x[is_good_chain], samples), tree.map(lambda x: x[is_bad_chain], samples)

# NOTE: THIS CELL REQUIRES USER INPUT!
bad_chains = []  # Put the indices of any nonconvergent chains here to remove them. This will change run to run.

xi_samples, _ = remove_bad_chains(
    samples=xi_samples_all_chains, 
    bad_chain_ind=bad_chains, 
    num_chains=NUM_CHAINS
)

Show code cell source

Hide code cell source

if len(bad_chains) > 0:
    # Visualize the posterior samples
    samples_dataset = az.from_dict(posterior=xi_samples)
    az.plot_trace(samples_dataset, backend_kwargs={'tight_layout': True});

    # Visualize rhat
    plot_rhats(samples_dataset);

And let’s plot the epistemic and aleatoric uncertainty in the cars’ vertical position (as a function of time):

Show code cell source

Hide code cell source

# Concatenate the chains
# Each array in `xi_samples` begins with the dimensions (num_chains, num_samples).
# Each array in `xi_samples_combined` combines those two leading dimensions.
xi_samples_combined = tree.map(lambda x: x.reshape(x.shape[0] * x.shape[1], *x.shape[2:]), xi_samples)

# Transform the samples to the original parameter ranges
samples = vmap(constrain)(xi_samples_combined)

# Recenter the non-centered parameters (if applicable)
if PARAMETERIZATION == 'noncentered':
    recenter = lambda x, mu, tau: mu + tau*x
    samples['log_zeta'] = vmap(recenter)(samples['log_zeta_noncentered'], samples['mu_zeta'], samples['tau_zeta'])
    samples['log_omega'] = vmap(recenter)(samples['log_omega_noncentered'], samples['mu_omega'], samples['tau_omega'])

# Pick a dataset to plot
data_idx = 20
t_i, y_obs_i = times, y_obs[:, data_idx]

# Extract the samples for the chosen dataset
zeta_samples = jnp.exp(samples['log_zeta'][:, data_idx])
omega_samples = jnp.exp(samples['log_omega'][:, data_idx])
x0_samples = samples['x0']

# Propagate samples through the ODE
t_plt = jnp.linspace(0, 4.0, 200)
solver = lambda zeta, omega, x0: damped_harmonic_oscillator(t=t_plt, x0=x0, v0=0.0, zeta=zeta, omega=omega)
x_samples = vmap(solver)(zeta_samples, omega_samples, x0_samples)

# Simulate measurements
y_samples = x_samples + jr.normal(key, shape=x_samples.shape)*MEASUREMENT_NOISE

# Compute statistics
x05, x95 = jnp.quantile(x_samples, q=jnp.array([0.05, 0.95]), axis=0)
y05, y95 = jnp.quantile(y_samples, q=jnp.array([0.05, 0.95]), axis=0)

fig, ax = plt.subplots(figsize=FIGURE_SIZES["half_standard"])
ax.fill_between(t_plt, y05, y95, facecolor='0.92', edgecolor='0.35', lw=0.35, hatch='////', label='90% observation predictive interval')
ax.fill_between(t_plt, x05, x95, facecolor='0.60', edgecolor='0.20', lw=0.4, alpha=0.55, label='90% latent predictive interval')
ax.scatter(times, y_obs_i, s=8, alpha=0.8, color='k', label=f"Data", zorder=10)
ax.set_xlabel("Time")
ax.set_ylabel("Position")
ax.legend(loc='lower center', bbox_to_anchor=(0.5, 1.02), ncol=1, fontsize=7)
finalize_axes(keep_box=False)

array([<Axes: xlabel='Time', ylabel='Position'>], dtype=object)

Posterior predictive distribution for a new car

For each posterior draw of the population parameters, we draw one new pair \((\zeta,\omega)\). These samples approximate the posterior predictive distribution for a new car; propagating them through the oscillator gives the predictive displacements shown next.

The latent predictive interval combines uncertainty in the population parameters with variation between individual cars. Adding measurement noise gives the observation predictive interval. These nested predictive bands are not a decomposition into epistemic and aleatoric variances.

# Get the samples for the population parameters
mu_zeta = samples['mu_zeta']
tau_zeta = samples['tau_zeta']
mu_omega = samples['mu_omega']
tau_omega = samples['tau_omega']

# Sample the posterior predictive distribution for a new car
key, key_zeta, key_omega = jr.split(key, 3)
log_zeta_pop_samples = dist.Normal(mu_zeta, tau_zeta).rsample(key_zeta)
log_omega_pop_samples = dist.Normal(mu_omega, tau_omega).rsample(key_omega)

# Transform to physical space
zeta_pop_samples = jnp.exp(log_zeta_pop_samples)
omega_pop_samples = jnp.exp(log_omega_pop_samples)

Show code cell source

Hide code cell source

############################################################################################
# Histograms
############################################################################################

_df = pd.DataFrame({r'$\log(\zeta)$': log_zeta_pop_samples, r'$\log(\omega)$': log_omega_pop_samples})
g = sns.jointplot(data=_df, x=r'$\log(\zeta)$', y=r'$\log(\omega)$', kind='hist', fill=True, ratio=2, bins=30, cmap='Greys')

_df = pd.DataFrame({r'$\zeta$': zeta_pop_samples, r'$\omega$': omega_pop_samples})
g = sns.jointplot(data=_df, x=r'$\zeta$', y=r'$\omega$', kind='hist', fill=True, ratio=2, bins=30, cmap='Greys')


############################################################################################
# Time series plot
############################################################################################

# Get the initial condition samples
x0_samples = samples['x0']

# Propagate through the ODE
t_plt = jnp.linspace(0, 4.0, 200)
solver = lambda zeta, omega, x0: damped_harmonic_oscillator(t=t_plt, x0=x0, v0=0.0, zeta=zeta, omega=omega)
x_samples = vmap(solver)(zeta_pop_samples, omega_pop_samples, x0_samples)

# Simulate measurements
y_samples = x_samples + jr.normal(key, shape=x_samples.shape)*MEASUREMENT_NOISE

# Compute statistics
x05, x95 = jnp.quantile(x_samples, q=jnp.array([0.05, 0.95]), axis=0)
y05, y95 = jnp.quantile(y_samples, q=jnp.array([0.05, 0.95]), axis=0)

fig, ax = plt.subplots(figsize=FIGURE_SIZES["full_standard"])
ax.plot(t_plt, x_samples[0], alpha=0.8, lw=0.5, color='0.25', label=r'Posterior predictive samples')
ax.plot(t_plt, x_samples[1:20].T, alpha=0.8, lw=0.5, color='0.25')
ax.fill_between(t_plt, y05, y95, facecolor='0.92', edgecolor='0.35', lw=0.35, hatch='////', label='90% observation predictive interval')
ax.fill_between(t_plt, x05, x95, facecolor='0.60', edgecolor='0.20', lw=0.4, alpha=0.55, label='90% latent predictive interval')
ax.scatter(times, y_obs[:, 0], s=2, color='k', alpha=0.5, zorder=100, label='Data')
for i in range(1, N_INDIVIDUALS):
    ax.scatter(times, y_obs[:, i], s=2, color='k', alpha=0.5, zorder=100)
ax.set_xlabel("Time")
ax.set_ylabel("Position")
ax.legend(loc='lower center', bbox_to_anchor=(0.5, 1.02), ncol=1, fontsize=7)
finalize_axes(keep_box=False)

array([<Axes: xlabel='Time', ylabel='Position'>], dtype=object)

The posterior predictive samples cover the broad range of observed trajectories. This visual check is encouraging, but it is not evidence that the fitted population family is exact; the data are synthetic, and their generator was deliberately chosen outside that family.

Population-level predictions

Suppose there is another bump farther down the road, and a construction team will smooth it if more than 10% of cars cross the displacement threshold \(x=-3\) cm. We model the new bump’s initial displacement as \(x^\text{new}_0 \sim \mathcal{N}(5,1^2)\) cm, but no camera is available at that location. The fitted hierarchy allows us to propagate posterior uncertainty to this new setting.

We first visualize trajectories drawn from the posterior predictive distribution:

Show code cell source

Hide code cell source

# Get the initial condition samples
key, subkey = jr.split(key)
x0_samples = dist.Normal(5, 1).rsample(subkey, sample_shape=(zeta_pop_samples.shape[0],))

# Propagate through the ODE
t_plt = jnp.linspace(0, 4.0, 200)
solver = lambda zeta, omega, x0: damped_harmonic_oscillator(t=t_plt, x0=x0, v0=0.0, zeta=zeta, omega=omega)
x_samples = vmap(solver)(zeta_pop_samples, omega_pop_samples, x0_samples)

# Simulate measurements
y_samples = x_samples + jr.normal(key, shape=x_samples.shape)*MEASUREMENT_NOISE

# Compute statistics
x05, x95 = jnp.quantile(x_samples, q=jnp.array([0.05, 0.95]), axis=0)
y05, y95 = jnp.quantile(y_samples, q=jnp.array([0.05, 0.95]), axis=0)

# Check which population samples hit the threshold
hits_threshold = jnp.any(x_samples < -3, axis=1)

fig, ax = plt.subplots(figsize=FIGURE_SIZES["half_standard"])
for i in range(300):
    if hits_threshold[i]:
        ax.plot(t_plt, x_samples[i], alpha=0.24, lw=0.6, color='0.10', linestyle='-')
    else:
        ax.plot(t_plt, x_samples[i], alpha=0.35, lw=0.55, color='0.65', linestyle=(0, (3, 2)))
ax.plot([], [], color='0.65', lw=1.2, linestyle=(0, (3, 2)), label='Does not hit threshold')
ax.plot([], [], color='0.10', lw=1.2, linestyle='-', label='Hits threshold')
ax.axhline(y=-3, color='black', linestyle=':', lw=1.2, label='Threshold')
ax.set_xlabel("Time")
ax.set_ylabel("Position")
ax.legend()
finalize_axes(keep_box=False)

array([<Axes: xlabel='Time', ylabel='Position'>], dtype=object)

The following histogram shows the minimum position of each trajectory, \(\min_t\{x(t;x^\text{new}_0,\zeta,\omega)\}\), under the posterior predictive distribution:

Show code cell source

Hide code cell source

x_min = jnp.min(x_samples, axis=1)

fig, ax = plt.subplots(figsize=FIGURE_SIZES["half_standard"])
ax.hist(x_min, bins=40, color='0.75', edgecolor='0.15', linewidth=0.5)
ax.axvline(x=-3, color='black', linestyle=':', linewidth=1.2, label='Threshold')
ax.set_xlabel("Position")
ax.set_ylabel("Number of samples")
finalize_axes(keep_box=False)

array([<Axes: xlabel='Position', ylabel='Number of samples'>],
      dtype=object)

Finally, we estimate the posterior predictive probability that a car will cross the threshold \(x=-3\) cm. If this value is greater than 0.1, we will send a construction team to smooth out the bump.

Show code cell source

Hide code cell source

# Compute the probability that a car will hit the threshold
prob_hit_threshold = jnp.mean(hits_threshold)
print(f"Probability that a car will hit the threshold is {prob_hit_threshold:.2f}.")

Probability that a car will hit the threshold is 0.18.

Exercises

• Use a smaller dataset (\(N_\mathrm{cars}=10\); N_INDIVIDUALS in the companion notebook). Do we still get a good approximation of the population distribution?

• Use fewer time points (\(N_\mathrm{obs}=8\); N_TIMES in the companion notebook). Do the MCMC chains all converge to the same posterior distribution? Why or why not?

• Increase the measurement noise (\(\sigma=0.3\); MEASUREMENT_NOISE in the companion notebook). Do the MCMC chains all converge to the same posterior distribution? Why or why not?