numpyro-forecast 0.1.0__py3-none-any.whl

numpyro_forecast/__init__.py

@@ -0,0 +1,45 @@
+"""numpyro_forecast: a JAX/NumPyro port of Pyro's forecasting module."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from jaxtyping import install_import_hook
+
+with install_import_hook("numpyro_forecast", "beartype.beartype"):
+    from numpyro_forecast import (  # noqa: F401
+        datasets,
+        evaluate,
+        forecaster,
+        functional,
+        metrics,
+        util,
+    )
+
+from numpyro_forecast.evaluate import (
+    BacktestResult,
+    backtest,
+    eval_coverage,
+    eval_crps,
+    eval_mae,
+    eval_rmse,
+)
+from numpyro_forecast.forecaster import Forecaster, ForecastingModel, HMCForecaster
+from numpyro_forecast.functional import forecasting_model
+
+try:
+    __version__ = version("numpyro_forecast")
+except PackageNotFoundError:  # pragma: no cover - package not installed
+    __version__ = "0.0.0+unknown"
+
+__all__ = [
+    "BacktestResult",
+    "Forecaster",
+    "ForecastingModel",
+    "HMCForecaster",
+    "__version__",
+    "backtest",
+    "eval_coverage",
+    "eval_crps",
+    "eval_mae",
+    "eval_rmse",
+    "forecasting_model",
+]

numpyro_forecast/datasets.py

@@ -0,0 +1,104 @@
+"""BART ridership dataset helpers.
+
+Thin wrappers around :func:`numpyro.examples.datasets.load_bart_od` that return
+arrays in the package convention (time at axis ``-2``) for the two examples.
+"""
+
+import jax.numpy as jnp
+from jaxtyping import Float
+
+from numpyro_forecast.typing import Array
+
+HOURS_PER_WEEK = 24 * 7
+
+
+def bart_available() -> bool:
+    """Return whether the BART dataset can be loaded (download succeeds).
+
+    Returns
+    -------
+    bool
+        ``True`` if :func:`load_bart_od` loads without error.
+    """
+    try:
+        _load_counts()
+    except Exception:
+        return False
+    return True
+
+
+def _load_counts() -> tuple[Array, list[str]]:
+    """Load raw hourly origin-destination counts ``(time, origin, destin)``."""
+    from numpyro.examples.datasets import load_bart_od
+
+    dataset = load_bart_od()
+    counts = jnp.asarray(dataset["counts"])
+    stations = [str(name) for name in dataset["stations"]]
+    return counts, stations
+
+
+def load_bart_weekly() -> Float[Array, " weeks 1"]:
+    """Load total weekly BART ridership (log scale) for the univariate example.
+
+    Hourly counts are summed over all origin-destination pairs, aggregated into
+    non-overlapping weeks, and log-transformed.
+
+    Returns
+    -------
+    Float[Array, " weeks 1"]
+        Log weekly totals with time at axis ``-2`` and a single observation dim.
+    """
+    counts, _ = _load_counts()
+    hourly_total = counts.sum(axis=(1, 2))
+    num_weeks = hourly_total.shape[0] // HOURS_PER_WEEK
+    weekly = hourly_total[: num_weeks * HOURS_PER_WEEK]
+    weekly = weekly.reshape(num_weeks, HOURS_PER_WEEK).sum(axis=1)
+    return jnp.log(weekly)[:, None]
+
+
+def load_bart_hierarchical(
+    train_days: int = 90,
+    test_weeks: int = 2,
+) -> tuple[Float[Array, " origin time destin"], int, list[str]]:
+    """Load the windowed hierarchical BART panel for the hierarchical example.
+
+    The counts are ``log1p``-transformed and transposed to the
+    ``(origin, time, destin)`` convention, then restricted to a ``train_days``
+    training window followed by a ``test_weeks`` test window.
+
+    Parameters
+    ----------
+    train_days
+        Number of training days (24 hours each).
+    test_weeks
+        Number of test weeks (``24 * 7`` hours each).
+
+    Returns
+    -------
+    y : Float[Array, " origin time destin"]
+        Log counts over the train+test window with time at axis ``-2``.
+    split : int
+        Index along the time axis separating train from test.
+    stations : list[str]
+        Station names.
+
+    Raises
+    ------
+    ValueError
+        If the requested ``train_days`` + ``test_weeks`` window exceeds the
+        available history (which would otherwise wrap a negative slice index).
+    """
+    counts, stations = _load_counts()
+    log_counts = jnp.log1p(jnp.transpose(counts, (1, 0, 2)))
+    t_total = log_counts.shape[1]
+    t1 = t_total - test_weeks * HOURS_PER_WEEK
+    t0 = t1 - train_days * 24
+    if t0 < 0 or t1 <= t0:
+        msg = (
+            f"requested window (train_days={train_days}, test_weeks={test_weeks}) "
+            f"exceeds available history of {t_total} hours"
+        )
+        raise ValueError(msg)
+    y = log_counts[:, t0:t_total, :]
+    split = t1 - t0
+    return y, split, stations

numpyro_forecast/evaluate.py

@@ -0,0 +1,279 @@
+"""Backtesting and evaluation metrics.
+
+This is the JAX/NumPyro port of ``pyro.contrib.forecast.evaluate``. Unlike Pyro
+there is no global parameter store, so each backtest window is a pure call that
+fits its own forecaster.
+"""
+
+from collections.abc import Callable, Mapping
+from dataclasses import asdict, dataclass, field
+from time import perf_counter
+from typing import Any, cast
+
+import jax.numpy as jnp
+from jax import random
+
+from numpyro_forecast.forecaster import Forecaster
+from numpyro_forecast.metrics import crps_empirical
+from numpyro_forecast.typing import Array, ForecasterFactory, Metric, ModelFactory
+
+
+def eval_mae(pred: Array, truth: Array) -> float:
+    """Mean absolute error using the forecast sample median as point estimate.
+
+    Parameters
+    ----------
+    pred
+        Forecast samples with the sample axis first.
+    truth
+        Ground-truth values (matching ``pred`` without the sample axis).
+
+    Returns
+    -------
+    float
+        The mean absolute error.
+    """
+    point = jnp.median(pred, axis=0)
+    return float(jnp.abs(point - truth).mean())
+
+
+def eval_rmse(pred: Array, truth: Array) -> float:
+    """Root mean squared error using the forecast sample mean as point estimate.
+
+    Parameters
+    ----------
+    pred
+        Forecast samples with the sample axis first.
+    truth
+        Ground-truth values (matching ``pred`` without the sample axis).
+
+    Returns
+    -------
+    float
+        The root mean squared error.
+    """
+    point = pred.mean(axis=0)
+    return float(jnp.sqrt(jnp.square(point - truth).mean()))
+
+
+def eval_crps(pred: Array, truth: Array) -> float:
+    """Empirical CRPS averaged over all data elements.
+
+    Parameters
+    ----------
+    pred
+        Forecast samples with the sample axis first.
+    truth
+        Ground-truth values (matching ``pred`` without the sample axis).
+
+    Returns
+    -------
+    float
+        The mean empirical CRPS.
+    """
+    return float(crps_empirical(pred, truth).mean())
+
+
+def eval_coverage(pred: Array, truth: Array, *, alpha: float = 0.9) -> float:
+    """Empirical coverage of the central ``alpha`` prediction interval.
+
+    The central ``alpha`` interval is bounded by the ``(1 - alpha) / 2`` and
+    ``1 - (1 - alpha) / 2`` quantiles of the forecast samples; the metric is the
+    fraction of ground-truth values that fall inside it. A well-calibrated
+    forecast has coverage close to ``alpha``.
+
+    Parameters
+    ----------
+    pred
+        Forecast samples with the sample axis first.
+    truth
+        Ground-truth values (matching ``pred`` without the sample axis).
+    alpha
+        Nominal interval level in ``(0, 1)`` (defaults to ``0.9``).
+
+    Returns
+    -------
+    float
+        The fraction of ground-truth values inside the central ``alpha`` interval.
+    """
+    tail = (1.0 - alpha) / 2.0
+    lo = jnp.quantile(pred, tail, axis=0)
+    hi = jnp.quantile(pred, 1.0 - tail, axis=0)
+    return float(((truth >= lo) & (truth <= hi)).mean())
+
+
+DEFAULT_METRICS: dict[str, Metric] = {
+    "mae": eval_mae,
+    "rmse": eval_rmse,
+    "crps": eval_crps,
+    "coverage": eval_coverage,
+}
+"""Default metrics used by :func:`backtest`."""
+
+
+@dataclass(frozen=True)
+class BacktestResult:
+    """Per-window result of a :func:`backtest` run.
+
+    Attributes
+    ----------
+    t0, t1, t2
+        Train-begin, train/test split, and test-end time indices.
+    num_samples
+        Number of forecast samples drawn.
+    train_walltime, test_walltime
+        Wall-clock seconds for fitting and forecasting.
+    metrics
+        Mapping of metric name to value for the window.
+    params
+        Mapping of scalar parameter name to value (when available).
+    """
+
+    t0: int
+    t1: int
+    t2: int
+    num_samples: int
+    train_walltime: float
+    test_walltime: float
+    metrics: dict[str, float]
+    params: dict[str, float] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Return a flat dictionary view (Pyro-style access).
+
+        Returns
+        -------
+        dict[str, Any]
+            All fields as a plain dictionary.
+        """
+        return asdict(self)
+
+
+def _scalar_params(forecaster: object) -> dict[str, float]:
+    """Extract scalar variational parameters from a fitted forecaster, if any."""
+    params = getattr(forecaster, "params", None)
+    if not isinstance(params, Mapping):
+        return {}
+    return {name: float(value) for name, value in params.items() if jnp.size(value) == 1}
+
+
+def backtest(
+    rng_key: Array,
+    data: Array,
+    covariates: Array,
+    model_fn: ModelFactory,
+    *,
+    forecaster_fn: ForecasterFactory = Forecaster,
+    metrics: Mapping[str, Metric] | None = None,
+    transform: Callable[[Array, Array], tuple[Array, Array]] | None = None,
+    train_window: int | None = None,
+    min_train_window: int = 1,
+    test_window: int | None = None,
+    min_test_window: int = 1,
+    stride: int = 1,
+    num_samples: int = 100,
+    batch_size: int | None = None,
+    forecaster_options: Mapping[str, Any] | Callable[..., Mapping[str, Any]] | None = None,
+) -> list[BacktestResult]:
+    """Backtest a forecasting model on a moving window of ``(train, test)`` data.
+
+    Parameters
+    ----------
+    rng_key
+        Base PRNG key (used for every window, matching Pyro).
+    data
+        Dataset with time at axis ``-2``.
+    covariates
+        Covariates with time at axis ``-2`` (same duration as ``data``).
+    model_fn
+        Factory returning a fresh :class:`ForecastingModel` per window.
+    forecaster_fn
+        Factory returning a fitted forecaster (defaults to :class:`Forecaster`).
+    metrics
+        Mapping of metric name to function; defaults to :data:`DEFAULT_METRICS`.
+    transform
+        Optional ``(pred, truth) -> (pred, truth)`` applied before metrics.
+    train_window
+        Training window size; if ``None`` the window expands from the start.
+    min_train_window
+        Minimum training window size when ``train_window`` is ``None``.
+    test_window
+        Test window size; if ``None`` forecasts to the end of the data.
+    min_test_window
+        Minimum test window size when ``test_window`` is ``None``.
+    stride
+        Step between successive train/test splits.
+    num_samples
+        Number of forecast samples per window.
+    batch_size
+        Optional forecast-sampling chunk size.
+    forecaster_options
+        Options dict passed to ``forecaster_fn``, or a callable
+        ``(t0, t1, t2) -> dict`` returning per-window options.
+
+    Returns
+    -------
+    list[BacktestResult]
+        One result per backtest window.
+    """
+    if data.shape[-2] != covariates.shape[-2]:
+        msg = "data and covariates must share the time axis length"
+        raise ValueError(msg)
+    metrics = DEFAULT_METRICS if metrics is None else metrics
+
+    def options_for(t0: int, t1: int, t2: int) -> Mapping[str, Any]:
+        if forecaster_options is None:
+            return {}
+        if callable(forecaster_options) and not isinstance(forecaster_options, Mapping):
+            return forecaster_options(t0=t0, t1=t1, t2=t2)
+        return cast("Mapping[str, Any]", forecaster_options)
+
+    duration = data.shape[-2]
+    stop = duration - (min_test_window if test_window is None else test_window) + 1
+    start = min_train_window if train_window is None else train_window
+
+    results: list[BacktestResult] = []
+    for t1 in range(start, stop, stride):
+        t0 = 0 if train_window is None else t1 - train_window
+        t2 = duration if test_window is None else t1 + test_window
+
+        train_data = data[..., t0:t1, :]
+        train_covariates = covariates[..., t0:t1, :]
+        test_covariates = covariates[..., t0:t2, :]
+        truth = data[..., t1:t2, :]
+
+        key_fit, key_forecast = random.split(rng_key)
+        options = options_for(t0, t1, t2)
+
+        fit_start = perf_counter()
+        model = model_fn()
+        forecaster = forecaster_fn(key_fit, model, train_data, train_covariates, **options)
+        train_walltime = perf_counter() - fit_start
+
+        forecast_start = perf_counter()
+        pred = forecaster(
+            key_forecast,
+            train_data,
+            test_covariates,
+            num_samples,
+            batch_size=batch_size,
+        )
+        test_walltime = perf_counter() - forecast_start
+
+        if transform is not None:
+            pred, truth = transform(pred, truth)
+
+        results.append(
+            BacktestResult(
+                t0=t0,
+                t1=t1,
+                t2=t2,
+                num_samples=num_samples,
+                train_walltime=train_walltime,
+                test_walltime=test_walltime,
+                metrics={name: fn(pred, truth) for name, fn in metrics.items()},
+                params=_scalar_params(forecaster),
+            )
+        )
+
+    return results