pypomp.panel.panel.PanelPomp.train

PanelPomp.train(J: int, M: int, eta: LearningRate, chunk_size: int = 1, 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, key: Array | None = None, theta: PanelParameters | dict[str, DataFrame | None] | list[dict[str, DataFrame | None]] | None = None, alpha_cooling: float = 1.0)

Estimate parameters using chunked gradient-descent optimization (SGD/Adam).

This method performs stochastic gradient descent (or Adam) iterations over the likelihood of the panel POMP. It operates by drawing particles for a subset of units (defined by chunk_size), calculating gradients for both shared and unit-specific parameters, and updating estimates.

Parameters:

  • J (int) – Number of particles per unit.

  • M (int) – Number of training iterations.

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

  • chunk_size (int, optional) – Number of units to process per gradient calculation step.

  • optimizer (Optimizer, optional) – The optimizer configuration object to use (e.g., pp.Adam(), pp.SGD(), pp.FullMatrixAdam(), etc.). Defaults to pp.Adam(). Hyperparameters like gradient clipping (clip_norm) or Adam beta values are configured directly inside the optimizer instance.

  • alpha (float, optional) – Learning rate decay factor per iteration.

  • key (jax.Array, optional) – JAX PRNG key. If None, uses the fresh_key attribute.

  • theta (PanelParameters, optional) – Initial parameter estimates. If None, uses the current theta attribute.

  • 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).

Returns:

Updates self.theta and appends result to self.results_history.

Return type:

None