ODE Lotka-Volterra With Bayesian Inference in Multiple Ways

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
import pytensor
import pytensor.tensor as pt

from numba import njit
from pymc.ode import DifferentialEquation
from pytensor.compile.ops import as_op
from scipy.integrate import odeint
from scipy.optimize import least_squares

print(f"Running on PyMC v{pm.__version__}")

Running on PyMC v5.1.2+24.gf3ce16f26

%load_ext watermark
az.style.use("arviz-darkgrid")
rng = np.random.default_rng(1234)

Purpose

The purpose of this notebook is to demonstrate how to perform Bayesian inference on a system of ordinary differential equations (ODEs), both with and without gradients. The accuracy and efficiency of different samplers are compared.

We will first present the Lotka-Volterra predator-prey ODE model and example data. Next, we will solve the ODE using scipy.odeint and (non-Bayesian) least squares optimization. Next, we perform Bayesian inference in PyMC using non-gradient-based samplers. Finally, we use gradient-based samplers and compare results.

Key Conclusions

Based on the experiments in this notebook, the most simple and efficient method for performing Bayesian inference on the Lotka-Volterra equations was to specify the ODE system in Scipy, wrap the function as a Pytensor op, and use a Differential Evolution Metropolis (DEMetropolis) sampler in PyMC.

Background

Motivation

