Univariate forecasting with pymc_forecast

This notebook ports the upstream NumPyro univariate forecasting example to the PyMC API. We forecast weekly BART ridership with a random-walk local level, annual Fourier seasonality, and a Student-T likelihood, fit with variational inference and evaluated with the continuous ranked probability score (CRPS).

The data is derived from the complete hourly BART origin-destination panel used by the upstream example. All predictive results use arviz.InferenceData, and all time handling uses labeled xarray coordinates.

Prepare notebook

import logging
import os
from functools import partial

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pymc as pm
import pytensor.tensor as pt
import xarray as xr

from pymc_forecast import (
    Forecaster,
    ForecastingModel,
    Horizon,
    backtest,
    build_model,
    eval_coverage,
    eval_crps,
    evaluate_forecast,
    fourier_features,
    innovations,
    predict,
    results_to_dataframe,
)
from pymc_forecast.data import FUTURE_DIM, TIME_DIM
from pymc_forecast.datasets import load_bart_weekly

HDI_PROB = 0.94

az.style.use("arviz-darkgrid")
az.rcParams["stats.ci_kind"] = "hdi"
az.rcParams["stats.ci_prob"] = HDI_PROB
plt.rcParams["figure.figsize"] = [10, 6]
plt.rcParams["figure.dpi"] = 100
plt.rcParams["figure.facecolor"] = "white"

%config InlineBackend.figure_format = "retina"

logging.getLogger("pymc").setLevel(logging.ERROR)
logging.getLogger("pytensor").setLevel(logging.ERROR)

SEED = 42

# CI executes every example notebook end-to-end with reduced settings.
SMOKE_TEST = os.environ.get("PYMC_FORECAST_SMOKE_TEST", "0") == "1"
NUM_STEPS = 300 if SMOKE_TEST else 20_000
NUM_SAMPLES = 50 if SMOKE_TEST else 1_000
BACKTEST_SAMPLES = 50 if SMOKE_TEST else 500

Read data

We use the same data as the upstream notebook: total weekly BART ridership. The package downloads and verifies the complete hourly origin-destination dataset, then sums over station pairs and non-overlapping weeks before applying a log transform.

y = load_bart_weekly()
duration = y.sizes[TIME_DIM]
print("weekly observations:", duration)

fig, ax = plt.subplots()
ax.plot(y[TIME_DIM].values, y.values, color="black", lw=1)
ax.set(title="Total weekly BART ridership", xlabel="week", ylabel="log(# rides)")
plt.show()

weekly observations: 469

Train-test split

Hold out the final 52 weeks, matching the one-year forecasting horizon in the upstream example. The integer week coordinates are preserved in the predictive groups.

TEST_WINDOW = 52
T1 = duration - TEST_WINDOW

y_train = y.isel({TIME_DIM: slice(None, T1)})
y_test = y.isel({TIME_DIM: slice(T1, None)})
print("train:", y_train.sizes[TIME_DIM], "test:", y_test.sizes[TIME_DIM])

fig, ax = plt.subplots()
ax.plot(y_train[TIME_DIM].values, y_train.values, color="C0", label="train")
ax.plot(y_test[TIME_DIM].values, y_test.values, color="C1", label="test")
ax.axvline(y_test[TIME_DIM].values[0], color="gray", ls="--", label="train/test split")
ax.legend(loc="lower left")
ax.set(title="Train / test split", xlabel="week", ylabel="log(# rides)")
plt.show()

train: 417 test: 52

Seasonal features

The BART series is weekly, so the seasonal design uses a 365.25 / 7 week period and 26 Fourier harmonics, matching the upstream example. Feature and observation coordinates remain aligned by week.

NUM_TERMS = 26
covariates = fourier_features(duration, period=365.25 / 7, num_terms=NUM_TERMS)
covariates = covariates.rename({"fourier": "covariate"})
covariates_train = covariates.isel({TIME_DIM: slice(None, T1)})
print("covariates:", covariates.sizes)

fig, ax = plt.subplots()
for name in covariates["covariate"].values[:3]:
    ax.plot(covariates[TIME_DIM].values, covariates.sel(covariate=name).values, label=name)
ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))
ax.set(title="First annual Fourier modes", xlabel="week")
plt.show()

covariates: Frozen({'time': 469, 'covariate': 52})

Model specification

