pypomp.functional.panel_train¶

pypomp.functional.panel_train(struct: PanelPompStruct, shared_array: Array, unit_array: Array, J: int, optimizer: Optimizer, M: int, eta_shared: Array, eta_spec: Array, alpha: float, keys: Array, alpha_cooling: float, chunk_size: int = 1) → tuple[Array, Array, Array][source]¶

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

This function performs Maximum Likelihood Estimation (MLE) for Panel POMP models by treating the particle filter as a differentiable computational graph (Tan et al. 2024 [1]). It computes gradients of the log-likelihood with respect to parameters across units, and updates them using an optimizer (e.g. Adam, SGD).

A pure functional implementation of the optimization (gradient-descent) algorithm, intended for composition within custom JAX code.

Parameters:

struct (PanelPompStruct) – Compiled structural representation of the Panel POMP model.
shared_array (jax.Array) – Array of initial shared parameters of shape (n_reps, n_shared) on the natural scale.
unit_array (jax.Array) – Array of initial unit-specific parameters of shape (n_reps, U, n_spec) on the natural scale.
J (int) – Number of particles.
optimizer (Optimizer) – Optimizer configuration object (e.g. Adam, SGD, Newton).
M (int) – Number of iterations.
eta_shared (jax.Array) – Learning rates array for shared parameters of shape (M, n_shared), aligned with the canonical order of struct.shared_param_names along the last axis.
eta_spec (jax.Array) – Learning rates array for unit-specific parameters of shape (M, n_spec), aligned with the canonical order of struct.unit_param_names along the last axis.
alpha (float) – Discount factor for MOP updates.
keys (jax.Array) – Random keys of shape (n_reps, M, U, ...).
alpha_cooling (float) – Cooling factor for discount factor alpha.
chunk_size (int, optional) – Number of units to process per gradient step. Defaults to 1.

Returns:

logliks_history (jax.Array) – Average negative log-likelihood trace across iterations of shape (n_reps, M + 1).
shared_history_natural (jax.Array) – Shared parameter history trace of shape (n_reps, M + 1, n_shared) on the natural scale.
unit_history_natural (jax.Array) – Unit-specific parameter history trace of shape (n_reps, M + 1, U, n_spec) on the natural scale.

Notes

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