pypomp.core.pomp.Pomp.train

Pomp.train(J: int, M: int, eta: LearningRate, key: Array | None = None, theta: PompParameters | None = None, optimizer: Optimizer = Adam(clip_norm=None, scale=False, ls=False, c=0.1, max_ls_itn=10, beta1=0.9, beta2=0.999, epsilon=1e-08), alpha: float = 0.97, thresh: float = 0.0, alpha_cooling: float = 1.0, n_monitors: int = 1, track_time: bool = True) → None[source]

Optimizes parameters for a continuous-state model using a differentiable particle filter and gradient-based methods.

This method performs Maximum Likelihood Estimation (MLE) using MOP, a differentiable particle filter for continuous-state POMPs. 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).

It bears repeating that this optimizer is only valid for continuous-state POMPs! For discrete-state models, use Pomp.mif() or Pomp.dpop_train().

This implementation leverages JAX to efficiently vectorize the algorithm across multiple initial parameter sets simultaneously. Results are automatically stored in the model’s history and can be accessed using Pomp.results().

Parameters:

  • J (int) – The number of particles in the MOP objective for obtaining the gradient and/or Hessian.

  • M (int) – Maximum iteration for the gradient descent optimization.

  • eta (LearningRate) – Learning rates per parameter as a LearningRate object.

  • key (jax.Array, optional) – The random key for reproducibility. Defaults to self.fresh_key.

  • theta (PompParameters, optional) – Parameters involved in the POMP model. Defaults to self.theta. Providing a PompParameters object with multiple parameter sets enables faster, vectorized execution across all parameter sets.

  • optimizer (Optimizer, optional) – The optimizer configuration object to use (e.g., pypomp.Adam(), pypomp.SGD(), pypomp.Newton(), pypomp.FullMatrixAdam(), etc.). Defaults to pypomp.Adam(). Hyperparameters like learning rate scaling, line search (scale, ls, c, max_ls_itn), gradient clipping (clip_norm), or Adam beta values are configured directly inside the optimizer instance.

  • alpha (float, optional) – Discount factor for MOP.

  • thresh (int, optional) – Threshold value to determine whether to resample particles.

  • alpha_cooling (float, optional) – Cooling factor for the MOP discount factor (alpha) using cosine decay. This factor represents the multiplier for the distance of alpha from 1.0 by the end of training (i.e., alpha approaches 1.0). Defaults to 1.0 (no cooling).

  • n_monitors (int, optional) – Number of particle filter runs to average for log-likelihood estimation.

  • track_time (bool, optional) – Boolean flag controlling whether to track the execution time.

Returns:

None. Updates Pomp.results_history with a PompTrainResult containing the log-likelihoods, parameter traces, and optimizer details from the training run.