The mean has three pieces: a global intercept, a random-walk local level, and a linear regression on Fourier features. The forecast horizon is represented by separate *_future variables through innovations, so fitting and forecasting use one model body. The random walk is non-centered for stable PyMC variational inference while retaining the upstream model’s priors.

pymc_forecast offers two equivalent ways to write a model:

• Object-oriented: subclass ForecastingModel and implement model(self, covariates, data=None). The bound helpers self.innovations and self.predict thread the current horizon (available as self.horizon), and priors can be injected through default_priors or the priors= constructor argument.

• Functional: a plain function (covariates, data=None) that builds h = Horizon.from_data(covariates, data) and calls innovations(h, ...) and predict(h, ...).

Both build the same PyMC model, and every forecaster, fitter and backtest accepts either form. This notebook uses the object-oriented form; the functional twin is at the end.

class LocalLevelSeasonal(ForecastingModel):
    """Local level + Fourier regression with Student-T observations."""

    def model(self, covariates, data=None) -> None:
        bias = pm.Normal("bias", 0.0, 10.0)
        weight = pm.Normal("weight", 0.0, 0.1, dims="covariate")
        drift_scale = pm.LogNormal("drift_scale", -20.0, 5.0, initval=0.01)
        nu = pm.Gamma("nu", alpha=10.0, beta=2.0)
        sigma = pm.LogNormal("sigma", -5.0, 5.0, initval=0.05)

        drift_raw = self.innovations("drift_raw", pm.Normal.dist(0.0, 1.0))
        level = pt.cumsum(drift_raw * drift_scale)
        regression = pt.dot(covariates.values, weight)
        prediction = bias + level + regression

        self.predict(pm.StudentT.dist(nu=nu, sigma=sigma), prediction)


model = LocalLevelSeasonal()
train_model = build_model(model, y_train, covariates_train)
train_model

\[\begin{split} \begin{array}{rcl} \text{bias} &\sim & \operatorname{Normal}(0,~10)\\\text{weight} &\sim & \operatorname{Normal}(0,~0.1)\\\text{drift\_scale} &\sim & \operatorname{LogNormal}(-20,~5)\\\text{nu} &\sim & \operatorname{Gamma}(10,~\text{<constant>})\\\text{sigma} &\sim & \operatorname{LogNormal}(-5,~5)\\\text{drift\_raw} &\sim & \operatorname{Normal}(0,~1)\\\text{mu} &= &\operatorname{Deterministic}(f(\text{bias},~\text{weight},~\text{drift\_raw},~\text{drift\_scale}))\\\text{obs} &\sim & \operatorname{StudentT}(\text{nu},~f(\text{bias},~\text{weight},~\text{drift\_raw},~\text{drift\_scale}),~\text{sigma}) \end{array} \end{split}\]

Plot helper

def attach_constant_data(idata, *, name: str, values: np.ndarray, dim: str) -> None:
    """Add an ArviZ constant-data group used by plot_lm."""
    idata["constant_data"] = xr.Dataset(
        {name: (dim, np.asarray(values, dtype=float))},
        coords={dim: idata["observed_data"][dim].values},
    )


def lm_idata_from_forecast(
    samples: xr.DataArray,
    truth: xr.DataArray,
    *,
    x_name: str,
    x_values: np.ndarray,
):
    """Represent held-out forecasts as an ArviZ posterior predictive container."""
    return az.from_dict(
        {
            "posterior_predictive": {"obs": samples.values},
            "observed_data": {"obs": truth.values},
            "constant_data": {x_name: np.asarray(x_values, dtype=float)},
        },
        coords={TIME_DIM: truth[FUTURE_DIM].values},
        dims={"obs": [TIME_DIM], x_name: [TIME_DIM]},
    )

Prior predictive checks

prior_model = build_model(model, y_train, covariates_train)
prior_idata = pm.sample_prior_predictive(
    draws=NUM_SAMPLES,
    model=prior_model,
    var_names=["obs"],
    random_seed=SEED,
)
attach_constant_data(
    prior_idata,
    name="week",
    values=y_train[TIME_DIM].values,
    dim=TIME_DIM,
)

