aimz.ImpactModel.fit#

ImpactModel.fit(X, y=None, *, num_samples=1000, rng_key=None, progress=True, batch_size=None, epochs=1, shuffle=True, **kwargs)[source]#

Fit the impact model to the provided data using epoch-based training.

This method implements an epoch-based training loop, where the data is iterated over in minibatches for a specified number of epochs. Variational inference is performed by repeatedly updating the model parameters on each minibatch, and then posterior samples are drawn from the fitted model.

Parameters:

  • X (ArrayLike | ArrayLoader) – Input data. If array-like, the leading axis is the sample axis. Alternatively, a data loader that holds all array-like objects and handles batching internally.

  • y (ArrayLike | None) – Output data. The leading axis is the sample axis. Must be None if X is a data loader.

  • num_samples (int) – The number of posterior samples to draw.

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

  • batch_size (int | None) – The number of data points processed at each step of variational inference. If None, the entire dataset is used as a single batch in each epoch. Ignored if X is a data loader, in which case the data loader is expected to handle batching internally.

  • epochs (int) – The number of epochs for variational inference optimization.

  • shuffle (bool) – Whether to shuffle the data at each epoch. Ignored if X is a data loader.

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

Returns:

The fitted model instance, enabling method chaining.

Raises:

TypeError – If the inference method is MCMC.

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. It does not check whether the model or guide is written to support subsampling semantics (e.g., using NumPyro’s subsample() or similar constructs).