pypomp.functional.train¶

pypomp.functional.train(struct: PompStruct, thetas_array: Array, J: int, optimizer: Optimizer, M: int, eta: Array, alpha: float | Array, keys: Array, alpha_cooling: float = 1.0, thresh: float = 0.0, n_monitors: int = 1) → tuple[Array, Array][source]¶

Optimize parameters via a differentiable particle filter (MOP).

Performs Maximum Likelihood Estimation using the Measurement Off-Parameter (MOP) particle filter (Tan et al. 2024 [1]), treating the particle filter as a differentiable computation graph and applies gradient-based optimizers (e.g. Adam, SGD, Newton) via JAX reverse-mode automatic differentiation.

Pure-functional implementation intended for users who need to compose the algorithm within custom JAX loops or higher-order functions. For the standard interface, see pypomp.Pomp.train().

JAX vectorizes the computation across all starting parameter sets simultaneously.

Parameters:

struct (PompStruct) – Compiled structural representation of the POMP model. Obtain via to_struct().
thetas_array (jax.Array) – Initial parameter array of shape (n_reps, n_params) on the natural scale. Must be aligned with struct.param_names.
J (int) – Number of particles.
optimizer (Optimizer) – Optimizer configuration object (e.g. Adam, SGD).
M (int) – Maximum number of gradient steps.
eta (jax.Array) – Per-parameter learning rate array of shape (M, n_params). Must be aligned with struct.param_names along the last axis.
alpha (float or jax.Array) – MOP discount factor.
keys (jax.Array) – Random keys of shape (n_reps, ...).
alpha_cooling (float, optional) – Cosine cooling multiplier for alpha. Defaults to 1.0.
thresh (float, optional) – ESS-based resampling threshold. Defaults to 0.0.
n_monitors (int, optional) – Number of unperturbed filter runs for log-likelihood monitoring. Defaults to 1.

Returns:

Negative log-likelihood history of shape (n_reps, M).
Parameter trace history of shape (n_reps, M+1, n_params).

Return type:

tuple of (jax.Array, jax.Array)

Notes

To align and stack input parameter dictionaries into the correct canonical ordering, use pypomp.functional.align_params().