Ordinary differential equation models (ODEs) are used in a variety of science and engineering domains to model the time evolution of physical variables. A natural choice to estimate the values and uncertainty of model parameters given experimental data is Bayesian inference. However, ODEs can be challenging to specify and solve in the Bayesian setting, therefore, this notebook steps through multiple methods for solving an ODE inference problem using PyMC. The Lotka-Volterra model used in this example has often been used for benchmarking Bayesian inference methods (e.g., in this Stan case study, and in Chapter 16 of Statistical Rethinking [McElreath, 2018].

Lotka-Volterra Predator-Prey Model

The Lotka-Volterra model describes the interaction between a predator and prey species. This ODE given by:

\[\begin{split} \begin{aligned} \frac{d x}{dt} &=\alpha x -\beta xy \\ \frac{d y}{dt} &=-\gamma y + \delta xy \end{aligned} \end{split}\]

The state vector \(X(t)=[x(t),y(t)]\) comprises the densities of the prey and the predator species respectively. Parameters \(\boldsymbol{\theta}=[\alpha,\beta,\gamma,\delta, x(0),y(0)]\) are the unknowns that we wish to infer from experimental observations. \(x(0), y(0)\) are the initial values of the states needed to solve the ODE, and \(\alpha,\beta,\gamma\), and \(\delta\) are unknown model parameters which represent the following:

• \(\alpha\) is the growing rate of prey when there’s no predator.

• \(\beta\) is the dying rate of prey due to predation.

• \(\gamma\) is the dying rate of predator when there is no prey.

• \(\delta\) is the growing rate of predator in the presence of prey.

The Hudson’s Bay Company data

The Lotka-Volterra predator prey model has been used to successfully explain the dynamics of natural populations of predators and prey, such as the lynx and snowshoe hare data of the Hudson’s Bay Company. Since the dataset is small, we will hand-enter the values.

# fmt: off
data = pd.DataFrame(dict(
    year = np.arange(1900., 1921., 1),
    lynx = np.array([4.0, 6.1, 9.8, 35.2, 59.4, 41.7, 19.0, 13.0, 8.3, 9.1, 7.4,
                8.0, 12.3, 19.5, 45.7, 51.1, 29.7, 15.8, 9.7, 10.1, 8.6]),
    hare = np.array([30.0, 47.2, 70.2, 77.4, 36.3, 20.6, 18.1, 21.4, 22.0, 25.4, 
                 27.1, 40.3, 57.0, 76.6, 52.3, 19.5, 11.2, 7.6, 14.6, 16.2, 24.7])))
data.head()
# fmt: on

	year	lynx	hare
0	1900.0	4.0	30.0
1	1901.0	6.1	47.2
2	1902.0	9.8	70.2
3	1903.0	35.2	77.4
4	1904.0	59.4	36.3

# plot data function for reuse later
def plot_data(ax, lw=2, title="Hudson's Bay Company Data"):
    ax.plot(data.year, data.lynx, color="b", lw=lw, marker="o", markersize=12, label="Lynx (Data)")
    ax.plot(data.year, data.hare, color="g", lw=lw, marker="+", markersize=14, label="Hare (Data)")
    ax.legend(fontsize=14, loc="center left", bbox_to_anchor=(1, 0.5))
    ax.set_xlim([1900, 1920])
    ax.set_ylim(0)
    ax.set_xlabel("Year", fontsize=14)
    ax.set_ylabel("Pelts (Thousands)", fontsize=14)
    ax.set_xticks(data.year.astype(int))
    ax.set_xticklabels(ax.get_xticks(), rotation=45)
    ax.set_title(title, fontsize=16)
    return ax

_, ax = plt.subplots(figsize=(12, 4))
plot_data(ax);

Problem Statement

The purpose of this analysis is to estimate, with uncertainty, the parameters for the Lotka-Volterra model for the Hudson’s Bay Company data from 1900 to 1920.

Scipy odeint

Here, we make a Python function that represents the right-hand-side of the ODE equations with the call signature needed for the odeint function. Note that Scipy’s solve_ivp could also be used, but the older odeint function was faster in speed tests and is therefore used in this notebook.

# define the right hand side of the ODE equations in the Scipy odeint signature
from numba import njit


@njit
def rhs(X, t, theta):
    # unpack parameters
    x, y = X
    alpha, beta, gamma, delta, xt0, yt0 = theta
    # equations
    dx_dt = alpha * x - beta * x * y
    dy_dt = -gamma * y + delta * x * y
    return [dx_dt, dy_dt]

To get a feel for the model and make sure the equations are working correctly, let’s run the model once with reasonable values for \(\theta\) and plot the results.

# plot model function
def plot_model(
    ax,
    x_y,
    time=np.arange(1900, 1921, 0.01),
    alpha=1,
    lw=3,
    title="Hudson's Bay Company Data and\nExample Model Run",
):
    ax.plot(time, x_y[:, 1], color="b", alpha=alpha, lw=lw, label="Lynx (Model)")
    ax.plot(time, x_y[:, 0], color="g", alpha=alpha, lw=lw, label="Hare (Model)")
    ax.legend(fontsize=14, loc="center left", bbox_to_anchor=(1, 0.5))
    ax.set_title(title, fontsize=16)
    return ax

# note theta = alpha, beta, gamma, delta, xt0, yt0
theta = np.array([0.52, 0.026, 0.84, 0.026, 34.0, 5.9])
time = np.arange(1900, 1921, 0.01)

# call Scipy's odeint function
x_y = odeint(func=rhs, y0=theta[-2:], t=time, args=(theta,))

# plot
_, ax = plt.subplots(figsize=(12, 4))
plot_data(ax, lw=0)
plot_model(ax, x_y);

Looks like the odeint function is working as expected.

Least Squares Solution

Now, we can solve the ODE using least squares. Make a function that calculates the residual error.

# function that calculates residuals based on a given theta
def ode_model_resid(theta):
    return (
        data[["hare", "lynx"]] - odeint(func=rhs, y0=theta[-2:], t=data.year, args=(theta,))
    ).values.flatten()

Feed the residual error function to the Scipy least_squares solver.

# calculate least squares using the Scipy solver
results = least_squares(ode_model_resid, x0=theta)

# put the results in a dataframe for presentation and convenience
df = pd.DataFrame()
parameter_names = ["alpha", "beta", "gamma", "delta", "h0", "l0"]
df["Parameter"] = parameter_names
df["Least Squares Solution"] = results.x
df.round(2)

	Parameter	Least Squares Solution
0	alpha	0.48
1	beta	0.02
2	gamma	0.93
3	delta	0.03
4	h0	34.91
5	l0	3.86

Plot

time = np.arange(1900, 1921, 0.01)
theta = results.x
x_y = odeint(func=rhs, y0=theta[-2:], t=time, args=(theta,))
fig, ax = plt.subplots(figsize=(12, 4))
plot_data(ax, lw=0)
plot_model(ax, x_y, title="Least Squares Solution");

Looks right. If we didn’t care about uncertainty, then we would be done. But we do care about uncertainty, so let’s move on to Bayesian inference.

PyMC Model Specification for Gradient-Free Bayesian Inference

Like other Numpy or Scipy-based functions, the scipy.integrate.odeint function cannot be used directly in a PyMC model because PyMC needs to know the variable input and output types to compile. Therefore, we use a Pytensor wrapper to give the variable types to PyMC. Then the function can be used in PyMC in conjunction with gradient-free samplers.

Convert Python Function to a Pytensor Operator using @as_op decorator

We tell PyMC the input variable types and the output variable types using the @as_op decorator. odeint returns Numpy arrays, but we tell PyMC that they are Pytensor double float tensors for this purpose.

# decorator with input and output types a Pytensor double float tensors
@as_op(itypes=[pt.dvector], otypes=[pt.dmatrix])
def pytensor_forward_model_matrix(theta):
    return odeint(func=rhs, y0=theta[-2:], t=data.year, args=(theta,))

PyMC Model

Now, we can specify the PyMC model using the ode solver! For priors, we will use the results from the least squares calculation (results.x) to assign priors that start in the right range. These are empirically derived weakly informative priors. We also make them positive-only for this problem.

We will use a normal likelihood on untransformed data (i.e., not log transformed) to best fit the peaks of the data.

theta = results.x  # least squares solution used to inform the priors
with pm.Model() as model:
    # Priors
    alpha = pm.TruncatedNormal("alpha", mu=theta[0], sigma=0.1, lower=0, initval=theta[0])
    beta = pm.TruncatedNormal("beta", mu=theta[1], sigma=0.01, lower=0, initval=theta[1])
    gamma = pm.TruncatedNormal("gamma", mu=theta[2], sigma=0.1, lower=0, initval=theta[2])
    delta = pm.TruncatedNormal("delta", mu=theta[3], sigma=0.01, lower=0, initval=theta[3])
    xt0 = pm.TruncatedNormal("xto", mu=theta[4], sigma=1, lower=0, initval=theta[4])
    yt0 = pm.TruncatedNormal("yto", mu=theta[5], sigma=1, lower=0, initval=theta[5])
    sigma = pm.HalfNormal("sigma", 10)

    # Ode solution function
    ode_solution = pytensor_forward_model_matrix(
        pm.math.stack([alpha, beta, gamma, delta, xt0, yt0])
    )

    # Likelihood
    pm.Normal("Y_obs", mu=ode_solution, sigma=sigma, observed=data[["hare", "lynx"]].values)

pm.model_to_graphviz(model=model)

Plotting Functions

A couple of plotting functions that we will reuse below.

def plot_model_trace(ax, trace_df, row_idx, lw=1, alpha=0.2):
    cols = ["alpha", "beta", "gamma", "delta", "xto", "yto"]
    row = trace_df.iloc[row_idx, :][cols].values

    # alpha, beta, gamma, delta, Xt0, Yt0
    time = np.arange(1900, 1921, 0.01)
    theta = row
    x_y = odeint(func=rhs, y0=theta[-2:], t=time, args=(theta,))
    plot_model(ax, x_y, time=time, lw=lw, alpha=alpha);

def plot_inference(
    ax,
    trace,
    num_samples=25,
    title="Hudson's Bay Company Data and\nInference Model Runs",
    plot_model_kwargs=dict(lw=1, alpha=0.2),
):
    trace_df = az.extract(trace, num_samples=num_samples).to_dataframe()
    plot_data(ax, lw=0)
    for row_idx in range(num_samples):
        plot_model_trace(ax, trace_df, row_idx, **plot_model_kwargs)
    handles, labels = ax.get_legend_handles_labels()
    ax.legend(handles[:2], labels[:2], loc="center left", bbox_to_anchor=(1, 0.5))
    ax.set_title(title, fontsize=16)

Gradient-Free Sampler Options

Having good gradient free samplers can open up the models that can be fit within PyMC. There are five options for gradient-free samplers in PyMC that are applicable to this problem:

• Slice - the default gradient-free sampler

• DEMetropolisZ - a differential evolution Metropolis sampler that uses the past to inform sampling jumps

• DEMetropolis - a differential evolution Metropolis sampler

• Metropolis - the vanilla Metropolis sampler

• SMC - Sequential Monte Carlo

Let’s give them a shot.

A few notes on running these inferences. For each sampler, the number of tuning steps and draws have been reduced to run the inference in a reasonable amount of time (on the order of minutes). This is not a sufficient number of draws to get a good inferences, in some cases, but it works for demonstration purposes. In addition, multicore processing was not working for the Pytensor op function on all machines, so inference is performed on one core.

Slice Sampler

# Variable list to give to the sample step parameter
vars_list = list(model.values_to_rvs.keys())[:-1]

# Specify the sampler
sampler = "Slice Sampler"
tune = draws = 2000

# Inference!
with model:
    trace_slice = pm.sample(step=[pm.Slice(vars_list)], tune=tune, draws=draws)
trace = trace_slice
az.summary(trace)

Multiprocess sampling (4 chains in 4 jobs)
CompoundStep
>Slice: [alpha]
>Slice: [beta]
>Slice: [gamma]
>Slice: [delta]
>Slice: [xto]
>Slice: [yto]
>Slice: [sigma]

100.00% [16000/16000 02:00<00:00 Sampling 4 chains, 0 divergences]

Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 120 seconds.
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
alpha	0.478	0.025	0.433	0.526	0.002	0.002	115.0	254.0	1.04
beta	0.025	0.001	0.022	0.027	0.000	0.000	253.0	497.0	1.01
gamma	0.937	0.054	0.835	1.039	0.005	0.004	109.0	241.0	1.04
delta	0.028	0.002	0.025	0.031	0.000	0.000	109.0	242.0	1.05
xto	34.945	0.823	33.386	36.472	0.023	0.016	1269.0	2646.0	1.00
yto	3.837	0.476	2.958	4.730	0.036	0.026	169.0	491.0	1.03
sigma	4.111	0.487	3.263	5.038	0.007	0.005	5141.0	5579.0	1.00

az.plot_trace(trace, kind="rank_bars")
plt.suptitle(f"Trace Plot {sampler}");

fig, ax = plt.subplots(figsize=(12, 4))
plot_inference(ax, trace, title=f"Data and Inference Model Runs\n{sampler} Sampler");

Notes:
The Slice sampler was slow and resulted in a low effective sample size. Despite this, the results are starting to look reasonable!

DE MetropolisZ Sampler

sampler = "DEMetropolisZ"
tune = draws = 5000
with model:
    trace_DEMZ = pm.sample(step=[pm.DEMetropolisZ(vars_list)], tune=tune, draws=draws)
trace = trace_DEMZ
az.summary(trace)

Multiprocess sampling (4 chains in 4 jobs)
DEMetropolisZ: [alpha, beta, gamma, delta, xto, yto, sigma]

100.00% [40000/40000 00:16<00:00 Sampling 4 chains, 0 divergences]

Sampling 4 chains for 5_000 tune and 5_000 draw iterations (20_000 + 20_000 draws total) took 17 seconds.

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
alpha	0.482	0.024	0.434	0.523	0.001	0.001	747.0	1341.0	1.01
beta	0.025	0.001	0.022	0.028	0.000	0.000	821.0	1415.0	1.00
gamma	0.927	0.051	0.834	1.023	0.002	0.001	896.0	1547.0	1.01
delta	0.028	0.002	0.025	0.031	0.000	0.000	783.0	1432.0	1.01
xto	34.938	0.847	33.314	36.479	0.029	0.021	855.0	1201.0	1.00
yto	3.887	0.473	2.983	4.724	0.017	0.012	777.0	1156.0	1.01
sigma	4.129	0.477	3.266	5.029	0.017	0.012	799.0	1466.0	1.00

az.plot_trace(trace, kind="rank_bars")
plt.suptitle(f"Trace Plot {sampler}");

fig, ax = plt.subplots(figsize=(12, 4))
plot_inference(ax, trace, title=f"Data and Inference\n{sampler} Sampler")

Notes:
DEMetropolisZ sampled much quicker than the Slice sampler and therefore had a higher ESS per minute spent sampling. The parameter estimates are similar. A “final” inference would still need to beef up the number of samples.

DEMetropolis Sampler

In these experiments, DEMetropolis sampler was not accepting tune and requiring chains to be at least 8. We set draws at 5000, lower number like 3000 produce bad mixing.

sampler = "DEMetropolis"
chains = 8
draws = 6000
with model:
    trace_DEM = pm.sample(step=[pm.DEMetropolis(vars_list)], draws=draws, chains=chains)
trace = trace_DEM
az.summary(trace)

Population sampling (8 chains)
DEMetropolis: [alpha, beta, gamma, delta, xto, yto, sigma]
Attempting to parallelize chains to all cores. You can turn this off with `pm.sample(cores=1)`.

100.00% [8/8 00:00<00:00]

100.00% [7000/7000 00:39<00:00]

Sampling 8 chains for 1_000 tune and 6_000 draw iterations (8_000 + 48_000 draws total) took 40 seconds.

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
alpha	0.483	0.021	0.443	0.520	0.000	0.000	1820.0	2647.0	1.00
beta	0.025	0.001	0.023	0.027	0.000	0.000	1891.0	3225.0	1.00
gamma	0.924	0.045	0.837	1.008	0.001	0.001	1818.0	2877.0	1.00
delta	0.027	0.001	0.025	0.030	0.000	0.000	1628.0	2469.0	1.00
xto	34.890	0.707	33.523	36.176	0.018	0.013	1484.0	2862.0	1.01
yto	3.897	0.403	3.126	4.644	0.010	0.007	1756.0	2468.0	1.00
sigma	4.042	0.405	3.335	4.836	0.011	0.008	1437.0	2902.0	1.00

az.plot_trace(trace, kind="rank_bars")
plt.suptitle(f"Trace Plot {sampler}");

fig, ax = plt.subplots(figsize=(12, 4))
plot_inference(ax, trace, title=f"Data and Inference Model Runs\n{sampler} Sampler");

Notes:
KDEs looks too wiggly, but ESS is high R-hat is good and rank_plots also look good

Metropolis Sampler

sampler = "Metropolis"
tune = draws = 5000
with model:
    trace_M = pm.sample(step=[pm.Metropolis(vars_list)], tune=tune, draws=draws)
trace = trace_M
az.summary(trace)

Multiprocess sampling (4 chains in 4 jobs)
CompoundStep
>Metropolis: [alpha]
>Metropolis: [beta]
>Metropolis: [gamma]
>Metropolis: [delta]
>Metropolis: [xto]
>Metropolis: [yto]
>Metropolis: [sigma]

100.00% [40000/40000 01:46<00:00 Sampling 4 chains, 0 divergences]

/home/osvaldo/anaconda3/envs/pymc/lib/python3.10/site-packages/scipy/integrate/_odepack_py.py:248: ODEintWarning: Excess work done on this call (perhaps wrong Dfun type). Run with full_output = 1 to get quantitative information.
  warnings.warn(warning_msg, ODEintWarning)
/home/osvaldo/anaconda3/envs/pymc/lib/python3.10/site-packages/scipy/integrate/_odepack_py.py:248: ODEintWarning: Excess work done on this call (perhaps wrong Dfun type). Run with full_output = 1 to get quantitative information.
  warnings.warn(warning_msg, ODEintWarning)
Sampling 4 chains for 5_000 tune and 5_000 draw iterations (20_000 + 20_000 draws total) took 106 seconds.
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
alpha	0.481	0.024	0.437	0.523	0.004	0.003	44.0	112.0	1.10
beta	0.025	0.001	0.023	0.027	0.000	0.000	123.0	569.0	1.05
gamma	0.928	0.052	0.836	1.022	0.008	0.005	44.0	93.0	1.10
delta	0.028	0.002	0.025	0.031	0.000	0.000	47.0	113.0	1.09
xto	34.928	0.833	33.396	36.513	0.029	0.021	808.0	1128.0	1.00
yto	3.892	0.492	3.026	4.878	0.055	0.039	81.0	307.0	1.04
sigma	4.116	0.496	3.272	5.076	0.009	0.007	2870.0	3372.0	1.00

az.plot_trace(trace, kind="rank_bars")
plt.suptitle(f"Trace Plot {sampler}");

fig, ax = plt.subplots(figsize=(12, 4))
plot_inference(ax, trace, title=f"Data and Inference Model Runs\n{sampler} Sampler");

Notes:
The old-school Metropolis sampler is less reliable and slower than the DEMetroplis samplers. Not recommended.

SMC Sampler

The Sequential Monte Carlo (SMC) sampler can be used to sample a regular Bayesian model or to run model without a likelihood (Aproximate Bayesian Computation). Let’s try first with a regular model,

SMC with a Likelihood Function

sampler = "SMC with Likelihood"
draws = 2000
with model:
    trace_SMC_like = pm.sample_smc(draws)
trace = trace_SMC_like
az.summary(trace)

Initializing SMC sampler...
Sampling 4 chains in 4 jobs

100.00% [100/100 00:00<? Stage: 7 Beta: 1.000]

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
alpha	0.482	0.025	0.436	0.527	0.000	0.000	8093.0	7636.0	1.0
beta	0.025	0.001	0.022	0.027	0.000	0.000	8090.0	7582.0	1.0
gamma	0.927	0.053	0.826	1.023	0.001	0.000	8064.0	8142.0	1.0
delta	0.028	0.002	0.025	0.031	0.000	0.000	8028.0	8016.0	1.0
xto	34.893	0.843	33.324	36.500	0.009	0.007	8060.0	7716.0	1.0
yto	3.889	0.480	2.997	4.796	0.005	0.004	7773.0	7884.0	1.0
sigma	4.123	0.497	3.243	5.057	0.006	0.004	8169.0	7971.0	1.0

trace.sample_stats._t_sampling

64.09551501274109

az.plot_trace(trace, kind="rank_bars")
plt.suptitle(f"Trace Plot {sampler}");

fig, ax = plt.subplots(figsize=(12, 4))
plot_inference(ax, trace, title=f"Data and Inference Model Runs\n{sampler} Sampler");

Notes:
At this number of samples and tuning scheme, the SMC algorithm results in wider uncertainty bounds compared with the other samplers.

SMC Using pm.Simulator Epsilon=1

As outlined in the SMC tutorial on PyMC.io, the SMC sampler can be used for Aproximate Bayesian Computation, i.e. we can use a pm.Simulator instead of a explicit likelihood. Here is a rewrite of the PyMC - odeint model for SMC-ABC.

The simulator function needs to have the correct signature (e.g., accept an rng argument first).

# simulator function based on the signature rng, parameters, size.
def simulator_forward_model(rng, alpha, beta, gamma, delta, xt0, yt0, sigma, size=None):
    theta = alpha, beta, gamma, delta, xt0, yt0
    mu = odeint(func=rhs, y0=theta[-2:], t=data.year, args=(theta,))
    return rng.normal(mu, sigma)

Here is the model with the simulator function. Instead of a explicit likelihood function, the simulator uses distance metric (defaults to gaussian) between the simulated and observed values. When using a simulator we also need to specify epsilon, that is a tolerance value for the discrepancy between simulated and observed values. If epsilon is too low, SMC will not be able to move away from the initial values or a few values. We can easily see this with az.plot_trace. If epsilon is too high, the posterior will virtually be the prior. So

with pm.Model() as model:
    # Specify prior distributions for model parameters
    alpha = pm.TruncatedNormal("alpha", mu=theta[0], sigma=0.1, lower=0, initval=theta[0])
    beta = pm.TruncatedNormal("beta", mu=theta[1], sigma=0.01, lower=0, initval=theta[1])
    gamma = pm.TruncatedNormal("gamma", mu=theta[2], sigma=0.1, lower=0, initval=theta[2])
    delta = pm.TruncatedNormal("delta", mu=theta[3], sigma=0.01, lower=0, initval=theta[3])
    xt0 = pm.TruncatedNormal("xto", mu=theta[4], sigma=1, lower=0, initval=theta[4])
    yt0 = pm.TruncatedNormal("yto", mu=theta[5], sigma=1, lower=0, initval=theta[5])
    sigma = pm.HalfNormal("sigma", 10)

    # ode_solution
    pm.Simulator(
        "Y_obs",
        simulator_forward_model,
        params=(alpha, beta, gamma, delta, xt0, yt0, sigma),
        epsilon=1,
        observed=data[["hare", "lynx"]].values,
    )

Inference. Note the progressbar was throwing an error so it is turned off.

sampler = "SMC_epsilon=1"
draws = 2000
with model:
    trace_SMC_e1 = pm.sample_smc(draws=draws, progressbar=False)
trace = trace_SMC_e1
az.summary(trace)

Initializing SMC sampler...
Sampling 4 chains in 4 jobs

The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
alpha	0.474	0.012	0.460	0.492	0.006	0.004	5.0	5.0	3.41
beta	0.024	0.000	0.024	0.025	0.000	0.000	5.0	4.0	4.01
gamma	0.946	0.023	0.918	0.986	0.011	0.008	4.0	4.0	3.43
delta	0.028	0.001	0.028	0.029	0.000	0.000	4.0	4.0	4.19
xto	34.734	0.582	33.747	35.194	0.289	0.221	4.0	4.0	7.21
yto	3.814	0.214	3.429	3.966	0.101	0.077	4.0	5.0	3.93
sigma	1.899	0.357	1.369	2.206	0.173	0.132	4.0	8000.0	4.65

az.plot_trace(trace, kind="rank_bars")
plt.suptitle(f"Trace Plot {sampler}");

/home/osvaldo/proyectos/00_BM/arviz/arviz/stats/density_utils.py:487: UserWarning: Your data appears to have a single value or no finite values
  warnings.warn("Your data appears to have a single value or no finite values")

fig, ax = plt.subplots(figsize=(12, 4))
plot_inference(ax, trace, title=f"Data and Inference Model Runs\n{sampler} Sampler");

Notes:
We can see that if epsilon is too low plot_trace will clearly show it.

SMC with Epsilon = 10

with pm.Model() as model:
    # Specify prior distributions for model parameters
    alpha = pm.TruncatedNormal("alpha", mu=theta[0], sigma=0.1, lower=0, initval=theta[0])
    beta = pm.TruncatedNormal("beta", mu=theta[1], sigma=0.01, lower=0, initval=theta[1])
    gamma = pm.TruncatedNormal("gamma", mu=theta[2], sigma=0.1, lower=0, initval=theta[2])
    delta = pm.TruncatedNormal("delta", mu=theta[3], sigma=0.01, lower=0, initval=theta[3])
    xt0 = pm.TruncatedNormal("xto", mu=theta[4], sigma=1, lower=0, initval=theta[4])
    yt0 = pm.TruncatedNormal("yto", mu=theta[5], sigma=1, lower=0, initval=theta[5])
    sigma = pm.HalfNormal("sigma", 10)

    # ode_solution
    pm.Simulator(
        "Y_obs",
        simulator_forward_model,
        params=(alpha, beta, gamma, delta, xt0, yt0, sigma),
        epsilon=10,
        observed=data[["hare", "lynx"]].values,
    )

sampler = "SMC epsilon=10"
draws = 2000
with model:
    trace_SMC_e10 = pm.sample_smc(draws=draws)
trace = trace_SMC_e10
az.summary(trace)

Initializing SMC sampler...
Sampling 4 chains in 4 jobs

100.00% [100/100 00:00<? Stage: 5 Beta: 1.000]

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
alpha	0.483	0.035	0.416	0.548	0.000	0.000	7612.0	7414.0	1.0
beta	0.025	0.003	0.020	0.030	0.000	0.000	7222.0	7768.0	1.0
gamma	0.927	0.072	0.795	1.063	0.001	0.001	7710.0	7361.0	1.0
delta	0.028	0.002	0.023	0.032	0.000	0.000	7782.0	7565.0	1.0
xto	34.888	0.965	33.145	36.781	0.011	0.008	7921.0	7521.0	1.0
yto	3.902	0.723	2.594	5.319	0.008	0.006	7993.0	7835.0	1.0
sigma	1.450	1.080	0.024	3.409	0.013	0.009	7490.0	7172.0	1.0

az.plot_trace(trace, kind="rank_bars")
plt.suptitle(f"Trace Plot {sampler}");

fig, ax = plt.subplots(figsize=(12, 4))
plot_inference(ax, trace, title=f"Data and Inference Model Runs\n{sampler} Sampler");

Notes:
Now that we set a larger value for epsilon we can see that the SMC sampler (plus simulator) provides good results. Choosing a value for epsilon will always involve some trial and error. So, what to do in practice? As epsilon is the scale of the distance function. If you don’t have any idea of how much error do you expected to get between simulated and observed values then a rule of thumb for picking an initial guess for epsilon is to use a number smaller than the standard deviation of the observed data, how much smaller maybe one order of magnitude or so.

Posterior Correlations

As an aside, it is worth pointing out that the posterior parameter space is a difficult geometry for sampling.

az.plot_pair(trace_DEM, figsize=(8, 6), scatter_kwargs=dict(alpha=0.01), marginals=True)
plt.suptitle("Pair Plot Showing Posterior Correlations", size=18);

The major observation here is that the posterior shape is pretty difficult for a sampler to handle, with positive correlations, negative correlations, crecent-shapes, and large variations in scale. This contributes to the slow sampling (in addition to the computational overhead in solving the ODE thousands of times). This is also fun to look at for understanding how the model parameters impact each other.

Bayesian Inference with Gradients

NUTS, the PyMC default sampler can only be used if gradients are supplied to the sampler. In this section, we will solve the system of ODEs within PyMC in two different ways that supply the sampler with gradients. The first is the built-in pymc.ode.DifferentialEquation solver, and the second is to forward simulate using pytensor.scan, which allows looping. Note that there may be other better and faster ways to perform Bayesian inference with ODEs using gradients, such as the sunode project, and diffrax, which relies on JAX.

PyMC ODE Module

Pymc.ode uses scipy.odeint under the hood to estimate a solution and then estimate the gradient through finite differences.

The pymc.ode API is similar to scipy.odeint. The right-hand-side equations are put in a function and written as if y and p are vectors, as follows. (Even when your model has one state and/or one parameter, you should explicitly write y[0] and/or p[0].)

def rhs_pymcode(y, t, p):
    dX_dt = p[0] * y[0] - p[1] * y[0] * y[1]
    dY_dt = -p[2] * y[1] + p[3] * y[0] * y[1]
    return [dX_dt, dY_dt]

DifferentialEquation takes as arguments:

• func: A function specifying the differential equation (i.e. \(f(\mathbf{y},t,\mathbf{p})\)),

• times: An array of times at which data was observed,

• n_states: The dimension of \(f(\mathbf{y},t,\mathbf{p})\) (number of output parameters),

• n_theta: The dimension of \(\mathbf{p}\) (number of input parameters),

• t0: Optional time to which the initial condition belongs,

as follows:

ode_model = DifferentialEquation(
    func=rhs_pymcode, times=data.year.values, n_states=2, n_theta=4, t0=data.year.values[0]
)

Once the ODE is specified, we can use it in our PyMC model.

Inference with NUTS

pymc.ode is quite slow, so for demonstration purposes, we will only draw a few samples.

with pm.Model() as model:
    # Priors
    alpha = pm.TruncatedNormal("alpha", mu=theta[0], sigma=0.1, lower=0, initval=theta[0])
    beta = pm.TruncatedNormal("beta", mu=theta[1], sigma=0.01, lower=0, initval=theta[1])
    gamma = pm.TruncatedNormal("gamma", mu=theta[2], sigma=0.1, lower=0, initval=theta[2])
    delta = pm.TruncatedNormal("delta", mu=theta[3], sigma=0.01, lower=0, initval=theta[3])
    xt0 = pm.TruncatedNormal("xto", mu=theta[4], sigma=1, lower=0, initval=theta[4])
    yt0 = pm.TruncatedNormal("yto", mu=theta[5], sigma=1, lower=0, initval=theta[5])
    sigma = pm.HalfNormal("sigma", 10)

    # ode_solution
    ode_solution = ode_model(y0=[xt0, yt0], theta=[alpha, beta, gamma, delta])

    # Likelihood
    pm.Normal("Y_obs", mu=ode_solution, sigma=sigma, observed=data[["hare", "lynx"]].values)

sampler = "NUTS PyMC ODE"
tune = draws = 15
with model:
    trace_pymc_ode = pm.sample(tune=tune, draws=draws)

Only 15 samples in chain.
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [alpha, beta, gamma, delta, xto, yto, sigma]

100.00% [120/120 00:59<00:00 Sampling 4 chains, 0 divergences]

 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.2340415876362D-14
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.2340415876362D-14

/home/osvaldo/anaconda3/envs/pymc/lib/python3.10/site-packages/scipy/integrate/_odepack_py.py:248: ODEintWarning: Excess work done on this call (perhaps wrong Dfun type). Run with full_output = 1 to get quantitative information.
  warnings.warn(warning_msg, ODEintWarning)

 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.7324477632756D-17
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.5527744901481D-17
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.1105548980296D-16
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.1105548980296D-16
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.1105548980296D-16
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.1105548980296D-15
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.1105548980296D-15
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.1105548980296D-15
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.1105548980296D-15
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.1105548980296D-14
 lsoda--  above warning has been issued i1 times.         it will not be issued again for this problem      in above message,  i1 =        10
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.5241493348134D-15
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.5241493348134D-15
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.1088694571877D-13
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.1088694571877D-13
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.1088694571877D-13
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.4463323525725D-13
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.3388355031231D-13
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.3388355031231D-13
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.3388355031231D-13
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.6776710062462D-13
 lsoda--  above warning has been issued i1 times.         it will not be issued again for this problem      in above message,  i1 =        10
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.4158457835953D-42
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.4158457835953D-42
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.1370912617246D-40
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.1370912617246D-40
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.1370912617246D-40
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.5948148309049D-40
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.4374718724784D-40
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.3628059832750D-40
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.3628059832750D-40
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.3628059832750D-40
 lsoda--  above warning has been issued i1 times.         it will not be issued again for this problem      in above message,  i1 =        10

/home/osvaldo/anaconda3/envs/pymc/lib/python3.10/site-packages/scipy/integrate/_odepack_py.py:248: ODEintWarning: Excess work done on this call (perhaps wrong Dfun type). Run with full_output = 1 to get quantitative information.
  warnings.warn(warning_msg, ODEintWarning)
/home/osvaldo/anaconda3/envs/pymc/lib/python3.10/site-packages/scipy/integrate/_odepack_py.py:248: ODEintWarning: Excess work done on this call (perhaps wrong Dfun type). Run with full_output = 1 to get quantitative information.
  warnings.warn(warning_msg, ODEintWarning)

 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.8146615298522D-82
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.8146615298522D-82
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.8146615298522D-78
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.8146615298522D-78
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.8146615298522D-78
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.8146615298522D-77
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.8146615298522D-77
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.8146615298522D-77
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.8146615298522D-77
 lsoda--  warning..internal t (=r1) and h (=r2) are       such that in the machine, t + h = t on the next step  
       (h = step size). solver will continue anyway      in above,  r1 =  0.1900000000000D+04   r2 =  0.7775771408140D-76
 lsoda--  above warning has been issued i1 times.         it will not be issued again for this problem      in above message,  i1 =        10

Sampling 4 chains for 15 tune and 15 draw iterations (60 + 60 draws total) took 60 seconds.
The number of samples is too small to check convergence reliably.

trace = trace_pymc_ode
az.summary(trace)

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
alpha	0.472	0.031	0.389	0.506	0.008	0.006	17.0	33.0	1.36
beta	0.026	0.003	0.022	0.032	0.001	0.001	12.0	40.0	1.54
gamma	0.959	0.080	0.868	1.151	0.025	0.018	11.0	33.0	1.59
delta	0.029	0.003	0.026	0.035	0.001	0.001	14.0	37.0	1.33
xto	34.907	0.852	33.526	36.300	0.099	0.071	98.0	43.0	1.21
yto	3.347	0.772	1.742	4.342	0.278	0.205	10.0	16.0	1.78
sigma	6.117	4.425	3.502	16.420	1.353	0.984	9.0	16.0	1.87

az.plot_trace(trace, kind="rank_bars")
plt.suptitle(f"Trace Plot {sampler}");

_, ax = plt.subplots(figsize=(12, 4))
plot_inference(ax, trace, title=f"Data and Inference Model Runs\n{sampler} Sampler");

Notes:
NUTS is starting to find to the correct posterior, but would need a whole lot more time to make a good inference.

Simulate with Pytensor Scan

Finally, we can write the system of ODEs as a forward simulation solver within PyMC. The way to write for-loops in PyMC is with pytensor.scan. Gradients are then supplied to the sampler via autodifferentiation.

First, we should test that the time steps are sufficiently small to get a reasonable estimate.

Check Time Steps

Create a function that accepts different numbers of time steps for testing. The function also demonstrates how pytensor.scan is used.

# Lotka-Volterra forward simulation model using scan
def lv_scan_simulation_model(theta, steps_year=100, years=21):
    # variables to control time steps
    n_steps = years * steps_year
    dt = 1 / steps_year

    # PyMC model
    with pm.Model() as model:
        # Priors (these are static for testing)
        alpha = theta[0]
        beta = theta[1]
        gamma = theta[2]
        delta = theta[3]
        xt0 = theta[4]
        yt0 = theta[5]

        # Lotka-Volterra calculation function
        ## Similar to the right-hand-side functions used earlier
        ## but with dt applied to the equations
        def ode_update_function(x, y, alpha, beta, gamma, delta):
            x_new = x + (alpha * x - beta * x * y) * dt
            y_new = y + (-gamma * y + delta * x * y) * dt
            return x_new, y_new

        # Pytensor scan looping function
        ## The function argument names are not intuitive in this context!
        result, updates = pytensor.scan(
            fn=ode_update_function,  # function
            outputs_info=[xt0, yt0],  # initial conditions
            non_sequences=[alpha, beta, gamma, delta],  # parameters
            n_steps=n_steps,  # number of loops
        )

        # Put the results together and track the result
        pm.Deterministic("result", pm.math.stack([result[0], result[1]], axis=1))

    return model

Run the simulation for various time steps and plot the results.

_, ax = plt.subplots(figsize=(12, 4))

steps_years = [12, 100, 1000, 10000]
for steps_year in steps_years:
    time = np.arange(1900, 1921, 1 / steps_year)
    model = lv_scan_simulation_model(theta, steps_year=steps_year)
    with model:
        prior = pm.sample_prior_predictive(1)
    ax.plot(time, prior.prior.result[0][0].values, label=str(steps_year) + " steps/year")
ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))
ax.set_title("Lotka-Volterra Forward Simulation Model with different step sizes");

