from __future__ import annotations from dataclasses import dataclass from typing import Iterator, List, Optional import numpy as np from gluonts.dataset import Dataset as GluonDataset from gluonts.model import Forecast from gluonts.model.forecast import QuantileForecast from gluonts.model.predictor import RepresentablePredictor @dataclass class _ForecastConfig: quantile_levels: Optional[List[float]] forecast_keys: List[str] intervals: Optional[List[int]] @classmethod def from_quantiles(cls, quantile_levels: Optional[List[float]]): forecast_keys = ["mean"] intervals = None if not quantile_levels: return cls(quantile_levels=None, forecast_keys=forecast_keys, intervals=intervals) intervals = sorted( { round(200 * (max(q, 1 - q) - 0.5)) for q in quantile_levels } ) forecast_keys = ["mean"] + [str(q) for q in quantile_levels] return cls( quantile_levels=quantile_levels, forecast_keys=forecast_keys, intervals=intervals, ) class SeasonalNaivePredictor(RepresentablePredictor): """Lightweight seasonal naive baseline for pipeline smoke tests.""" def __init__( self, prediction_length: int, season_length: int = 1, quantile_levels: Optional[List[float]] = None, ): super().__init__(prediction_length=prediction_length) self.season_length = max(1, season_length) self.config = _ForecastConfig.from_quantiles(quantile_levels) def predict(self, dataset: GluonDataset, **kwargs) -> Iterator[Forecast]: for entry in dataset: target = np.asarray(entry["target"], dtype=np.float64) univariate = target.ndim == 1 if univariate: working = target.reshape(1, -1) else: working = target history = working[:, -self.season_length :] if history.shape[-1] == 0: raise ValueError("Seasonal-Naive requires at least one observation") tiled = np.tile( history, (1, int(np.ceil(self.prediction_length / history.shape[-1]))) ) mean = tiled[:, : self.prediction_length] if univariate: mean = mean[0] arrays = [mean] if self.config.quantile_levels: for q in self.config.quantile_levels: arrays.append(mean + np.abs(mean) * (q - 0.5) * 0.05) forecast_arrays = np.stack(arrays, axis=0) yield QuantileForecast( forecast_arrays=forecast_arrays, forecast_keys=self.config.forecast_keys, start_date=entry["start"] + target.shape[-1], item_id=entry["item_id"], ) class RepeatedPredictionQuantileMixin: """Interface for deriving quantiles from repeated deterministic predictions. Models that do not natively produce quantiles can implement ``point_forecast_1d``. The default ``sample_forecast_1d`` repeatedly perturbs that point forecast with bootstrapped in-sample residuals, then converts the resulting sample paths into quantile forecasts. """ quantile_levels: Optional[List[float]] num_samples: int random_seed: int def point_forecast_1d(self, history: np.ndarray) -> np.ndarray: raise NotImplementedError def residuals_1d(self, history: np.ndarray, fitted: np.ndarray | None = None) -> np.ndarray: if fitted is not None and fitted.size == history.size: residuals = history - fitted else: residuals = np.diff(history) residuals = residuals[np.isfinite(residuals)] if residuals.size == 0: scale = np.nanstd(history) if not np.isfinite(scale) or scale == 0: scale = 1e-6 return np.array([scale], dtype=np.float64) return residuals.astype(np.float64) def sample_forecast_1d(self, history: np.ndarray) -> np.ndarray: mean = self.point_forecast_1d(history).astype(np.float64) residuals = self.residuals_1d(history) rng = np.random.default_rng(self.random_seed + int(abs(np.nansum(history))) % 100000) centered = residuals - float(np.nanmean(residuals)) if centered.size == 0 or np.nanstd(centered) == 0: scale = float(np.nanstd(history)) centered = rng.normal(0.0, scale if scale > 0 else 1e-6, size=max(8, self.num_samples)) innovations = rng.choice(centered, size=(self.num_samples, self.prediction_length), replace=True) return mean[None, :] + innovations def quantile_forecast_arrays(self, history: np.ndarray) -> np.ndarray: samples = self.sample_forecast_1d(history) mean = np.mean(samples, axis=0) arrays = [mean] if self.quantile_levels: arrays.extend(np.quantile(samples, q, axis=0) for q in self.quantile_levels) return np.stack(arrays, axis=0).astype(np.float64) class StatisticalBaselinePredictor(RepeatedPredictionQuantileMixin, RepresentablePredictor): """Base class for lightweight statistical baselines with quantile fallback.""" def __init__( self, prediction_length: int, quantile_levels: Optional[List[float]] = None, num_samples: int = 200, random_seed: int = 0, ): super().__init__(prediction_length=prediction_length) self.quantile_levels = quantile_levels self.num_samples = max(20, int(num_samples)) self.random_seed = int(random_seed) self.config = _ForecastConfig.from_quantiles(quantile_levels) def predict(self, dataset: GluonDataset, **kwargs) -> Iterator[Forecast]: for entry in dataset: target = np.asarray(entry["target"], dtype=np.float64) univariate = target.ndim == 1 working = target.reshape(1, -1) if univariate else target rows = [] for variate_idx in range(working.shape[0]): series = working[variate_idx] series = series[np.isfinite(series)] if series.size == 0: series = np.zeros(1, dtype=np.float64) rows.append(self.quantile_forecast_arrays(series)) stacked = np.stack(rows, axis=0) forecast_arrays = stacked[0] if univariate else stacked yield QuantileForecast( forecast_arrays=forecast_arrays, forecast_keys=self.config.forecast_keys, start_date=entry["start"] + target.shape[-1], item_id=entry["item_id"], ) class MovingAveragePredictor(StatisticalBaselinePredictor): """Rolling mean baseline with bootstrap quantile intervals.""" def __init__(self, prediction_length: int, window: int = 24, **kwargs): super().__init__(prediction_length=prediction_length, **kwargs) self.window = max(1, int(window)) def point_forecast_1d(self, history: np.ndarray) -> np.ndarray: window = history[-min(self.window, history.size):] value = float(np.nanmean(window)) if window.size else 0.0 return np.full(self.prediction_length, value, dtype=np.float64) class ARIMAPredictor(StatisticalBaselinePredictor): """ARIMA baseline using statsmodels, with residual bootstrap quantiles.""" def __init__(self, prediction_length: int, order: tuple[int, int, int] = (1, 1, 1), **kwargs): super().__init__(prediction_length=prediction_length, **kwargs) self.order = order self._last_fitted: np.ndarray | None = None def point_forecast_1d(self, history: np.ndarray) -> np.ndarray: self._last_fitted = None if history.size < 8 or np.nanstd(history) == 0: return np.full(self.prediction_length, float(history[-1]), dtype=np.float64) try: from statsmodels.tsa.arima.model import ARIMA fit = ARIMA(history, order=self.order).fit() fitted = np.asarray(fit.fittedvalues, dtype=np.float64) if fitted.size == history.size: self._last_fitted = fitted forecast = np.asarray(fit.forecast(steps=self.prediction_length), dtype=np.float64) if forecast.size == self.prediction_length and np.all(np.isfinite(forecast)): return forecast except Exception: pass return np.full(self.prediction_length, float(history[-1]), dtype=np.float64) def residuals_1d(self, history: np.ndarray, fitted: np.ndarray | None = None) -> np.ndarray: return super().residuals_1d(history, self._last_fitted) class ETSPredictor(StatisticalBaselinePredictor): """Exponential smoothing baseline with residual bootstrap quantiles.""" def point_forecast_1d(self, history: np.ndarray) -> np.ndarray: if history.size < 4 or np.nanstd(history) == 0: return np.full(self.prediction_length, float(history[-1]), dtype=np.float64) try: from statsmodels.tsa.holtwinters import ExponentialSmoothing trend = "add" if history.size >= 8 else None fit = ExponentialSmoothing( history, trend=trend, seasonal=None, initialization_method="estimated", ).fit(optimized=True) forecast = np.asarray(fit.forecast(self.prediction_length), dtype=np.float64) if forecast.size == self.prediction_length and np.all(np.isfinite(forecast)): return forecast except Exception: pass return np.full(self.prediction_length, float(history[-1]), dtype=np.float64)