pypomp.core.pomp.Pomp.train

Pomp.train(J: int, M: int, eta: LearningRate, key: Array | None = None, theta: Mapping[str, int | float | number | Array] | Sequence[Mapping[str, int | float | number | Array]] | pypomp.core.parameters.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: int = 0, alpha_cooling: float = 1.0, n_monitors: int = 1, track_time: bool = True) → None[source]

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

This method 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. Results are automatically stored in the model’s history and can be accessed using self.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 (ThetaInput, optional) – Parameters involved in the POMP model. Defaults to self.theta. Accepts: - A single dictionary: dict[str, Numeric] - A list of dictionaries: list[dict[str, Numeric]] - An existing PompParameters object Providing a list or PompParameters object enables faster, vectorized execution across all parameter sets.

  • optimizer (Optimizer, optional) – The optimizer configuration object to use (e.g., pp.Adam(), pp.SGD(), pp.Newton(), pp.FullMatrixAdam(), etc.). Defaults to pp.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 self.results_history with a PompTrainResult containing the log-likelihoods, parameter traces, and optimizer details from the training run.