aimz.ImpactModel.sample_prior_predictive#

ImpactModel.sample_prior_predictive(X, *, num_samples=1000, rng_key=None, return_sites=None, shard_axis='obs', batch_size=None, output_dir=None, progress=True, **kwargs)[source]#

Draw samples from the prior predictive distribution.

Results are written to disk in the Zarr format, with computing and file writing decoupled and executed concurrently.

Parameters:

X (ArrayLike) – Input data. The leading axis is the observation axis.
num_samples (int) – The number of samples to draw.
rng_key (Array | None) – A pseudo-random number generator key. By default, an internal key is used and split as needed.
return_sites (str | Iterable[str] | None) – Names of variables (sites) to return. If None, samples param_output and deterministic sites.
shard_axis (Literal['obs', 'draw']) – Multi-device sharding strategy. "obs" (default) shards the input across devices and replicates the drawn samples. "draw" shards the drawn samples across devices and replicates the input.
batch_size (int | None) – Size of each batch, taken from the input under shard_axis="obs" and from the draws under shard_axis="draw". Also used as the chunk size when storing results. If None, it is determined automatically from the input size and number of samples.
output_dir (str | Path | None) – The directory where the outputs will be saved. If the specified directory does not exist, it will be created automatically. If None, a model-owned temporary directory is used. A subdirectory is generated within this directory to store the outputs; its path is recorded in the returned tree’s artifact_path attribute (on the root and the group node). The temporary directory is removed by cleanup(), cleanup_models(), or when the model is garbage-collected. Pass an explicit output_dir to keep results beyond the model’s lifetime.
progress (bool) – Whether to display a progress bar.
**kwargs (object) – Additional arguments passed to the model.

Returns:

Prior predictive samples. Posterior samples are included if available.

Raises:

TypeError – If param_output is passed as an argument.
ValueError – If shard_axis is not "obs" or "draw".
NotImplementedError – If a return site’s axis-1 size does not match the input batch size (shard_axis="obs" only).

Return type:

xr.DataTree