Sampling: []
Sampling: []
Sampling: []
Sampling: []

Notice how the lower resolution simulations are less accurate over time. Based on this check, 100 time steps per year is sufficiently accurate. 12 steps per year has too much “numerical diffusion” over 20 years of simulation.

Inference Using NUTs

Now that we are OK with 100 time steps per year, we write the model with indexing to align the data with the results.

def lv_scan_inference_model(theta, steps_year=100, years=21):
    # variables to control time steps
    n_steps = years * steps_year
    dt = 1 / steps_year

    # variables to control indexing to get annual values
    segment = [True] + [False] * (steps_year - 1)
    boolist_idxs = []
    for _ in range(years):
        boolist_idxs += segment

    # PyMC model
    with pm.Model() as model:
        # Priors
        alpha = pm.TruncatedNormal("alpha", mu=theta[0], sigma=0.1, lower=0, initval=theta[0])
        beta = pm.TruncatedNormal("beta", mu=theta[1], sigma=0.01, lower=0, initval=theta[1])
        gamma = pm.TruncatedNormal("gamma", mu=theta[2], sigma=0.1, lower=0, initval=theta[2])
        delta = pm.TruncatedNormal("delta", mu=theta[3], sigma=0.01, lower=0, initval=theta[3])
        xt0 = pm.TruncatedNormal("xto", mu=theta[4], sigma=1, lower=0, initval=theta[4])
        yt0 = pm.TruncatedNormal("yto", mu=theta[5], sigma=1, lower=0, initval=theta[5])
        sigma = pm.HalfNormal("sigma", 10)

        # Lotka-Volterra calculation function
        def ode_update_function(x, y, alpha, beta, gamma, delta):
            x_new = x + (alpha * x - beta * x * y) * dt
            y_new = y + (-gamma * y + delta * x * y) * dt
            return x_new, y_new

        # Pytensor scan is a looping function
        result, updates = pytensor.scan(
            fn=ode_update_function,  # function
            outputs_info=[xt0, yt0],  # initial conditions
            non_sequences=[alpha, beta, gamma, delta],  # parameters
            n_steps=n_steps,
        )  # number of loops

        # Put the results together
        final_result = pm.math.stack([result[0], result[1]], axis=1)
        # Filter the results down to annual values
        annual_value = final_result[np.array(boolist_idxs), :]

        # Likelihood function
        pm.Normal("Y_obs", mu=annual_value, sigma=sigma, observed=data[["hare", "lynx"]].values)
    return model