prior_pc = az.plot_lm(
    prior_idata,
    y="obs",
    x="week",
    group="prior_predictive",
    ci_kind="hdi",
    ci_prob=(0.5, 0.94),
    smooth=False,
    visuals={
        "ci_band": {"color": "C0"},
        "observed_scatter": False,
        "pe_line": False,
    },
    figure_kwargs={"figsize": (10, 6)},
)
ax = prior_pc.viz["figure"].item().axes[0]
ax.plot(y_train[TIME_DIM].values, y_train.values, color="black", lw=1, label="train")
ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))
ax.set(title="Prior predictive check", xlabel="week", ylabel="log(# rides)")
plt.show()

Inference with ADVI

Forecaster fits the model with PyMC variational inference. The non-centered random walk keeps the high-dimensional weekly latent trajectory tractable for ADVI.

forecaster = Forecaster(
    model,
    y_train,
    covariates_train,
    optimizer=0.005,
    num_steps=NUM_STEPS,
    random_seed=SEED,
    progressbar=False,
)

fig, ax = plt.subplots()
ax.plot(forecaster.losses)
ax.set(title="ELBO loss", xlabel="ADVI step", ylabel="loss")
plt.show()

/Users/juanitorduz/.paseo/worktrees/1dtaxlkh/docs-example-notebook-visual-style/pymc_forecast/fit.py:251: UserWarning: ADVI has not converged after 20000 steps: the ELBO loss is still descending (median over the last 2000 steps improved by 869 since mid-run, more than the within-window fluctuation 19.5). The forecast may be confidently wrong — increase num_steps, raise the learning rate, or use HMCForecaster; inspect the loss history via the `losses` attribute.
  _check_vi_convergence(losses, num_steps)

Posterior predictive check and forecast

train_pp_idata = forecaster.predict_in_sample(
    num_samples=NUM_SAMPLES,
    random_seed=SEED,
    progressbar=False,
)
forecast_idata = forecaster.forecast(
    covariates,
    num_samples=NUM_SAMPLES,
    random_seed=SEED,
    progressbar=False,
)
attach_constant_data(
    train_pp_idata,
    name="week",
    values=y_train[TIME_DIM].values,
    dim=TIME_DIM,
)

train_pp = train_pp_idata["posterior_predictive"]["obs"]
forecast = forecast_idata["predictions"]["forecast"]
forecast_lm_idata = lm_idata_from_forecast(
    forecast,
    y_test.rename({TIME_DIM: FUTURE_DIM}),
    x_name="week",
    x_values=y_test[TIME_DIM].values,
)

metrics = {
    "crps": eval_crps,
    "coverage_50": partial(eval_coverage, alpha=0.5),
    "coverage_94": partial(eval_coverage, alpha=0.94),
}
train_scores = evaluate_forecast(train_pp, y_train, metrics=metrics)
test_scores = evaluate_forecast(forecast, y_test.rename({TIME_DIM: FUTURE_DIM}), metrics=metrics)
print("train:", train_scores)
print("test: ", test_scores)

train: {'crps': 0.02855502688926316, 'coverage_50': 0.6378896882494005, 'coverage_94': 0.9640287769784173}
test:  {'crps': 0.03481793016505621, 'coverage_50': 0.6346153846153846, 'coverage_94': 0.9038461538461539}

Forecast visualization

forecast_title = (
    f"Forecast (train CRPS: {train_scores['crps']:.3f}, test CRPS: {test_scores['crps']:.3f})"
)

pc = az.plot_lm(
    train_pp_idata,
    y="obs",
    x="week",
    ci_kind="hdi",
    ci_prob=(0.5, 0.94),
    smooth=False,
    visuals={"ci_band": {"color": "C0"}, "observed_scatter": False, "pe_line": False},
    figure_kwargs={"figsize": (12, 6)},
)
az.plot_lm(
    forecast_lm_idata,
    y="obs",
    x="week",
    plot_collection=pc,
    ci_kind="hdi",
    ci_prob=(0.5, 0.94),
    smooth=False,
    visuals={"ci_band": {"color": "C1"}, "observed_scatter": False, "pe_line": False},
)
ax = pc.viz["figure"].item().axes[0]
ax.plot(y[TIME_DIM].values, y.values, color="black", lw=1, label="observed")
ax.axvline(T1, color="gray", ls="--", label="train/test split")
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.12), ncol=3)
ax.set(title=forecast_title, xlabel="week", ylabel="log(# rides)")
plt.show()

Rolling-origin backtesting

The supported PyMC API refits expanding windows one fold at a time. Each fold forecasts one year, advances one year, and keeps predictions labeled by their real week coordinates. The upstream JAX-only vectorized backtest is intentionally omitted.

