PyMC-Forecast

Bayesian time-series forecasting with PyMC: you write the generative model; the package handles the train/forecast plumbing, inference, backtesting, and evaluation.

A PyMC port of numpyro_forecast (itself a port of Pyro's pyro.contrib.forecast) — redesigned around PyMC idioms rather than a 1:1 translation.

Status: early development. The design and roadmap live in PLAN.md and the issue tracker.

Documentation: https://pymc-labs.github.io/pymc-forecast/ — API reference and executed example notebooks (univariate, hierarchical, covariates, state-space, VAR, M5).

Quickstart

Write a model with one predict() call, fit it, and forecast:

import numpy as np, pandas as pd, pymc as pm, pytensor.tensor as pt
from pymc_forecast import Forecaster, ForecastingModel, backtest, evaluate_forecast

# a trending weekly series; hold out the last 8 weeks
dates = pd.date_range("2024-01-07", periods=60, freq="W")
y = pd.Series(np.cumsum(np.random.default_rng(0).normal(0.2, 1.0, 60)) + 10, index=dates)
train, test = y.iloc[:52], y.iloc[52:]


class LocalLevel(ForecastingModel):
    def model(self, covariates, data=None):
        # a per-step drift latent; innovations adds the matching `_future` latent
        drift = self.innovations("drift", pm.Normal.dist(0.0, 0.5))
        sigma = pm.HalfNormal("sigma", 1.0)
        self.predict(pm.Normal.dist(0.0, sigma), pt.cumsum(drift))  # local-linear trend


model = LocalLevel()
fc = Forecaster(model, train, num_steps=5_000, random_seed=0)  # ADVI
idata = fc.forecast(horizon=8, num_samples=500, random_seed=0)
forecast = idata["predictions"]["forecast"]  # dims: (chain, draw, time_future)
# Outputs stay draw-level (never reduced to means/quantiles); prediction_samples(idata)
# extracts the samples Dataset from a forecast or in-sample result alike.

# score against the held-out weeks (aligned by dim name, not axis position)
truth = test.to_xarray().rename({"index": "time_future"})
print(evaluate_forecast(forecast, truth))  # {'mae': ..., 'rmse': ..., 'crps': ..., 'coverage': ...}

# rolling-origin backtest over the whole series
results = backtest(
    y,
    None,
    model,
    min_train_window=48,
    test_window=4,
    stride=4,
    num_samples=200,
    forecaster_options={"num_steps": 3_000},
    random_seed=0,
)

Swap Forecaster for HMCForecaster (NUTS, with nuts_sampler="nutpie"/"numpyro"/...) or PathfinderForecaster (pymc-extras) — the fit/forecast interface is identical. For models with real covariates, pass full-horizon covariates to .forecast() instead of horizon=, or hand the post-training rows alone to future_covariates= for a forecast conditioned on them. Covariate-free models can also forecast over an exact — even irregular — later time index with future_index=; the horizon length is derived from it at forecast time. See ssoe for observation/error-driven recursions such as ARMA and Holt-Winters, markov_series for state-space latents, and predict_mvn for observation noise correlated across time.

Functional API

The same model can be written as a plain function: it builds the current Horizon itself and passes it to the module-level innovations / predict primitives that the ForecastingModel helpers wrap. The functional fitters return a FitResult; draw the posterior from it and forecast:

from pymc_forecast import (
    Horizon,
    draw_posterior,
    fit_vi,
    forecast,
    innovations,
    null_covariates,
    predict,
)


def local_level(covariates, data=None):
    h = Horizon.from_data(covariates, data)
    drift = innovations(h, "drift", pm.Normal.dist(0.0, 0.5))
    sigma = pm.HalfNormal("sigma", 1.0)
    predict(h, pm.Normal.dist(0.0, sigma), pt.cumsum(drift))


result = fit_vi(local_level, train, num_steps=5_000, random_seed=0)  # ADVI
posterior = draw_posterior(result, 500, random_seed=0)
idata = forecast(local_level, posterior, train, null_covariates(dates), random_seed=0)

Both forms build the same PyMC model. The forecaster classes are thin wrappers that call fit_vi / fit_mcmc / fit_pathfinder, and every forecaster, fitter, and backtest accepts either model form.

pymc-extras statespace structural models (level/trend, seasonality, SARIMAX, ...) are first-class citizens too: define one as a StatespaceModel and fit it with StatespaceForecaster — the same forecast (including exogenous-regression covariates), predict_in_sample, backtest, and metrics calls apply, with the Kalman filter marginalizing the latent states instead of sampling them (see docs/examples/scan_vs_statespace_local_level.ipynb). The pymc-extras integrations (PathfinderForecaster / fit_pathfinder, StatespaceForecaster) are an optional extra: install with pip install 'pymc-forecast[extras]'.

Priors stay user-injectable on the model object: pass pymc_extras.prior.Prior objects directly to innovations/predict, or declare named default_priors on a ForecastingModel/StatespaceModel subclass and let callers override any subset with Model(priors={...}) — nested hyper-priors are shared across the train/forecast split so replay semantics hold.

Design principles

• One model trains and forecasts. In-sample time latents are fitted; the forecast horizon uses separate {name}_future variables that pm.sample_posterior_predictive draws from the prior while replaying the posterior for everything else.
• Dims and coords everywhere. No positional axis conventions: variables carry named dims ("time", "time_future", "obs", batch dims), results are arviz.InferenceData / xarray objects with real coordinates, and metrics are dim-aware. The prediction output schema is documented, test-locked public API.
• Not AutoML. No model zoo, no automatic feature pipelines — a clean path from a hand-written PyMC model to probabilistic forecasts and scores.
• Leverage the ecosystem. ADVI/NUTS from PyMC core, Pathfinder and state-space models from pymc-extras, diagnostics and storage from ArviZ.

Development

Requires Python >= 3.11 and uv:

uv sync --all-extras
uv run pytest
uv run ruff check .

To build the documentation (the example notebooks are committed fully executed; CI re-executes them with reduced sampling settings):

uv sync --all-extras --group docs
uv run sphinx-build -b html docs docs/_build/html

License

Apache-2.0. Portions derived from numpyro_forecast (Apache-2.0).