This is also quite slow, so we will just pull a few samples for demonstration purposes.

steps_year = 100
model = lv_scan_inference_model(theta, steps_year=steps_year)
sampler = "NUTS Pytensor Scan"
tune = draws = 50
with model:
    trace_scan = pm.sample(tune=tune, draws=draws)

Only 50 samples in chain.
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [alpha, beta, gamma, delta, xto, yto, sigma]

100.00% [400/400 01:29<00:00 Sampling 4 chains, 0 divergences]

Sampling 4 chains for 50 tune and 50 draw iterations (200 + 200 draws total) took 89 seconds.
The number of samples is too small to check convergence reliably.

trace = trace_scan
az.summary(trace)

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
alpha	0.480	0.025	0.432	0.526	0.003	0.002	77.0	94.0	1.02
beta	0.025	0.001	0.023	0.027	0.000	0.000	147.0	155.0	1.03
gamma	0.933	0.054	0.832	1.030	0.007	0.005	70.0	80.0	1.04
delta	0.028	0.002	0.024	0.031	0.000	0.000	70.0	94.0	1.04
xto	34.877	0.764	33.232	36.118	0.046	0.032	265.0	110.0	1.04
yto	3.987	0.504	2.887	4.749	0.069	0.049	58.0	102.0	1.06
sigma	4.173	0.488	3.361	5.005	0.056	0.039	83.0	104.0	1.03

