Functional API

The pypomp.functional module provides a collection of pure, stateless JAX functions for model simulation and inference. While the object-oriented Pomp class is recommended for most users, the functional API is intended for advanced users who need to compose algorithms within custom JAX loops, scan, or higher-order functions.

Data Structs

To use the functional API, you must first export your model’s structural data and compiled functions into a PompStruct or PanelPompStruct. This can be done using the to_struct() or to_struct() methods.

PompStruct(ys, dt_array_extended, ...)

Lightweight immutable JAX PyTree containing a compiled POMP model.

PanelPompStruct(ys_per_unit, ...)

A lightweight, immutable JAX PyTree holding the static data and compiled simulator functions for a PanelPOMP model.

Core Algorithms

pfilter(struct, thetas_array, J, keys[, ...])

Run the bootstrap particle filter on a POMP model struct.

panel_pfilter(struct, thetas_array, J, keys)

Evaluate panel POMP log-likelihood via particle filtering.

mif(struct, thetas_array, sigmas_array, ...)

Run the Iterated Filtering 2 (IF2) algorithm on a POMP model struct.

panel_mif(struct, shared_array, unit_array, ...)

Estimate panel POMP parameters using Panel Iterated Filtering.

simulate(struct, thetas_array, nsim, keys[, ...])

Simulate latent states and observations from a POMP model struct.

pmcmc(struct, thetas_array, proposal, ...)

Functional PMCMC entry point.

abc(struct, thetas_array, proposal, dprior, ...)

Functional ABC-MCMC entry point.

Differentiable Particle Filtering

These functions are primarily used for gradient-based parameter estimation. mop and dpop are designed to be fully differentiable with respect to the model parameters.

train(struct, thetas_array, J, optimizer, M, ...)

Optimize parameters via a differentiable particle filter (MOP).

panel_train(struct, shared_array, ...[, ...])

Optimize panel POMP parameters via a differentiable particle filter (MOP).

mop(struct, thetas_array, J, alpha, keys)

MOP differentiable particle filter log-likelihood objective.

dpop(struct, thetas_array, J, alpha, ...)

DPOP differentiable particle filter log-likelihood objective.

Utilities

align_params(params, names[, axis])

Align and stack parameter arrays into the canonical ordering for a model struct.