aimz.ImpactModel.fit_on_batch#

ImpactModel.fit_on_batch(X, y, *, num_steps=10000, num_samples=1000, rng_key=None, progress=True, **kwargs)[source]#

Fit the impact model to the provided batch of data.

This method behaves differently depending on the inference method specified at th initialization:

  • SVI

    Runs variational inference on the provided batch by invoking the run() method of the SVI instance from NumPyro to estimate the posterior distribution, then draws samples from it.

  • MCMC

    Runs posterior sampling by invoking the run() method of the MCMC instance from NumPyro.

Parameters:

  • X (ArrayLike) – Input data. The leading axis is the sample axis.

  • y (ArrayLike) – Output data. The leading axis is the sample axis.

  • num_steps (int) – Number of steps for variational inference optimization. Ignored if the inference method is MCMC.

  • num_samples (int) – The number of posterior samples to draw. Ignored if the inference method is MCMC.

  • rng_key (Array | None) – A pseudo-random number generator key. By default, an internal key is used and split as needed.

  • progress (bool) – Whether to display a progress bar. Ignored if the inference method is MCMC.

  • **kwargs (object) – Additional arguments passed to the model.

Returns:

The fitted model instance, enabling method chaining.

Return type:

Self

Note

This method continues training from the existing SVI state if available. To start training from scratch, create a new model instance.