Skip to main content

evaluate.backtest()

Backtest a forecasting model on a moving window of (train, test) data.

Usage

Source

evaluate.backtest(
    rng_key,
    model_fn,
    data,
    covariates,
    *,
    forecast_fn,
    in_sample_fn=None,
    metrics=None,
    per_window_metrics=None,
    transform=None,
    window_type=None,
    train_window=None,
    min_train_window=1,
    test_window=None,
    min_test_window=1,
    stride=1,
    num_samples=100,
    batch_size=None,
    eval_train=False,
    keep_predictions=False,
    reuse_model=True
)

Fitting and forecasting are delegated entirely to user-supplied closures rather than an OOP forecaster: forecast_fn fits model on the training window and forecasts the test horizon, and the optional in_sample_fn fits and scores the in-sample fit. Both closures own their own inference backend (SVI, MCMC, or anything else), so backtest itself has no dependency on how a model is fit.

forecast_fn has the call signature (see ForecastFn):

forecast_fn(
    rng_key, model, train_data, train_covariates, full_covariates,
    num_samples, *, batch_size=None,
) -> draws  # shape (num_samples, *batch, t2 - t1, obs)

where full_covariates spans the full window, covariates[..., t0:t2, :] (train followed by test), matching what the model needs to run the forecast horizon. The optional in_sample_fn has the call signature (see InSampleFn):

in_sample_fn(
    rng_key, model, train_data, train_covariates, num_samples, *, batch_size=None,
) -> draws  # shape (num_samples, *batch, t1 - t0, obs)

batch_size is forwarded unchanged into both closures so a chunked implementation can bound its own device memory. A closure may return draws committed to host memory (e.g. via device="host", to cap peak accelerator usage): every metric in DEFAULT_METRICS accepts a host-committed pred or truth (or both), in any mix and regardless of batch_size, moving a host-committed operand to device memory first where needed. Returning draws already on-device still avoids the extra host-to-device hop for metrics scored every window.

A minimal forecast_fn built on plain NumPyro (AutoNormal + SVI.run + Predictive):

import numpyro
from jax import random
from numpyro.infer import SVI, Predictive, Trace_ELBO
from numpyro.infer.autoguide import AutoNormal


def forecast_fn(
    rng_key,
    model,
    train_data,
    train_covariates,
    full_covariates,
    num_samples,
    *,
    batch_size=None,
):
    guide = AutoNormal(model)
    svi = SVI(model, guide, numpyro.optim.Adam(0.01), Trace_ELBO())
    key_fit, key_post, key_pred = random.split(rng_key, 3)
    state = svi.run(key_fit, 1_000, train_covariates, train_data, progress_bar=False)
    posterior = guide.sample_posterior(key_post, state.params, sample_shape=(num_samples,))
    predictive = Predictive(model, posterior_samples=posterior, return_sites=["forecast"])
    return predictive(key_pred, full_covariates, train_data)["forecast"]

Parameters

rng_key: Array

Base PRNG key, used for every window (matching Pyro). Each window splits it once, forecast_key, train_key = random.split(rng_key): forecast_fn receives forecast_key and in_sample_fn receives train_key. The split does not depend on eval_train, so turning in-sample scoring on does not change the out-of-sample draws.

model_fn: ModelFactory

Factory returning a fresh ForecastModel per window.

data: Array

Dataset with time at axis -2.

covariates: Array

Covariates with time at axis -2 (same duration as data).

forecast_fn: ForecastFn

Closure that fits model on the training window and forecasts the test horizon (see ForecastFn and the contract above).

in_sample_fn: InSampleFn | None = None

Optional closure that fits model on the training window and scores its in-sample fit (see InSampleFn and the contract above). Required when eval_train=True. It receives its own child of rng_key (see rng_key), never a key that forecast_fn can derive from its own.

metrics: Mapping[str, Metric] | None = None

Mapping of metric name to function; defaults to DEFAULT_METRICS. Each function takes (pred, truth) and returns a scalar array (see Metric); bind any metric-specific parameters with functools.partial(), e.g. {**DEFAULT_METRICS, "coverage": partial(eval_coverage, alpha=0.8)}. Each window is scored through evaluate_forecast(), so the integer draws of a count model need no cast and every metric receives floating arrays. transform receives the draws as forecast_fn returned them; keep_predictions stores them after transform.

per_window_metrics: Callable[[int, int, int], Mapping[str, Metric]] | None = None

Optional (t0, t1, t2) -> Mapping[str, Metric] callable producing extra metrics merged over metrics for each window. Use it for window-dependent metrics such as a MASE scaled by that window’s training data (numpyro_forecast.metrics.make_mase()).

transform: Callable[[Array, Array], tuple[Array, Array]] | None = None

Optional (pred, truth) -> (pred, truth) applied before metrics. It runs before the metrics and receives the forecast/in-sample closure’s draws as-is, so a transform that does its own array math against a device-resident operand must convert a host-committed pred or truth first, e.g. with numpy.asarray() (or move it back onto a device explicitly).

window_type: WindowType | None = None

Windowing strategy. If None (default) it is inferred from train_window: "expanding" when train_window is None and "rolling" when it is set, matching the historical behavior. Pass "expanding" to always train on all history from t0 = 0, or "rolling" to hold the training length fixed at train_window and slide it forward. "expanding" and train_window are mutually exclusive, and "rolling" requires train_window (both validated).

train_window: int | None = None

Training window size; if None the window expands from the start. Required for window_type="rolling".

min_train_window: int = 1

Minimum training window size for the expanding strategy (used when train_window is None).

test_window: int | None = None

Test window size; if None forecasts to the end of the data.

min_test_window: int = 1

Minimum test window size when test_window is None.

stride: int = 1

Step between successive train/test splits.

num_samples: int = 100

Number of forecast samples per window. Must be positive, and at least 2 when metrics contains eval_crps() (as DEFAULT_METRICS does) or a functools.partial() of it. Both rules are checked before the first window is fitted. A CRPS inside another wrapper (a lambda or a jax.jit function) or added by per_window_metrics is not checked up front: it fails only when it runs.

batch_size: int | None = None

Optional chunk size forwarded to forecast_fn and in_sample_fn (see the contract above).

eval_train: bool = False

If True, also score the in-sample posterior predictive over each training window with the same metrics and store them in BacktestResult.train_metrics. Requires in_sample_fn.

keep_predictions: bool = False

If True, store each window’s out-of-sample forecast samples (after transform) on BacktestResult.prediction. Defaults to False to avoid retaining large Monte Carlo arrays.

reuse_model: bool = True
When True (default) and the windowing strategy is rolling, the model instance returned by the first model_fn() call is reused for every window, so model_fn runs once. Ignored for expanding windows and when False (model_fn then runs once per window). It does not cache compiled programs: each forecast() or predict_in_sample() call that forecast_fn or in_sample_fn makes compiles once per window, and so does an SVI fit; for a single fused fit and forecast over all windows use backtest_vectorized(). Kept for API stability: it only matters for a model_fn that is expensive or stateful.

Returns

list[BacktestResult]
One result per backtest window.

Raises

ValueError

If data or covariates has fewer than two axes (time must be at axis -2), if their durations differ, if eval_train=True but in_sample_fn is None, or if num_samples is too small (see num_samples). Also if a metric returns a value that is not 0-d or does not have a bool, integer or floating dtype (see evaluate_forecast()): the error names the metric and comes when the first window is scored, after its fit, because a metric may run on the host here and so is not traced before the fit.

BacktestWindowError
If window_type and train_window disagree ("rolling" without a train_window, or "expanding" with one).