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]

This is a pure functional implementation of the optimization algorithm, intended for users who need to compose it within custom JAX loops or higher-order functions. For a more user-friendly (but impurely-functional) interface, see pypomp.core.pomp.Pomp.train().

This function performs Maximum Likelihood Estimation (MLE) by treating the particle filter as a differentiable computational graph. It computes gradients of the log-likelihood with respect to the parameters via reverse-mode automatic differentiation (using JAX), and updates the parameters using optimizers (e.g., Adam, SGD).

This implementation leverages JAX to efficiently vectorize the algorithm across multiple initial parameter sets simultaneously.

Parameters:

  • struct (PompStruct) – The compiled structural representation of the POMP model.

  • thetas_array (jax.Array) – Array of initial parameters. Shape (n_reps, n_params). Must be aligned with the canonical order of struct.param_names (e.g. prepared via align_params).

  • J (int) – Number of particles.

  • optimizer (Optimizer) – Optimizer configuration object.

  • M (int) – Number of iterations.

  • eta (jax.Array) – Learning rates array. Shape (M, n_params). Must be aligned with the canonical order of struct.param_names along the last axis.

  • alpha (float | jax.Array) – Alpha parameter.

  • keys (jax.Array) – Random keys. Shape (n_reps, …).

  • alpha_cooling (float) – Alpha cooling factor.

  • thresh (float) – Resampling threshold.

  • n_monitors (int) – Number of monitors.

Returns:

Negative logLik history: Shape (n_reps, M) Theta history: Shape (n_reps, M+1, n_params)

Return type:

tuple[jax.Array, jax.Array]

Note

To align and stack input parameter dictionaries/scalars into the correct canonical ordering required by these arrays, use pypomp.functional.align_params().