backtest_model = LocalLevelSeasonal()
results = backtest(
    y,
    covariates,
    backtest_model,
    metrics=metrics,
    test_window=52,
    stride=52,
    min_train_window=104,
    min_test_window=52,
    num_samples=BACKTEST_SAMPLES,
    eval_train=True,
    keep_predictions=True,
    forecaster_options={"optimizer": 0.005, "num_steps": NUM_STEPS},
    random_seed=SEED,
)
summary = results_to_dataframe(results)
summary

/Users/juanitorduz/.paseo/worktrees/1dtaxlkh/docs-example-notebook-visual-style/pymc_forecast/fit.py:251: UserWarning: ADVI has not converged after 20000 steps: the ELBO loss is still descending (median over the last 2000 steps improved by 108 since mid-run, more than the within-window fluctuation 9.13). The forecast may be confidently wrong — increase num_steps, raise the learning rate, or use HMCForecaster; inspect the loss history via the `losses` attribute.
  _check_vi_convergence(losses, num_steps)

/Users/juanitorduz/.paseo/worktrees/1dtaxlkh/docs-example-notebook-visual-style/pymc_forecast/fit.py:251: UserWarning: ADVI has not converged after 20000 steps: the ELBO loss is still descending (median over the last 2000 steps improved by 165 since mid-run, more than the within-window fluctuation 12.9). The forecast may be confidently wrong — increase num_steps, raise the learning rate, or use HMCForecaster; inspect the loss history via the `losses` attribute.
  _check_vi_convergence(losses, num_steps)

/Users/juanitorduz/.paseo/worktrees/1dtaxlkh/docs-example-notebook-visual-style/pymc_forecast/fit.py:251: UserWarning: ADVI has not converged after 20000 steps: the ELBO loss is still descending (median over the last 2000 steps improved by 378 since mid-run, more than the within-window fluctuation 14.2). The forecast may be confidently wrong — increase num_steps, raise the learning rate, or use HMCForecaster; inspect the loss history via the `losses` attribute.
  _check_vi_convergence(losses, num_steps)

/Users/juanitorduz/.paseo/worktrees/1dtaxlkh/docs-example-notebook-visual-style/pymc_forecast/fit.py:251: UserWarning: ADVI has not converged after 20000 steps: the ELBO loss is still descending (median over the last 2000 steps improved by 665 since mid-run, more than the within-window fluctuation 17). The forecast may be confidently wrong — increase num_steps, raise the learning rate, or use HMCForecaster; inspect the loss history via the `losses` attribute.
  _check_vi_convergence(losses, num_steps)

/Users/juanitorduz/.paseo/worktrees/1dtaxlkh/docs-example-notebook-visual-style/pymc_forecast/fit.py:251: UserWarning: ADVI has not converged after 20000 steps: the ELBO loss is still descending (median over the last 2000 steps improved by 875 since mid-run, more than the within-window fluctuation 20). The forecast may be confidently wrong — increase num_steps, raise the learning rate, or use HMCForecaster; inspect the loss history via the `losses` attribute.
  _check_vi_convergence(losses, num_steps)

	t0	t1	t2	num_samples	train_walltime	test_walltime	crps	coverage_50	coverage_94	train_crps	train_coverage_50	train_coverage_94
0	0	104	156	500	2.360420	0.162516	0.068290	0.442308	0.942308	0.024918	0.855769	0.990385
1	0	156	208	500	2.112400	0.135653	0.040307	0.519231	0.923077	0.034820	0.743590	0.980769
2	0	208	260	500	2.094671	0.149827	0.030875	0.557692	0.980769	0.031596	0.716346	0.971154
3	0	260	312	500	2.043023	0.136346	0.025316	0.673077	0.980769	0.028461	0.726923	0.973077
4	0	312	364	500	2.041257	0.140213	0.029328	0.557692	0.980769	0.028013	0.692308	0.967949
5	0	364	416	500	2.200452	0.135489	0.037924	0.423077	0.903846	0.027974	0.689560	0.969780
6	0	416	468	500	2.124175	0.143270	0.037585	0.461538	0.903846	0.028298	0.653846	0.966346

fig, ax = plt.subplots()
ax.plot(summary["t1"], summary["train_crps"], "o-", color="C0", label="in-sample CRPS")
ax.plot(summary["t1"], summary["crps"], "o-", color="C1", label="out-of-sample CRPS")
ax.legend()
ax.set(xlabel="train/test split week", ylabel="CRPS", title="CRPS per backtest fold")
plt.show()

