pypomp.panel.panel.PanelPomp.train

PanelPomp.train(J: int, M: int, eta: dict[str, float] | float, chunk_size: int = 1, optimizer: str = 'Adam', alpha: float = 0.97, key: Array | None = None, theta: PanelParameters | dict[str, DataFrame | None] | list[dict[str, DataFrame | None]] | None = None, eta_cooling: float = 1.0, alpha_cooling: float = 1.0, clip_norm: float | None = None)

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 (dict[str, float] | float) – Learning rate(s). Can be a float for a global learning rate or a dictionary mapping parameter names to rates.

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

  • optimizer (str, optional) – Optimizer type. Supported: ‘Adam’, ‘SGD’, ‘FullMatrixAdam’.

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

  • eta_cooling (float, optional) – Cooling factor for the learning rate (eta) using cosine decay. This represents the factor by which the original learning rate is multiplied by the end of training. Defaults to 1.0 (no cooling).

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

  • clip_norm (float, optional) – Clips gradient to [-clip_norm, clip_norm]. If None, no clipping is applied.

Returns:

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

Return type:

None