az.plot_trace(trace, kind="rank_bars")
plt.suptitle(f"Trace Plot {sampler}");

time = np.arange(1900, 1921, 0.01)
odeint(func=rhs, y0=theta[-2:], t=time, args=(theta,)).shape

(2100, 2)

_, ax = plt.subplots(figsize=(12, 4))
plot_inference(ax, trace, title=f"Data and Inference Model Runs\n{sampler} Sampler");

Notes:
The sampler is faster than the pymc.ode implementation, but still slower than scipy odeint combined with gradient-free inference methods.

Summary

Let’s compare inference results among these different methods. Recall that, in order to run this notebook in a reasonable amount of time, we have an insufficient number of samples for many inference methods. For a fair comparison, we would need to bump up the number of samples and run the notebook for longer. Regardless, let’s take a look.

# Make lists with variable for looping
var_names = [str(s).split("_")[0] for s in list(model.values_to_rvs.keys())[:-1]]
# Make lists with model results and model names for plotting
inference_results = [
    trace_slice,
    trace_DEMZ,
    trace_DEM,
    trace_M,
    trace_SMC_like,
    trace_SMC_e1,
    trace_SMC_e10,
    trace_pymc_ode,
    trace_scan,
]
model_names = [
    "Slice Sampler",
    "DEMetropolisZ",
    "DEMetropolis",
    "Metropolis",
    "SMC with Likelihood",
    "SMC e=1",
    "SMC e=10",
    "PyMC ODE NUTs",
    "Pytensor Scan NUTs",
]