pc = None
for result in results:
    if result.prediction is None:
        continue
    truth = y.isel({TIME_DIM: slice(result.t1, result.t2)}).rename({TIME_DIM: FUTURE_DIM})
    fold_idata = lm_idata_from_forecast(
        result.prediction,
        truth,
        x_name="week",
        x_values=np.arange(result.t1, result.t2),
    )
    if pc is None:
        pc = az.plot_lm(
            fold_idata,
            y="obs",
            x="week",
            ci_kind="hdi",
            ci_prob=(0.5, 0.94),
            smooth=False,
            visuals={"ci_band": {"color": "C1"}, "observed_scatter": False, "pe_line": False},
            figure_kwargs={"figsize": (12, 6)},
        )
    else:
        az.plot_lm(
            fold_idata,
            y="obs",
            x="week",
            plot_collection=pc,
            ci_kind="hdi",
            ci_prob=(0.5, 0.94),
            smooth=False,
            visuals={"ci_band": {"color": "C1"}, "observed_scatter": False, "pe_line": False},
        )

if pc is None:
    msg = "no folds were plotted"
    raise ValueError(msg)
ax = pc.viz["figure"].item().axes[0]
ax.plot(y[TIME_DIM].values, y.values, color="black", lw=1, label="observed")
for result in results:
    ax.axvline(result.t1, color="gray", ls="--", lw=0.6)
ax.legend(loc="lower left")
ax.set(title="Rolling-origin forecasts", xlabel="week", ylabel="log(# rides)")
plt.show()

Functional API

The same model can be written as a plain function (covariates, data=None). It builds the horizon itself with h = Horizon.from_data(covariates, data) and passes it to the module-level primitives innovations(h, ...) and predict(h, ...), which the class’s bound helpers wrap. The function is interchangeable with LocalLevelSeasonal(): pass it to build_model, any forecaster or backtest in its place.

Inference has a functional form too. The forecaster classes call fit_vi, fit_mcmc and fit_pathfinder (Forecaster, HMCForecaster and PathfinderForecaster respectively), which return a FitResult; draw_posterior and forecast then produce predictions:

result = fit_vi(
    local_level_seasonal,
    y_train,
    covariates_train,
    optimizer=0.005,
    num_steps=NUM_STEPS,
    random_seed=SEED,
)
posterior = draw_posterior(result, NUM_SAMPLES, random_seed=SEED)
samples = forecast(local_level_seasonal, posterior, y_train, covariates, random_seed=SEED)

The cell below defines the functional twin and checks that both forms register the same variables.

def local_level_seasonal(covariates, data=None) -> None:
    """Functional twin of ``LocalLevelSeasonal.model``."""
    h = Horizon.from_data(covariates, data)
    bias = pm.Normal("bias", 0.0, 10.0)
    weight = pm.Normal("weight", 0.0, 0.1, dims="covariate")
    drift_scale = pm.LogNormal("drift_scale", -20.0, 5.0, initval=0.01)
    nu = pm.Gamma("nu", alpha=10.0, beta=2.0)
    sigma = pm.LogNormal("sigma", -5.0, 5.0, initval=0.05)

    drift_raw = innovations(h, "drift_raw", pm.Normal.dist(0.0, 1.0))
    level = pt.cumsum(drift_raw * drift_scale)
    regression = pt.dot(covariates.values, weight)

    predict(h, pm.StudentT.dist(nu=nu, sigma=sigma), bias + level + regression)


oop_names = list(build_model(LocalLevelSeasonal(), y_train, covariates_train).named_vars)
functional_names = list(build_model(local_level_seasonal, y_train, covariates_train).named_vars)
print("same named variables:", oop_names == functional_names)
print(functional_names)

same named variables: True
['bias', 'weight', 'drift_scale', 'nu', 'sigma', 'drift_raw', 'obs', 'mu']

References

• Upstream NumPyro notebook: https://github.com/juanitorduz/numpyro_forecast/blob/main/docs/examples/forecasting_univariate.ipynb

• Pyro BART dataset mirror: https://github.com/pyro-ppl/datasets/tree/master/bart

• Juan Orduz, Univariate time series forecasting with NumPyro: https://juanitorduz.github.io/numpyro_forecasting-univariate/