This document showcases basic functions in pypomp. We build a simple partially observed Markov process (POMP) model from scratch, simulate data from it, estimate its likelihood with the particle filter, and estimate its parameters with iterated filtering 2 (IF2). Because the model is linear and Gaussian, its exact likelihood is available from the Kalman filter, so we can check the Monte Carlo answers from pypomp against the truth.
First, we import the necessary packages.
import jax
import jax.numpy as jnp
import numpy as np
import pandas as pd
import plotnine as p9
from scipy import optimize, stats
import pypomp as pp
from pypomp.types import *
print(f"pypomp version: {pp.__version__}")
print(f"JAX version: {jax.__version__}")pypomp version: 1.0.2
JAX version: 0.11.1We use a first-order autoregressive, or AR(1), process observed with noise. The latent state \(X_n\) evolves according to \[ X_n = \phi X_{n-1} + \sigma_X \varepsilon_n, \qquad \varepsilon_n \sim \mathcal{N}(0, 1), \] and at each time \(n = 1, \dots, N\) we observe \[ Y_n = X_n + \sigma_Y \eta_n, \qquad \eta_n \sim \mathcal{N}(0, 1), \] where all \(\varepsilon_n\) and \(\eta_n\) are independent. We require \(|\phi| < 1\) so that the process is stationary, and we draw the initial state \(X_0\) from the stationary distribution \(\mathcal{N}\big(0, \sigma_X^2 / (1 - \phi^2)\big)\). The model has three parameters: the autoregressive coefficient \(\phi\), the process noise standard deviation \(\sigma_X\), and the measurement noise standard deviation \(\sigma_Y\).
A POMP model in pypomp is specified by up to four functions, which we call the model mechanics:
rinit simulates the initial state \(X_0\).rproc simulates one step of the latent process, from \(X_{n-1}\) to \(X_n\).dmeas evaluates the log-density of an observation \(Y_n\) given the state \(X_n\).rmeas simulates an observation \(Y_n\) given the state \(X_n\).rinit, rproc, and rmeas return dictionaries keyed by the names of the state or observation variables, and dmeas returns a scalar log-density. Parameters, states, and observations are passed in as dictionaries as well, so the model code can refer to each quantity by name.
pypomp calls these functions internally, so it needs to know which argument is which. It matches arguments using the type hints in pypomp.types (such as ParamDict and RNGKey) or, for arguments without these hints, the default argument names (theta_, key, X_, Y_, covars, t, t0, and dt). We use both below. Every expected argument must be present, even if the function does not use it. Our model has no covariates and its dynamics do not depend on time, so covars, t, t0, and dt go unused here.
def rinit(theta_: ParamDict, key: RNGKey, covars: CovarDict, t0: InitialTimeFloat):
"""Draw X_0 from the stationary distribution of the AR(1) process."""
sd = theta_["sigma_x"] / jnp.sqrt(1 - theta_["phi"] ** 2)
return {"X": sd * jax.random.normal(key)}
def rproc(
X_: StateDict,
theta_: ParamDict,
key: RNGKey,
covars: CovarDict,
t: TimeFloat,
dt: StepSizeFloat,
):
"""Advance the latent AR(1) process by one step."""
X = theta_["phi"] * X_["X"] + theta_["sigma_x"] * jax.random.normal(key)
return {"X": X}
def dmeas(
Y_: ObservationDict,
X_: StateDict,
theta_: ParamDict,
covars: CovarDict,
t: TimeFloat,
):
"""Log-density of the observation given the latent state."""
return jax.scipy.stats.norm.logpdf(Y_["Y"], loc=X_["X"], scale=theta_["sigma_y"])
def rmeas(
X_: StateDict,
theta_: ParamDict,
key: RNGKey,
covars: CovarDict,
t: TimeFloat,
):
"""Simulate an observation given the latent state."""
return {"Y": X_["X"] + theta_["sigma_y"] * jax.random.normal(key)}Each function handles a single particle and a single parameter set. pypomp uses JAX to vectorize them over particles and parameter sets and to compile them, so there is no need to write loops.
The algorithms pypomp uses to search the parameter space assume that the parameters can take any real value. Our parameters are constrained: \(\phi\) lies in \((-1, 1)\) and the standard deviations are positive. We therefore write functions that map the parameters to an unconstrained “estimation scale” and back, and store them in a ParTrans object. Here, we use the inverse hyperbolic tangent for \(\phi\) and the logarithm for the standard deviations.
def to_est(theta: ParamDict) -> ParamDict:
return {
"phi": jnp.arctanh(theta["phi"]),
"sigma_x": jnp.log(theta["sigma_x"]),
"sigma_y": jnp.log(theta["sigma_y"]),
}
def from_est(theta: ParamDict) -> ParamDict:
return {
"phi": jnp.tanh(theta["phi"]),
"sigma_x": jnp.exp(theta["sigma_x"]),
"sigma_y": jnp.exp(theta["sigma_y"]),
}
ptrans = pp.ParTrans(to_est=to_est, from_est=from_est)We now construct the POMP model object with pp.Pomp. Parameters are supplied as a PompParameters object, which can hold one or more parameter sets. Observation times are given by the row index of the data frame of observations, and t0 sets the time of the initial state. nstep=1 tells pypomp to call rproc once per observation interval.
We do not have real data in this example, so we generate our own. We first construct the model with a placeholder data frame that has the desired observation times and column name.
theta_true = {"phi": 0.9, "sigma_x": 1.0, "sigma_y": 1.0}
N = 200
ys_placeholder = pd.DataFrame({"Y": np.zeros(N)}, index=np.arange(1.0, N + 1))
ar1_template = pp.Pomp(
ys=ys_placeholder,
theta=pp.PompParameters(theta_true),
statenames=["X"],
t0=0.0,
rinit=rinit,
rproc=rproc,
dmeas=dmeas,
rmeas=rmeas,
nstep=1,
par_trans=ptrans,
)JAX handles random numbers differently from NumPy. Instead of a global random state, every random function takes an explicit pseudorandom number generator (PRNG) key, and the same key always gives the same result. We create a key and split it whenever we need a new one.
Pomp.simulate() returns data frames of simulated states and observations.
key = jax.random.key(1)
key, subkey = jax.random.split(key)
X_sim, Y_sim = ar1_template.simulate(key=subkey)
Y_sim.head()| theta_idx | sim | time | Y | |
|---|---|---|---|---|
| 0 | 0 | 0 | 1.0 | 1.507848 |
| 1 | 0 | 0 | 2.0 | -0.490220 |
| 2 | 0 | 0 | 3.0 | 1.440536 |
| 3 | 0 | 0 | 4.0 | 1.415157 |
| 4 | 0 | 0 | 5.0 | -1.506160 |
(
p9.ggplot()
+ p9.geom_line(p9.aes(x="time", y="X", color='"latent state X"'), data=X_sim)
+ p9.geom_point(p9.aes(x="time", y="Y", color='"observation Y"'), data=Y_sim, size=1)
+ p9.scale_color_manual(values=["black", "#d95f02"])
+ p9.theme_minimal()
+ p9.labs(x="Time", y="Value", color="")
)
Calling simulate() again with the same key and as_pomp=True reproduces the same draw, but returns a new Pomp object with the simulated observations as its data. This is the model we fit for the rest of the tutorial.
ar1 = ar1_template.simulate(key=subkey, as_pomp=True)
np.allclose(ar1.ys["Y"], Y_sim["Y"])TrueThe likelihood of a POMP model is usually intractable, so pypomp estimates it with the particle filter, which is implemented as the Pomp method pfilter(). The argument J sets the number of particles and reps sets the number of independent replicate filters.
Unlike the R family of POMP packages, some Pomp methods, including pfilter(), mif(), and train(), store their results in the object instead of returning new objects. Results are appended to the list ar1.results_history, with one element per method call. Each element contains the outputs, such as log-likelihood estimates, and the inputs used for the call, so it is easy to keep track of how each result was produced. The method results() returns the results of one call as a tidy data frame (by default, the most recent call). The log-likelihood reported by pfilter() averages the replicates on the likelihood scale, and se is its Monte Carlo standard error.
Each method that takes a key also stores an unused child key under ar1.fresh_key, which later calls use by default when no key is given. Below, we give the first call a key and let the others use fresh_key.
key, subkey = jax.random.split(key)
ar1.pfilter(J=100, reps=10, key=subkey)
ar1.results()| theta_idx | logLik | se | phi | sigma_x | sigma_y | |
|---|---|---|---|---|---|---|
| 0 | 0 | -374.331915 | 0.483449 | 0.9 | 1.0 | 1.0 |
To check this estimate, we compute the exact log-likelihood with the Kalman filter, which is only a few lines of code for this model.
def kalman_loglik(theta: dict[str, float], y: np.ndarray) -> float:
"""Exact log-likelihood of the AR(1) plus noise model."""
phi, sigma_x, sigma_y = theta["phi"], theta["sigma_x"], theta["sigma_y"]
m, P = 0.0, sigma_x**2 / (1 - phi**2) # stationary distribution of X_0
loglik = 0.0
for y_n in y:
m, P = phi * m, phi**2 * P + sigma_x**2 # predict X_n
S = P + sigma_y**2 # predictive variance of Y_n
loglik += stats.norm.logpdf(y_n, m, np.sqrt(S))
K = P / S # Kalman gain
m, P = m + K * (y_n - m), (1 - K) * P # update given Y_n
return loglik
y = ar1.ys["Y"].to_numpy()
exact_loglik = kalman_loglik(theta_true, y)
exact_logliknp.float64(-373.2145225685712)The particle filter estimate becomes more precise as we increase the number of particles. The particle filter’s estimate of the likelihood is unbiased, but taking its logarithm introduces a small downward bias, which also shrinks as the number of particles grows.
for J in [100, 1_000, 10_000]:
ar1.pfilter(J=J, reps=10)
pd.DataFrame(
{
"J": [100, 1_000, 10_000],
"logLik": [ar1.results(i)["logLik"].item() for i in [-3, -2, -1]],
"se": [ar1.results(i)["se"].item() for i in [-3, -2, -1]],
"exact logLik": exact_loglik,
}
)| J | logLik | se | exact logLik | |
|---|---|---|---|---|
| 0 | 100 | -373.171349 | 1.257857 | -373.214523 |
| 1 | 1000 | -373.425199 | 0.099352 | -373.214523 |
| 2 | 10000 | -373.228027 | 0.055475 | -373.214523 |
We now pretend that we do not know the parameters and estimate them. We implement IF2 as the Pomp method mif(). IF2 adds random walk perturbations to the parameters and runs a particle filter on the perturbed model, gradually shrinking, or “cooling”, the perturbations over repeated iterations until the parameters settle near the maximum likelihood estimate (MLE).
Starting a search from several different parameter values is good practice, because it explores more of the parameter space and shows whether the searches agree. Pomp.sample_params() draws parameter sets from uniform distributions with bounds given as a dictionary of length-2 tuples and returns them in a PompParameters object. Methods such as mif() and pfilter() run on every parameter set in a PompParameters object at once, so there is no need to make a new Pomp object for each set. In fact, methods like mif() and pfilter() are optimized to run substantially faster when operating on multiple parameter sets at once instead of running them sequentially.
key, subkey = jax.random.split(key)
starts = pp.Pomp.sample_params(
{"phi": (0.0, 0.95), "sigma_x": (0.2, 3.0), "sigma_y": (0.2, 3.0)},
n=8,
key=subkey,
)The random walk standard deviations are set with an RWSigma object and apply on the estimation scale. Parameters listed in its optional init_names argument would be perturbed only at the initial time \(t_0\), which is useful for parameters that only affect the initial state. Our model has no such parameters, so we perturb all of them at every time step. geometric_cooling(a=0.5) gradually shrinks the perturbations, totaling to a factor of 0.5 every 50 iterations.
rw_sd = pp.RWSigma(
sigmas={"phi": 0.02, "sigma_x": 0.02, "sigma_y": 0.02}
).geometric_cooling(a=0.5)We pass the starting values through the theta argument, run M=100 iterations of IF2 with J=1000 particles, and then evaluate the resulting estimates with the particle filter. mif() also replaces ar1.theta with the parameter estimates from its last iteration, so the following pfilter() call evaluates these estimates by default.
ar1.mif(J=1_000, M=100, rw_sd=rw_sd, theta=starts)
ar1.pfilter(J=2_000, reps=10)
ar1.results()| theta_idx | logLik | se | phi | sigma_x | sigma_y | |
|---|---|---|---|---|---|---|
| 0 | 0 | -370.555998 | 0.102340 | 0.813221 | 1.038959 | 0.968086 |
| 1 | 1 | -370.626290 | 0.102300 | 0.824400 | 0.880811 | 1.065401 |
| 2 | 2 | -371.216794 | 0.123509 | 0.742447 | 1.017276 | 0.858695 |
| 3 | 3 | -370.140258 | 0.136056 | 0.768908 | 1.037989 | 0.925975 |
| 4 | 4 | -370.552617 | 0.120507 | 0.743149 | 0.995955 | 1.039094 |
| 5 | 5 | -370.553294 | 0.108907 | 0.789364 | 0.932984 | 1.057886 |
| 6 | 6 | -370.279510 | 0.125699 | 0.748558 | 1.105699 | 0.865771 |
| 7 | 7 | -370.526143 | 0.184539 | 0.778799 | 0.998189 | 0.931117 |
JAX’s GPU acceleration enables pypomp to easily handle particle counts (\(J\)) in the tens of thousands, many more IF2 iterations (\(M\)), more filter replicates (reps), and many more parallel parameter searches. This introductory example uses small values so that it runs quickly on a laptop CPU.
ar1.traces() returns the parameter values and log-likelihood estimates at each iteration of every search as a data frame, including the log-likelihoods from pfilter() calls. The log-likelihood is not computed before the first iteration, so iteration 0 has a missing value.
traces = ar1.traces()
traces.dropna(subset="logLik").tail()| theta_idx | iteration | method | logLik | se | phi | sigma_x | sigma_y | |
|---|---|---|---|---|---|---|---|---|
| 815 | 7 | 97 | mif | -370.873138 | NaN | 0.803632 | 1.034992 | 0.954330 |
| 816 | 7 | 98 | mif | -371.332825 | NaN | 0.824704 | 0.999449 | 0.922404 |
| 817 | 7 | 99 | mif | -372.033691 | NaN | 0.783048 | 1.050122 | 0.904608 |
| 818 | 7 | 100 | mif | -370.958527 | NaN | 0.778799 | 0.998189 | 0.931117 |
| 819 | 7 | 100 | pfilter | -370.526143 | 0.184539 | 0.778799 | 0.998189 | 0.931117 |
Since this model has an exact likelihood, we can also find the exact MLE by numerically maximizing the Kalman filter log-likelihood, working on the same estimation scale.
def neg_loglik(z):
theta = {"phi": np.tanh(z[0]), "sigma_x": np.exp(z[1]), "sigma_y": np.exp(z[2])}
return -kalman_loglik(theta, y)
opt = optimize.minimize(neg_loglik, x0=np.zeros(3), method="Nelder-Mead")
mle = {"phi": np.tanh(opt.x[0]), "sigma_x": np.exp(opt.x[1]), "sigma_y": np.exp(opt.x[2])}
mle_loglik = -opt.funThe plots below show the IF2 traces for each search. The dashed line marks the exact MLE, and the dotted line in the parameter plot marks the true parameter values. The log-likelihoods computed during IF2 are for the perturbed model, so they are noisier than, and not directly comparable to, those from pfilter().
(
p9.ggplot(
traces.query("method == 'mif' and iteration > 0"),
p9.aes(x="iteration", y="logLik", color="factor(theta_idx)"),
)
+ p9.geom_line()
+ p9.geom_hline(yintercept=mle_loglik, linetype="dashed")
+ p9.coord_cartesian(ylim=(mle_loglik - 25, mle_loglik + 5))
+ p9.theme_minimal()
+ p9.labs(x="Iteration", y="Log-likelihood", title="IF2 Convergence", color="Search")
)
param_names = list(theta_true)
traces_long = traces.query("method == 'mif'").melt(
id_vars=["iteration", "theta_idx"], value_vars=param_names
)
reference = pd.DataFrame(
{
"variable": param_names,
"Truth": [theta_true[k] for k in param_names],
"MLE": [mle[k] for k in param_names],
}
).melt(id_vars=["variable"], var_name="Reference", value_name="intercept")
(
p9.ggplot(traces_long, p9.aes(x="iteration", y="value"))
+ p9.geom_line(p9.aes(color="factor(theta_idx)"))
+ p9.geom_hline(
p9.aes(yintercept="intercept", linetype="Reference"),
data=reference,
)
+ p9.scale_linetype_manual(values={"Truth": "dotted", "MLE": "dashed"})
+ p9.facet_wrap("~variable", scales="free_y")
+ p9.theme_minimal()
+ p9.labs(
x="Iteration",
y="Value",
title="Parameter Traces",
color="Search",
linetype="Reference",
)
)
Finally, we compare the best IF2 estimate with the exact MLE and the true parameter values, evaluating the exact log-likelihood at each. The best IF2 estimate should be close to the exact MLE; any remaining gap reflects Monte Carlo error and the limited number of iterations. The gap between the MLE and the truth is statistical error, which would shrink with more data.
best = ar1.results().sort_values("logLik", ascending=False).iloc[0]
best_if2 = {k: best[k] for k in param_names}
pd.DataFrame(
[
{**theta_true, "exact logLik": kalman_loglik(theta_true, y)},
{**mle, "exact logLik": mle_loglik},
{**best_if2, "exact logLik": kalman_loglik(best_if2, y)},
],
index=["truth", "exact MLE", "best IF2 estimate"],
)| phi | sigma_x | sigma_y | exact logLik | |
|---|---|---|---|---|
| truth | 0.900000 | 1.000000 | 1.000000 | -373.214523 |
| exact MLE | 0.756467 | 1.080946 | 0.937785 | -370.046225 |
| best IF2 estimate | 0.768908 | 1.037989 | 0.925975 | -370.187874 |