# Loop through variable names
for var_name in var_names:
    axes = az.plot_forest(
        inference_results,
        model_names=model_names,
        var_names=var_name,
        kind="forestplot",
        legend=False,
        combined=True,
        figsize=(7, 3),
    )
    axes[0].set_title(f"Marginal Probability: {var_name}")
    # Clean up ytick labels
    ylabels = axes[0].get_yticklabels()
    new_ylabels = []
    for label in ylabels:
        txt = label.get_text()
        txt = txt.replace(": " + var_name, "")
        label.set_text(txt)
        new_ylabels.append(label)
    axes[0].set_yticklabels(new_ylabels)

    plt.show();

Notes:
If we ran the samplers for long enough to get good inferences, we would expect them to converge on the same posterior probability distributions. This is not necessarily true for Aproximate Bayssian Computation, unless we first ensure that the approximation too the likelihood is good enough. For instance SMCe=1 is providing a wrong result, we have been warning that this was most likely the case when we use plot_trace as a diagnostic. For SMC e=10, we see that posterior mean agrees with the other samplers, but the posterior is wider. This is expected with ABC methods. A smaller value of epsilon, maybe 5, should provide a posterior closer to the true one.

Key Conclusions

We performed Bayesian inference on a system of ODEs in 4 main ways:

