pymc_forecast.prediction#

Predictive drivers: forecasting and in-sample posterior prediction.

Both drivers rebuild the model via build_model() (forecasting with extended covariates, in-sample with the observed window) and run pm.sample_posterior_predictive over a posterior. Posteriors are accepted in any of the shapes the fitting paths produce — an xarray.DataTree or arviz.InferenceData with a posterior group, or a bare posterior Dataset.

Prediction outputs are draw-level by contract: every variable in the predictions (out-of-sample) and posterior_predictive (in-sample) groups carries chain/draw dims with the full posterior-predictive samples — nothing is reduced to means or quantiles on the default path. prediction_samples() extracts that samples Dataset from any result shape the drivers produce.

pymc_forecast.prediction.forecast(model_fn, posterior, data, covariates, *, num_samples=None, var_names=None, batch_size=None, random_seed=None, progressbar=False)[source]#

Sample probabilistic forecasts over the covariate horizon.

Rebuilds the model with covariates extending data along "time"; the surplus covariate steps are the forecast horizon. The posterior is replayed for in-sample latents while *_future variables (absent from it) are drawn fresh.

Parameters:

  • model_fn (callable or pymc_forecast.model.ForecastingModel) – The model body ((covariates, data=None) -> None or a ForecastingModel).

  • posterior (xarray.Dataset, xarray.DataTree or arviz.InferenceData) – A posterior Dataset or any object with a posterior group (posterior_dataset()); thinned only if num_samples is given.

  • data (xarray.DataArray, pandas.Series, pandas.DataFrame or array_like) – Observed series over the training window, normalized with as_dataarray() ("time" first; 2-D input gets a "series" dim).

  • covariates (xarray.DataArray, pandas.Series, pandas.DataFrame or array_like) – Covariates on the same "time" coordinate as data, spanning the training window plus the forecast horizon (normalized with as_dataarray; 2-D input gets a "covariate" dim); the steps past data define the horizon. Required: for covariate-free models pass null_covariates() over the full time index.

  • num_samples (int, optional) – Thin posterior to this many draws first (see thin_draws()); None uses every draw.

  • var_names (sequence of str, optional) – Variables to record (a bare string is not accepted as a single name). Default: "forecast", all *_future latents, and — for models registered through predict() — the noise-free "mu_future" predictor plus "expected_observation_future" when supplied by the model. On very wide panels, restricting this to ["forecast"] also shrinks the result’s memory footprint.

  • batch_size (int, optional) – Maximum posterior draws (per chain) per predictive pass. When set, the posterior is processed in consecutive blocks of at most this many draws and the blocks are concatenated along draw — bounding the working memory of each pass on very wide panels (the port of upstream’s chunked prediction, juanitorduz/numpyro_forecast#65). Per-block seeds are derived from random_seed, so a batched run is deterministic given the seed but draws different (equally valid) noise than an unbatched run. None runs a single pass.

  • random_seed (int or numpy.random.Generator, optional) – Seed for thinning and predictive sampling.

  • progressbar (bool, default False) – Show the sampling progress bar.

Returns:

xarray.DataTree or arviz.InferenceData – Result of pm.sample_posterior_predictive (a DataTree with current PyMC/ArviZ, InferenceData with older releases), with a predictions group carrying time_future coords.

Raises:
pymc_forecast.prediction.posterior_dataset(posterior)[source]#

Extract the posterior group as a plain xarray.Dataset.

Accepts an xarray.DataTree / arviz.InferenceData (or any object with a posterior item or attribute; uses its posterior group) or a bare Dataset (returned unchanged).

Parameters:

posterior (xarray.Dataset, xarray.DataTree or arviz.InferenceData) – The posterior container.

Returns:

xarray.Dataset – The posterior samples.

Raises:

TypeError – If posterior is not a Dataset and has no posterior group.

pymc_forecast.prediction.predict_in_sample(model_fn, posterior, data, covariates=None, *, num_samples=None, batch_size=None, random_seed=None, progressbar=False)[source]#

Sample the in-sample posterior predictive of "obs" and registered predictors.

The in-sample counterpart of forecast(): the model is rebuilt over the observed window only (no forecast horizon) and the observed variable is resampled given replayed latents. For models registered through predict(), the noise-free "mu" predictor is recorded alongside "obs"; "expected_observation" is also recorded when supplied by the model.

Parameters:
Returns:

xarray.DataTree or arviz.InferenceData – Result of pm.sample_posterior_predictive (a DataTree with current PyMC/ArviZ, InferenceData with older releases), with a posterior_predictive group holding "obs", plus "mu" and "expected_observation" when the model registers them.

Raises:
pymc_forecast.prediction.prediction_samples(result)[source]#

Extract the draw-level predictive samples from a prediction result.

Accepts any result shape the predictive drivers produce — an xarray.DataTree / arviz.InferenceData with a predictions group (from forecast()) or a posterior_predictive group (from predict_in_sample()) — or a bare Dataset (returned unchanged), and returns the samples as an xarray.Dataset whose variables retain the full chain / draw dims. Point summaries are the caller’s choice, e.g. prediction_samples(result)["forecast"].mean(("chain", "draw")).

Parameters:

result (xarray.Dataset, xarray.DataTree or arviz.InferenceData) – A prediction result; predictions is looked up before posterior_predictive.

Returns:

xarray.Dataset – The draw-level predictive samples.

Raises:

TypeError – If result carries none of the predictive groups.

pymc_forecast.prediction.thin_draws(posterior, num_samples, random_seed=None)[source]#

Subsample a posterior to num_samples draws (chain-flattened).

Draws are selected uniformly without replacement from the flattened (chain, draw) axes (with replacement only if more draws are requested than exist). The result is a posterior Dataset with chain=1, directly consumable by pm.sample_posterior_predictive.

Parameters:
Returns:

xarray.Dataset – Posterior with dims (chain, draw, ...), chain coord [0] and draw coord 0 .. num_samples - 1.

Raises:

  • ValueError – If num_samples is not positive.

  • TypeError – If posterior has no posterior group.