Pymc_forecast — a new Bayesian time-series forecasting toolkit for PyMC

Sharing
1

Hi all,

we just released pymc_forecast, a toolkit for Bayesian time-series forecasting with PyMC. The philosophy: you write the generative model, the package handles everything around it — the train/forecast plumbing, inference, rolling-origin backtesting, and probabilistic scoring. No model zoo, no AutoML — a clean path from a hand-written PyMC model to labeled probabilistic forecasts.

It’s a PyMC port of @juanitorduz numpyro_forecast (itself a port of Pyro’s pyro.contrib.forecast), redesigned around PyMC idioms rather than translated 1:1.

The core trick: one model trains and forecasts

In-sample time latents live on variables dimmed "time"; the forecast horizon lives on separate {name}_future variables dimmed "time_future". Those future variables are absent from the fitted posterior — so pm.sample_posterior_predictive replays the posterior for everything it knows and draws the horizon from the model, conditioned on the replayed parents. One model body, no duplicated forecast code:

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

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:]

def model(h, covariates):
    # a per-step drift latent; time_series adds the matching `_future` latent
    drift = time_series(h, "drift", lambda name, dims: pm.Normal(name, 0.0, 0.5, dims=dims))
    sigma = pm.HalfNormal("sigma", 1.0)
    predict(
        h,
        lambda name, mu, dims, obs: pm.Normal(name, mu, sigma, dims=dims, observed=obs),
        pt.cumsum(drift),
    )

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"]                     # (chain, draw, time_future)

truth = test.to_xarray().rename({"index": "time_future"})
print(evaluate_forecast(forecast, truth))   # mae / rmse / crps / coverage

What’s in the box

  • Dims and coords everywhere. No positional axis conventions: variables carry named dims ("time", batch dims like "origin"), results are InferenceData/xarray with real datetime coordinates, and metrics align by dim name.
  • Three inference backends, one interface: Forecaster (ADVI/full-rank), HMCForecaster (NUTS, incl. nutpie/numpyro backends), PathfinderForecaster (pymc-extras).
  • pymc-extras statespace interop: structural time series / SARIMAX models are first-class citizens in the same forecast/backtest/metrics API, with the Kalman filter marginalizing the latent states instead of sampling them.
  • Backtesting: expanding/rolling-origin windows with per-fold refits and dim-aware metrics (CRPS, pinball, interval score, coverage, MASE).
  • Hierarchical models via batch dims: time_series(..., dims=("origin",)) gives every series its own latents (and matching _future variables) — see the 50-station BART example with partial pooling.

Docs & examples

All example notebooks are committed fully executed and re-run in CI, so they stay honest:

:open_book: Docs: pymc_forecast — pymc_forecast :laptop: Code: GitHub - pymc-labs/pymc_forecast: Port of numpyro_forecast · GitHub

pip install pymc-forecast          # core
pip install 'pymc-forecast[extras]'  # + pymc-extras (Pathfinder, statespace)

Status & feedback

This is an early 0.0.1 — the API is still settling, which makes this exactly the right moment to tell us what you need from a forecasting toolkit. What would make you use this over rolling your own forecast loop? Issues and PRs very welcome: Issues · pymc-labs/pymc_forecast · GitHub

Big thanks again to @juanitorduz for numpyro_forecast, which this builds on.