• Scipy odeint wrapped in a Pytensor op and sampled with non-gradient-based samplers (comparing 5 different samplers).

• Scipy odeint wrapped in a pm.Simulator function and sampled with a non-likelihood-based sequential Monte Carlo (SMC) sampler.

• PyMC ode.DifferentialEquation sampled with NUTs.

• Forward simulation using pytensor.scan and sampled with NUTs.

The “winner” for this problem was the Scipy odeint solver with a differential evolution (DE) Metropolis sampler and SMC (for a model with a Likelihood) provide good results with SMC being somewhat slower (but also better diagnostics). The improved efficiency of the NUTS sampler did not make up for the inefficiency in using the slow ODE solvers with gradients. Both DEMetropolis and SMC enable the simplest workflow for a scientist with a working numeric model and the desire to perform Bayesian inference. Just wrapping the numeric model in a Pytensor op and plugging it into a PyMC model can get you a long way!

Authors

Organized and rewritten by Greg Brunkhorst from multiple legacy PyMC.io example notebooks by Sanmitra Ghosh, Demetri Pananos, and the PyMC Team (Approximate Bayesian Computation).

Osvaldo Martin added some clarification about SMC-ABC and minor fixes in Mar, 2023

References

1

Richard McElreath. Statistical rethinking: A Bayesian course with examples in R and Stan. Chapman and Hall/CRC, 2018.

Watermark

%watermark -n -u -v -iv -w

Last updated: Thu Mar 30 2023

Python implementation: CPython
Python version       : 3.10.10
IPython version      : 8.10.0

pytensor  : 2.10.1
pandas    : 1.5.3
matplotlib: 3.5.2
pymc      : 5.1.2+12.g67925df69
numpy     : 1.23.5
arviz     : 0.16.0.dev0

Watermark: 2.3.1

License notice

All the notebooks in this example gallery are provided under the MIT License which allows modification, and redistribution for any use provided the copyright and license notices are preserved.

Citing PyMC examples

To cite this notebook, use the DOI provided by Zenodo for the pymc-examples repository.

Important

Many notebooks are adapted from other sources: blogs, books… In such cases you should cite the original source as well.

Also remember to cite the relevant libraries used by your code.

Here is an citation template in bibtex:

@incollection{citekey,
  author    = "<notebook authors, see above>",
  title     = "<notebook title>",
  editor    = "PyMC Team",
  booktitle = "PyMC examples",
  doi       = "10.5281/zenodo.5654871"
}

which once rendered could look like: