From dcdd3d5f91b14938f8acffc4ab56f73cd07351c1 Mon Sep 17 00:00:00 2001 From: MMenchero Date: Sun, 31 May 2026 00:29:18 -0600 Subject: [PATCH 1/8] first native UCM implementation --- python/statsforecast/models.py | 124 +++++- python/statsforecast/ucm.py | 730 +++++++++------------------------ tests/test_ucm.py | 393 +++++++++--------- 3 files changed, 514 insertions(+), 733 deletions(-) diff --git a/python/statsforecast/models.py b/python/statsforecast/models.py index 317d828fd..52413c207 100644 --- a/python/statsforecast/models.py +++ b/python/statsforecast/models.py @@ -83,7 +83,7 @@ from .mstl import mstl from .tbats import _compute_sigmah, tbats_forecast, tbats_selection from .theta import auto_theta, forecast_theta, forward_theta -from .ucm import UCM, LocalLevel, LocalLinearTrend, SmoothTrend # noqa: F401 +from .ucm import ucm_forecast, ucm_model def _add_fitted_pi(res, se, level): @@ -6285,6 +6285,128 @@ def __init__( super().__init__(p, q=0, alias=alias) +class UCM(_TS): + r"""Unobserved Components Model (UCM). + + Also known as a Structural Time Series Model. + This model decomposes a univariate series into a local linear trend (stochastic level + stochastic slope), a dummy-variable seasonal component and an irregular term: + + ``` math + y_t = \mu_t + \gamma_t + \varepsilon_t + ``` + + The level and slope evolve as random walks and the seasonal effects are + constrained to sum to zero over a full cycle. The model is cast in state + space form and the component variances are estimated by maximum likelihood + using a univariate Kalman filter. + + References: + - [Harvey, A. C. (1989). "Forecasting, Structural Time Series Models and the Kalman Filter". Cambridge University Press.](https://www.cambridge.org/core/books/forecasting-structural-time-series-models-and-the-kalman-filter/CE5E112570A56960601760E786A5E631) + - [Durbin, J. and Koopman, S. J. (2012). "Time Series Analysis by State Space Methods". 2nd ed. Oxford University Press.](https://academic.oup.com/book/16563) + + Args: + season_length (int): Number of observations per seasonal cycle. Ex: 12 for monthly data. Use 1 for a trend-only model. + alias (str): Custom name of the model. + prediction_intervals (Optional[ConformalIntervals]): Information to compute conformal prediction intervals. + """ + + def __init__( + self, + season_length: int = 1, + alias: str = "UCM", + prediction_intervals: Optional[ConformalIntervals] = None, + ): + self.season_length = season_length + self.alias = alias + self.prediction_intervals = prediction_intervals + + def _no_pi(self, level: Optional[List[int]]): + if level is not None: + raise NotImplementedError( + "Prediction intervals are not yet supported for UCM." + ) + + def fit(self, y: np.ndarray, X: Optional[np.ndarray] = None): + r"""Fit the UCM model. + + Fit a UCM to a time series (numpy array) `y`. + + Args: + y (numpy.array): Clean time series of shape (t, ). + X (array-like): Optional exogenous of shape (t, n_x). Currently ignored. + + Returns: + self: UCM fitted model. + """ + y = _ensure_float(y) + self.model_ = ucm_model(y, season_length=self.season_length) + return self + + def predict( + self, h: int, X: Optional[np.ndarray] = None, level: Optional[List[int]] = None + ): + r"""Predict with fitted UCM. + + Args: + h (int): Forecast horizon. + X (array-like): Optional exogenous of shape (h, n_x). Currently ignored. + level (List[float]): Confidence levels (0-100) for prediction intervals. Not yet supported. + + Returns: + dict: Dictionary with entry `mean` for point predictions. + """ + self._no_pi(level) + fcst = ucm_forecast(self.model_, h) + return {"mean": fcst["mean"]} + + def predict_in_sample(self, level: Optional[List[int]] = None): + r"""Access fitted UCM insample predictions. + + Args: + level (List[float]): Confidence levels (0-100) for prediction intervals. Not yet supported. + + Returns: + dict: Dictionary with entry `fitted` for point predictions. + """ + self._no_pi(level) + return {"fitted": self.model_["fitted"]} + + def forecast( + self, + y: np.ndarray, + h: int, + X: Optional[np.ndarray] = None, + X_future: Optional[np.ndarray] = None, + level: Optional[List[int]] = None, + fitted: bool = False, + ): + r"""Memory Efficient UCM predictions. + + This method avoids memory burden due from object storage. + It is analogous to `fit_predict` without storing information. + It assumes you know the forecast horizon in advance. + + Args: + y (numpy.array): Clean time series of shape (n, ). + h (int): Forecast horizon. + X (array-like): Optional insample exogenous of shape (t, n_x). Currently ignored. + X_future (array-like): Optional exogenous of shape (h, n_x). Currently ignored. + level (List[float]): Confidence levels (0-100) for prediction intervals. Not yet supported. + fitted (bool): Whether or not to return insample predictions. + + Returns: + dict: Dictionary with entry `mean` for point predictions and optionally `fitted` insample predictions. + """ + self._no_pi(level) + y = _ensure_float(y) + mod = ucm_model(y, season_length=self.season_length) + fcst = ucm_forecast(mod, h) + res = {"mean": fcst["mean"]} + if fitted: + res["fitted"] = fcst["fitted"] + return res + + class SklearnModel(_TS): r"""scikit-learn model wrapper diff --git a/python/statsforecast/ucm.py b/python/statsforecast/ucm.py index fcdaee32b..7d0cf9c66 100644 --- a/python/statsforecast/ucm.py +++ b/python/statsforecast/ucm.py @@ -1,549 +1,187 @@ -""" -Unobserved Components Model (UCM) wrapper for statsforecast. - -This module provides a wrapper around statsmodels.tsa.statespace.structural.UnobservedComponents -to make it compatible with the statsforecast API. - -UCM decomposes a time series into: -- Level (intercept that varies over time) -- Trend (slope that varies over time) -- Seasonal (periodic patterns) -- Cycle (longer-term oscillations) -- Irregular (noise) - -References: - - Harvey, A. C. (1989). Forecasting, Structural Time Series Models and the Kalman Filter. - - statsmodels documentation: https://www.statsmodels.org/stable/generated/statsmodels.tsa.statespace.structural.UnobservedComponents.html -""" - -import warnings -from typing import Any, Dict, List, Optional, Union +__all__ = ['ucm_model', 'ucm_forecast'] import numpy as np - - -def _calculate_sigma(residuals: np.ndarray, n: int) -> float: - """Calculate standard error of residuals.""" - return np.sqrt(np.sum(residuals ** 2) / max(n - 1, 1)) - - -def _add_fitted_pi( - res: Dict[str, Any], results: Any, level: List[int] -) -> Dict[str, Any]: - """Add Kalman-filter-based prediction intervals to fitted values.""" - prediction = results.get_prediction(start=0, end=results.nobs - 1) - for lv in level: - alpha = 1 - lv / 100 - ci = prediction.conf_int(alpha=alpha) - if hasattr(ci, "iloc"): - res[f"lo-{lv}"] = ci.iloc[:, 0].values - res[f"hi-{lv}"] = ci.iloc[:, 1].values - else: - res[f"lo-{lv}"] = np.asarray(ci[:, 0]) - res[f"hi-{lv}"] = np.asarray(ci[:, 1]) - return res - - -class UCM: - r"""Unobserved Components Model (UCM). - - Also known as Structural Time Series Model. Decomposes a univariate time series - into trend, seasonal, cyclical, and irregular components using state space methods - and the Kalman filter. - - This is a wrapper around `statsmodels.tsa.statespace.structural.UnobservedComponents`. - - Args: - level (Union[bool, str], default="local level"): - Level component specification. Can be: - - False: No level component - - True or "local level": Random walk level - - "local linear trend" or "lltrend": Level + trend (both stochastic) - - "random walk with drift" or "rwdrift": Random walk + deterministic drift - - "smooth trend" or "strend": Smooth trend (integrated random walk) - - "random trend" or "rtrend": Random trend - See statsmodels documentation for full list. - trend (bool, default=False): - Whether to include a trend component. Only used if `level` is bool. - seasonal (Optional[int], default=None): - Period of seasonal component. If None, no seasonal component. - cycle (bool, default=False): - Whether to include a cycle component. - autoregressive (Optional[int], default=None): - Order of autoregressive component. If None, no AR component. - irregular (bool, default=True): - Whether to include an irregular (noise) component. - stochastic_level (bool, default=True): - Whether level component is stochastic (if included). - stochastic_trend (bool, default=True): - Whether trend component is stochastic (if included). - stochastic_seasonal (bool, default=True): - Whether seasonal component is stochastic (if included). - stochastic_cycle (bool, default=False): - Whether cycle component is stochastic (if included). - damped_cycle (bool, default=False): - Whether cycle component is damped. - cycle_period_bounds (Optional[tuple], default=None): - Bounds on cycle period (lower, upper). Default is (1.5, 12). - use_exact_diffuse (bool, default=False): - Whether to use exact diffuse initialization. - fit_method (str, default="lbfgs"): - Optimization method for fitting. - maxiter (int, default=500): - Maximum iterations for fitting. - alias (Optional[str], default=None): - Custom name for the model. - - Examples: - >>> from statsforecast.models import UCM - >>> import numpy as np - >>> # Local level model (random walk) - >>> model = UCM(level='local level') - >>> y = np.cumsum(np.random.randn(100)) + 50 - >>> model.fit(y) - >>> forecast = model.predict(h=10) - >>> - >>> # Local linear trend with seasonal - >>> model = UCM(level='local linear trend', seasonal=12) - >>> model.fit(y) - >>> forecast = model.predict(h=12, level=[90, 95]) - - References: - - Harvey, A. C. (1989). Forecasting, Structural Time Series Models and the Kalman Filter. - - https://www.statsmodels.org/stable/generated/statsmodels.tsa.statespace.structural.UnobservedComponents.html - """ - - uses_exog = True - - def __init__( - self, - level: Union[bool, str] = "local level", - trend: bool = False, - seasonal: Optional[int] = None, - cycle: bool = False, - autoregressive: Optional[int] = None, - irregular: bool = True, - stochastic_level: bool = True, - stochastic_trend: bool = True, - stochastic_seasonal: bool = True, - stochastic_cycle: bool = False, - damped_cycle: bool = False, - cycle_period_bounds: Optional[tuple] = None, - use_exact_diffuse: bool = False, - fit_method: str = "lbfgs", - maxiter: int = 500, - alias: Optional[str] = None, - ): - self.level = level - self.trend = trend - self.seasonal = seasonal - self.cycle = cycle - self.autoregressive = autoregressive - self.irregular = irregular - self.stochastic_level = stochastic_level - self.stochastic_trend = stochastic_trend - self.stochastic_seasonal = stochastic_seasonal - self.stochastic_cycle = stochastic_cycle - self.damped_cycle = damped_cycle - self.cycle_period_bounds = cycle_period_bounds - self.use_exact_diffuse = use_exact_diffuse - self.fit_method = fit_method - self.maxiter = maxiter - self.alias = alias if alias is not None else self._default_alias() - - # Will be set during fit - self.model_: Optional[Dict[str, Any]] = None - - def _default_alias(self) -> str: - """Generate default alias based on model specification.""" - parts = ["UCM"] - if isinstance(self.level, str): - parts.append(self.level.replace(" ", "_")) - elif self.level: - parts.append("level") - if self.trend: - parts.append("trend") - if self.seasonal: - parts.append(f"s{self.seasonal}") - if self.cycle: - parts.append("cycle") - if self.autoregressive: - parts.append(f"ar{self.autoregressive}") - return "_".join(parts) - - def __repr__(self) -> str: - return self.alias - - def new(self): - """Create a copy of the model.""" - b = type(self).__new__(type(self)) - b.__dict__.update(self.__dict__) - return b - - def _build_model(self, y: np.ndarray, X: Optional[np.ndarray] = None): - """Build statsmodels UCM model.""" - from statsmodels.tsa.statespace.structural import UnobservedComponents - - kwargs = { - "endog": y, - "stochastic_seasonal": self.stochastic_seasonal, - "stochastic_cycle": self.stochastic_cycle, - "damped_cycle": self.damped_cycle, - "use_exact_diffuse": self.use_exact_diffuse, - } - - # Handle level/trend specification - # When using string specification, don't pass irregular/stochastic params - # as they may conflict with the string spec defaults - if isinstance(self.level, str): - kwargs["level"] = self.level - dropped = [] - if not self.stochastic_level: - dropped.append("stochastic_level") - if not self.stochastic_trend: - dropped.append("stochastic_trend") - if self.trend: - dropped.append("trend") - if dropped: - warnings.warn( - f"Parameters {dropped} are ignored when `level` is a string " - f"specification ('{self.level}'). Use a bool `level` to control " - "these parameters explicitly.", - UserWarning, - stacklevel=3, - ) - # Only pass irregular if it's explicitly False (to override string default) - if not self.irregular: - kwargs["irregular"] = False - else: - kwargs["level"] = self.level - kwargs["trend"] = self.trend - kwargs["irregular"] = self.irregular - kwargs["stochastic_level"] = self.stochastic_level - kwargs["stochastic_trend"] = self.stochastic_trend - - # Optional components - if self.seasonal is not None: - kwargs["seasonal"] = self.seasonal - - if self.cycle: - kwargs["cycle"] = True - if self.cycle_period_bounds is not None: - kwargs["cycle_period_bounds"] = self.cycle_period_bounds - - if self.autoregressive is not None: - kwargs["autoregressive"] = self.autoregressive - - if X is not None: - kwargs["exog"] = X - - return UnobservedComponents(**kwargs) - - def fit( - self, - y: np.ndarray, - X: Optional[np.ndarray] = None, - ) -> "UCM": - r"""Fit the UCM model. - - Args: - y (numpy.array): Time series of shape (t,). - X (numpy.array, optional): Exogenous variables of shape (t, n_x). - - Returns: - UCM: Fitted UCM object. - """ - y = np.asarray(y, dtype=np.float64) - if X is not None: - X = np.asarray(X, dtype=np.float64) - - mod = self._build_model(y, X) - - try: - res = mod.fit(method=self.fit_method, maxiter=self.maxiter, disp=False) - except (np.linalg.LinAlgError, ValueError) as original_exc: - warnings.warn( - f"{self.fit_method} optimizer failed ({original_exc}); " - "falling back to powell.", - UserWarning, - stacklevel=2, - ) - try: - res = mod.fit(method="powell", maxiter=self.maxiter, disp=False) - except Exception: - raise original_exc - - # Handle both pandas and numpy returns - fitted_vals = res.fittedvalues - if hasattr(fitted_vals, 'values'): - fitted_vals = fitted_vals.values - fitted_vals = np.asarray(fitted_vals) - - self.model_ = { - "model": mod, - "results": res, - "fitted": fitted_vals, - "y": y, - "X": X, - } - - # Calculate sigma from residuals - residuals = y - fitted_vals - self.model_["sigma"] = _calculate_sigma(residuals, y.size) - - return self - - def predict( - self, - h: int, - X: Optional[np.ndarray] = None, - level: Optional[List[int]] = None, - ) -> Dict[str, Any]: - r"""Predict with fitted UCM. - - Args: - h (int): Forecast horizon. - X (numpy.array, optional): Future exogenous variables of shape (h, n_x). - level (List[int], optional): Confidence levels (0-100) for prediction intervals. - - Returns: - dict: Dictionary with entries `mean` for point predictions and `level_*` for intervals. - """ - if self.model_ is None: - raise ValueError("Model has not been fitted. Call fit() first.") - - res = self.model_["results"] - - # Get forecast - if X is not None: - X = np.asarray(X, dtype=np.float64) - forecast = res.get_forecast(steps=h, exog=X) - else: - forecast = res.get_forecast(steps=h) - - # Handle both numpy array and pandas Series returns - pred_mean = forecast.predicted_mean - if hasattr(pred_mean, 'values'): - pred_mean = pred_mean.values - result = {"mean": np.asarray(pred_mean)} - - # Add prediction intervals - if level is not None: - level = sorted(level) - for lv in level: - alpha = 1 - lv / 100 - ci = forecast.conf_int(alpha=alpha) - # Handle both numpy array and pandas DataFrame returns - if hasattr(ci, 'iloc'): - result[f"lo-{lv}"] = ci.iloc[:, 0].values - result[f"hi-{lv}"] = ci.iloc[:, 1].values - else: - result[f"lo-{lv}"] = np.asarray(ci[:, 0]) - result[f"hi-{lv}"] = np.asarray(ci[:, 1]) - - return result - - def predict_in_sample( - self, - level: Optional[List[int]] = None, - ) -> Dict[str, Any]: - r"""Access fitted UCM in-sample predictions. - - Args: - level (List[int], optional): Confidence levels (0-100) for prediction intervals. - - Returns: - dict: Dictionary with entries `fitted` for point predictions. - """ - if self.model_ is None: - raise ValueError("Model has not been fitted. Call fit() first.") - - result = {"fitted": self.model_["fitted"]} - - if level is not None: - level = sorted(level) - result = _add_fitted_pi( - res=result, results=self.model_["results"], level=level - ) - - return result - - def forecast( - self, - y: np.ndarray, - h: int, - X: Optional[np.ndarray] = None, - X_future: Optional[np.ndarray] = None, - level: Optional[List[int]] = None, - fitted: bool = False, - ) -> Dict[str, Any]: - r"""Memory-efficient UCM predictions. - - This method avoids memory burden from object storage. - It is analogous to `fit` + `predict` without storing information. - - Args: - y (numpy.array): Time series of shape (t,). - h (int): Forecast horizon. - X (numpy.array, optional): In-sample exogenous variables of shape (t, n_x). - X_future (numpy.array, optional): Future exogenous variables of shape (h, n_x). - level (List[int], optional): Confidence levels (0-100) for prediction intervals. - fitted (bool, default=False): Whether to return in-sample predictions. - - Returns: - dict: Dictionary with entries `mean` for point predictions and optionally `fitted`. - """ - y = np.asarray(y, dtype=np.float64) - if X is not None: - X = np.asarray(X, dtype=np.float64) - if X_future is not None: - X_future = np.asarray(X_future, dtype=np.float64) - - mod = self._build_model(y, X) - - try: - res = mod.fit(method=self.fit_method, maxiter=self.maxiter, disp=False) - except (np.linalg.LinAlgError, ValueError) as original_exc: - warnings.warn( - f"{self.fit_method} optimizer failed ({original_exc}); " - "falling back to powell.", - UserWarning, - stacklevel=2, - ) - try: - res = mod.fit(method="powell", maxiter=self.maxiter, disp=False) - except Exception: - raise original_exc - - # Get forecast - if X_future is not None: - forecast = res.get_forecast(steps=h, exog=X_future) +from scipy.optimize import minimize +from typing import Dict, Tuple + +# Diffuse initialization: initial state mean is 0 and the initial state +# covariance is a large multiple of the identity. +_DIFFUSE_VARIANCE = 1e6 + +def _build_matrices( + season_length: int, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Build the time-invariant state space matrices for the UCM.""" + s = int(season_length) + has_seasonal = s > 1 + n_seas = s - 1 if has_seasonal else 0 + k = 2 + n_seas + + # Observation matrix: y_t = mu_t + gamma_t + eps_t. + Z = np.zeros((1, k)) + Z[0, 0] = 1.0 # level + if has_seasonal: + Z[0, 2] = 1.0 # current seasonal effect + + # Transition matrix. + T = np.zeros((k, k)) + T[0, 0] = 1.0 # mu_{t+1} = mu_t + beta_t + T[0, 1] = 1.0 + T[1, 1] = 1.0 # beta_{t+1} = beta_t + if has_seasonal: + # gamma_t = -(gamma_{t-1} + ... + gamma_{t-s+1}) + omega_t + T[2, 2 : 2 + n_seas] = -1.0 + # The remaining seasonal states are lagged copies (companion form). + for i in range(1, n_seas): + T[2 + i, 2 + i - 1] = 1.0 + + # State-noise selection matrix: shocks enter the level, slope and (if + # present) the current seasonal state. + n_shocks = 3 if has_seasonal else 2 + R = np.zeros((k, n_shocks)) + R[0, 0] = 1.0 # level shock + R[1, 1] = 1.0 # slope shock + if has_seasonal: + R[2, 2] = 1.0 # seasonal shock + + return Z, T, R + + +def kalman_filter( + y: np.ndarray, + Z: np.ndarray, + T: np.ndarray, + R: np.ndarray, + Q: np.ndarray, + H: float, + a0: np.ndarray, + P0: np.ndarray, +) -> Tuple[float, np.ndarray, np.ndarray]: + """Run a univariate Kalman filter.""" + n = y.shape[0] + z = Z[0] + RQRt = R @ Q @ R.T + + a_pred = a0.astype(np.float64).copy() + P_pred = P0.astype(np.float64).copy() + + one_step_pred = np.empty(n, dtype=np.float64) + loglik = 0.0 + log_2pi = np.log(2.0 * np.pi) + + for t in range(n): + # Forecast for the current observation. + y_hat = z @ a_pred + one_step_pred[t] = y_hat + v = y[t] - y_hat + Pz = P_pred @ z + F = z @ Pz + H + + # Update step (skip if F is not usable, e.g. during diffuse start). + if np.isfinite(F) and F > 0: + loglik += -0.5 * (log_2pi + np.log(F) + v * v / F) + K = Pz / F + a_filt = a_pred + K * v + P_filt = P_pred - np.outer(K, Pz) else: - forecast = res.get_forecast(steps=h) - - # Handle both numpy array and pandas Series returns - pred_mean = forecast.predicted_mean - if hasattr(pred_mean, 'values'): - pred_mean = pred_mean.values - result = {"mean": np.asarray(pred_mean)} - - if fitted: - fitted_vals = res.fittedvalues - if hasattr(fitted_vals, 'values'): - fitted_vals = fitted_vals.values - result["fitted"] = np.asarray(fitted_vals) - - # Add prediction intervals - if level is not None: - level = sorted(level) - for lv in level: - alpha = 1 - lv / 100 - ci = forecast.conf_int(alpha=alpha) - # Handle both numpy array and pandas DataFrame returns - if hasattr(ci, 'iloc'): - result[f"lo-{lv}"] = ci.iloc[:, 0].values - result[f"hi-{lv}"] = ci.iloc[:, 1].values - else: - result[f"lo-{lv}"] = np.asarray(ci[:, 0]) - result[f"hi-{lv}"] = np.asarray(ci[:, 1]) - - return result - - def get_components(self) -> Dict[str, np.ndarray]: - r"""Get decomposed components from fitted model. - - Returns: - dict: Dictionary with available components (level, trend, seasonal, cycle). - """ - if self.model_ is None: - raise ValueError("Model has not been fitted. Call fit() first.") - - res = self.model_["results"] - components = {} - - # Try to extract each component - if hasattr(res, "level") and res.level is not None: - components["level"] = res.level.smoothed - - if hasattr(res, "trend") and res.trend is not None: - components["trend"] = res.trend.smoothed - - if hasattr(res, "seasonal") and res.seasonal is not None: - components["seasonal"] = res.seasonal.smoothed - - if hasattr(res, "cycle") and res.cycle is not None: - components["cycle"] = res.cycle.smoothed - - if hasattr(res, "autoregressive") and res.autoregressive is not None: - components["autoregressive"] = res.autoregressive.smoothed - - return components - - -# Convenience aliases for common UCM specifications -class LocalLevel(UCM): - """Local Level model (random walk). - - The simplest UCM model: y_t = level_t + eps_t, where level follows a random walk. - """ - - def __init__( - self, - seasonal: Optional[int] = None, - fit_method: str = "lbfgs", - maxiter: int = 500, - alias: Optional[str] = None, - ): - super().__init__( - level="local level", - seasonal=seasonal, - fit_method=fit_method, - maxiter=maxiter, - alias=alias or ("LocalLevel" + (f"_s{seasonal}" if seasonal else "")), - ) - - -class LocalLinearTrend(UCM): - """Local Linear Trend model. - - Level and trend both follow random walks: - - level_t = level_{t-1} + trend_{t-1} + eta_t - - trend_t = trend_{t-1} + zeta_t - """ - - def __init__( - self, - seasonal: Optional[int] = None, - fit_method: str = "lbfgs", - maxiter: int = 500, - alias: Optional[str] = None, - ): - super().__init__( - level="local linear trend", - seasonal=seasonal, - fit_method=fit_method, - maxiter=maxiter, - alias=alias or ("LocalLinearTrend" + (f"_s{seasonal}" if seasonal else "")), - ) - - -class SmoothTrend(UCM): - """Smooth Trend model (integrated random walk). - - The trend is an integrated random walk, producing smoother forecasts. - """ - - def __init__( - self, - seasonal: Optional[int] = None, - fit_method: str = "lbfgs", - maxiter: int = 500, - alias: Optional[str] = None, - ): - super().__init__( - level="smooth trend", - seasonal=seasonal, - fit_method=fit_method, - maxiter=maxiter, - alias=alias or ("SmoothTrend" + (f"_s{seasonal}" if seasonal else "")), - ) + a_filt = a_pred + P_filt = P_pred + + # Predict the next state. + a_pred = T @ a_filt + P_pred = T @ P_filt @ T.T + RQRt + + return loglik, a_filt, one_step_pred + + +def _make_QH(theta: np.ndarray, has_seasonal: bool) -> Tuple[np.ndarray, float]: + """Build the state-noise covariance Q and observation variance H.""" + H = theta[0] + if has_seasonal: + Q = np.diag(theta[1:4]) + else: + Q = np.diag(theta[1:3]) + return Q, H + + +def _nll( + theta: np.ndarray, + y: np.ndarray, + Z: np.ndarray, + T: np.ndarray, + R: np.ndarray, + a0: np.ndarray, + P0: np.ndarray, + has_seasonal: bool, +) -> float: + """Negative log-likelihood for a variance vector theta.""" + Q, H = _make_QH(theta, has_seasonal) + loglik, _, _ = kalman_filter(y, Z, T, R, Q, H, a0, P0) + if not np.isfinite(loglik): + return 1e10 + return -loglik + + +def ucm_model(y: np.ndarray, season_length: int = 1) -> Dict: + """Fit the minimal UCM by maximum likelihood.""" + y = np.asarray(y, dtype=np.float64) + s = int(season_length) + has_seasonal = s > 1 + + Z, T, R = _build_matrices(s) + k = T.shape[0] + a0 = np.zeros(k) + P0 = _DIFFUSE_VARIANCE * np.eye(k) + + # Starting values and bounds: variances scaled by the series variance for + # numerical conditioning across different data scales. + var_y = float(np.var(y)) + if var_y <= 0 or not np.isfinite(var_y): + var_y = 1.0 + n_params = 4 if has_seasonal else 3 + x0 = np.full(n_params, var_y / 10.0) + bounds = [(1e-8 * var_y, None)] * n_params + + result = minimize( + _nll, + x0, + args=(y, Z, T, R, a0, P0, has_seasonal), + method="L-BFGS-B", + bounds=bounds, + ) + params = result.x + + Q, H = _make_QH(params, has_seasonal) + _, a_n, fitted = kalman_filter(y, Z, T, R, Q, H, a0, P0) + + residuals = y - fitted + sigma = np.sqrt(np.sum(residuals**2) / max(len(y) - 1, 1)) + + return { + "params": params, + "Z": Z, + "T": T, + "R": R, + "Q": Q, + "H": H, + "a_n": a_n, + "fitted": fitted, + "sigma": sigma, + "season_length": s, + "message": result.message, + } + + +def ucm_forecast(mod: Dict, h: int) -> Dict: + """Produce h-step-ahead point forecasts from a fitted UCM.""" + Z = mod["Z"] + T = mod["T"] + z = Z[0] + a = mod["a_n"].astype(np.float64).copy() + + mean = np.empty(h, dtype=np.float64) + for i in range(h): + a = T @ a + mean[i] = z @ a + + return {"mean": mean, "fitted": mod["fitted"]} diff --git a/tests/test_ucm.py b/tests/test_ucm.py index 5bd5b0969..582600ef0 100644 --- a/tests/test_ucm.py +++ b/tests/test_ucm.py @@ -1,249 +1,270 @@ -""" -Tests for UCM (Unobserved Components Model) module. -""" -import numpy as np -import pytest - - -def test_ucm_import(): - """Test that UCM classes can be imported.""" - from statsforecast.models import UCM, LocalLevel, LocalLinearTrend, SmoothTrend - assert UCM is not None - assert LocalLevel is not None - assert LocalLinearTrend is not None - assert SmoothTrend is not None - - -def test_ucm_basic_fit_predict(): - """Test basic fit and predict workflow.""" - from statsforecast.models import UCM +"""Tests for the UCM (Unobserved Components Model).""" - np.random.seed(42) - y = np.cumsum(np.random.randn(100)) + 50 +import warnings - model = UCM(level='local level') - model.fit(y) - - forecast = model.predict(h=10) +import numpy as np +import pytest - assert 'mean' in forecast - assert len(forecast['mean']) == 10 - assert not np.any(np.isnan(forecast['mean'])) +from statsforecast.ucm import ( + _build_matrices, + kalman_filter, + ucm_forecast, + ucm_model, +) +warnings.simplefilter("ignore") -def test_ucm_with_seasonal(): - """Test UCM with seasonal component.""" - from statsforecast.models import UCM - np.random.seed(42) +@pytest.fixture +def trend_series(): + """Local-linear-trend series (no seasonality).""" + np.random.seed(0) n = 120 - t = np.arange(n) - seasonal = 10 * np.sin(2 * np.pi * t / 12) - y = 50 + seasonal + np.random.randn(n) * 2 - - model = UCM(level='local level', seasonal=12) - model.fit(y) + return 50.0 + 0.5 * np.arange(n) + np.cumsum(np.random.randn(n)) - forecast = model.predict(h=12) - assert len(forecast['mean']) == 12 +@pytest.fixture +def seasonal_series(): + """Trend + seasonal (period 12) series.""" + np.random.seed(0) + n = 144 + t = np.arange(n) + return ( + 20.0 + + 0.3 * t + + 10.0 * np.sin(2 * np.pi * t / 12) + + np.random.randn(n) + ) -def test_ucm_prediction_intervals(): - """Test prediction intervals.""" +# --------------------------------------------------------------------------- +# Computational layer +# --------------------------------------------------------------------------- +def test_build_matrices_trend_only(): + Z, T, R = _build_matrices(season_length=1) + assert Z.shape == (1, 2) + assert T.shape == (2, 2) + assert R.shape == (2, 2) + np.testing.assert_array_equal(Z[0], np.array([1.0, 0.0])) + # Local linear trend transition. + np.testing.assert_array_equal(T, np.array([[1.0, 1.0], [0.0, 1.0]])) + + +def test_build_matrices_seasonal(): + s = 12 + Z, T, R = _build_matrices(season_length=s) + k = 2 + (s - 1) + assert Z.shape == (1, k) + assert T.shape == (k, k) + assert R.shape == (k, 3) + # Picks level and current seasonal effect. + assert Z[0, 0] == 1.0 + assert Z[0, 2] == 1.0 + # Seasonal effects sum-to-zero recursion. + np.testing.assert_array_equal(T[2, 2 : 2 + (s - 1)], np.full(s - 1, -1.0)) + + +def test_kalman_filter_loglik_finite(trend_series): + y = trend_series + Z, T, R = _build_matrices(season_length=1) + Q = np.diag([1.0, 0.1]) + H = 1.0 + k = T.shape[0] + a0 = np.zeros(k) + P0 = 1e6 * np.eye(k) + + loglik, a_filt, one_step = kalman_filter(y, Z, T, R, Q, H, a0, P0) + + assert np.isfinite(loglik) + assert a_filt.shape == (k,) + assert one_step.shape == (len(y),) + assert np.all(np.isfinite(one_step)) + + +def test_ucm_model_params_trend_only(trend_series): + mod = ucm_model(trend_series, season_length=1) + # [sigma2_eps, sigma2_eta, sigma2_zeta] + assert len(mod["params"]) == 3 + assert np.all(mod["params"] >= 0) + assert len(mod["fitted"]) == len(trend_series) + assert mod["season_length"] == 1 + + +def test_ucm_model_params_seasonal(seasonal_series): + mod = ucm_model(seasonal_series, season_length=12) + # adds sigma2_omega + assert len(mod["params"]) == 4 + assert len(mod["fitted"]) == len(seasonal_series) + assert mod["season_length"] == 12 + + +def test_ucm_forecast_shape(trend_series): + mod = ucm_model(trend_series, season_length=1) + h = 10 + fcst = ucm_forecast(mod, h) + + assert isinstance(fcst, dict) + assert "mean" in fcst + assert "fitted" in fcst + assert len(fcst["mean"]) == h + assert np.all(np.isfinite(fcst["mean"])) + + +def test_ucm_forecast_trend_direction(trend_series): + """An upward local linear trend should extrapolate upward.""" + mod = ucm_model(trend_series, season_length=1) + fcst = ucm_forecast(mod, 12) + assert fcst["mean"][-1] > fcst["mean"][0] + + +# --------------------------------------------------------------------------- +# Model class API +# --------------------------------------------------------------------------- +def test_ucm_import(): from statsforecast.models import UCM - np.random.seed(42) - y = np.cumsum(np.random.randn(100)) + 50 - - model = UCM(level='local level') - model.fit(y) - - forecast = model.predict(h=10, level=[80, 95]) - - assert 'lo-80' in forecast - assert 'hi-80' in forecast - assert 'lo-95' in forecast - assert 'hi-95' in forecast - - # Check that wider CI contains narrower CI - assert np.all(forecast['lo-95'] <= forecast['lo-80']) - assert np.all(forecast['hi-95'] >= forecast['hi-80']) + assert UCM is not None -def test_ucm_forecast_method(): - """Test the forecast() method (fit + predict in one call).""" +def test_ucm_fit_predict(trend_series): from statsforecast.models import UCM - np.random.seed(42) - y = np.cumsum(np.random.randn(100)) + 50 + model = UCM(season_length=1) + model.fit(trend_series) + fcst = model.predict(h=10) - model = UCM(level='local linear trend') - result = model.forecast(y=y, h=10, level=[90], fitted=True) + assert "mean" in fcst + assert len(fcst["mean"]) == 10 + assert not np.any(np.isnan(fcst["mean"])) - assert 'mean' in result - assert 'fitted' in result - assert 'lo-90' in result - assert 'hi-90' in result - assert len(result['mean']) == 10 - assert len(result['fitted']) == 100 - -def test_ucm_predict_in_sample(): - """Test in-sample predictions.""" +def test_ucm_seasonal_fit_predict(seasonal_series): from statsforecast.models import UCM - np.random.seed(42) - y = np.cumsum(np.random.randn(100)) + 50 - - model = UCM(level='local level') - model.fit(y) + model = UCM(season_length=12) + model.fit(seasonal_series) + fcst = model.predict(h=12) - insample = model.predict_in_sample(level=[95]) + assert len(fcst["mean"]) == 12 - assert 'fitted' in insample - assert len(insample['fitted']) == 100 - assert 'lo-95' in insample - assert 'hi-95' in insample - -def test_ucm_get_components(): - """Test component extraction.""" +def test_ucm_forecast_method(trend_series): from statsforecast.models import UCM - np.random.seed(42) - n = 120 - t = np.arange(n) - y = 50 + 0.3 * t + 8 * np.sin(2 * np.pi * t / 12) + np.random.randn(n) * 2 - - model = UCM(level='local level', seasonal=12) - model.fit(y) + model = UCM(season_length=1) + res = model.forecast(y=trend_series, h=10, fitted=True) - components = model.get_components() + assert "mean" in res + assert "fitted" in res + assert len(res["mean"]) == 10 + assert len(res["fitted"]) == len(trend_series) - assert 'level' in components - assert 'seasonal' in components - assert len(components['level']) == n - assert len(components['seasonal']) == n - -def test_ucm_with_exogenous(): - """Test UCM with exogenous variables.""" +def test_ucm_predict_in_sample(trend_series): from statsforecast.models import UCM - np.random.seed(42) - n = 100 - y = np.cumsum(np.random.randn(n)) + 50 - X = np.column_stack([np.ones(n), np.arange(n)]) - X_future = np.column_stack([np.ones(10), np.arange(n, n+10)]) - - model = UCM(level='local level') - model.fit(y, X=X) - - forecast = model.predict(h=10, X=X_future) - - assert len(forecast['mean']) == 10 + model = UCM(season_length=1) + model.fit(trend_series) + insample = model.predict_in_sample() + assert "fitted" in insample + assert len(insample["fitted"]) == len(trend_series) -def test_local_level(): - """Test LocalLevel convenience class.""" - from statsforecast.models import LocalLevel - np.random.seed(42) - y = np.cumsum(np.random.randn(100)) + 50 - - model = LocalLevel() - model.fit(y) - - assert 'LocalLevel' in model.alias - - forecast = model.predict(h=5) - assert len(forecast['mean']) == 5 - - -def test_local_linear_trend(): - """Test LocalLinearTrend convenience class.""" - from statsforecast.models import LocalLinearTrend - - np.random.seed(42) - n = 100 - y = 50 + 0.5 * np.arange(n) + np.random.randn(n) * 2 - - model = LocalLinearTrend() - model.fit(y) - - assert 'LocalLinearTrend' in model.alias - - forecast = model.predict(h=5) - assert len(forecast['mean']) == 5 - - -def test_smooth_trend(): - """Test SmoothTrend convenience class.""" - from statsforecast.models import SmoothTrend - - np.random.seed(42) - y = np.cumsum(np.random.randn(100)) + 50 - - model = SmoothTrend() - model.fit(y) +def test_ucm_level_not_implemented(trend_series): + from statsforecast.models import UCM - assert 'SmoothTrend' in model.alias + model = UCM(season_length=1) + model.fit(trend_series) + with pytest.raises(NotImplementedError): + model.predict(h=5, level=[95]) + with pytest.raises(NotImplementedError): + model.forecast(y=trend_series, h=5, level=[95]) def test_ucm_alias(): - """Test custom alias.""" from statsforecast.models import UCM - model = UCM(level='local level', alias='MyCustomModel') - assert str(model) == 'MyCustomModel' - assert model.alias == 'MyCustomModel' + model = UCM(season_length=12, alias="MyUCM") + assert str(model) == "MyUCM" + assert model.alias == "MyUCM" def test_ucm_new(): - """Test the new() method for copying.""" from statsforecast.models import UCM - model1 = UCM(level='local level', seasonal=12) + model1 = UCM(season_length=12) model2 = model1.new() - assert model2.level == model1.level - assert model2.seasonal == model1.seasonal + assert model2.season_length == model1.season_length assert model2 is not model1 +# --------------------------------------------------------------------------- +# StatsForecast integration +# --------------------------------------------------------------------------- def test_statsforecast_integration(): - """Test integration with StatsForecast.""" - pytest.importorskip("statsforecast") - import pandas as pd from statsforecast import StatsForecast - from statsforecast.models import LocalLevel, LocalLinearTrend + from statsforecast.models import UCM - np.random.seed(42) + np.random.seed(0) n = 60 - dates = pd.date_range('2020-01-01', periods=n, freq='MS') + dates = pd.date_range("2020-01-01", periods=n, freq="MS") y = 50 + 0.3 * np.arange(n) + np.random.randn(n) * 2 + df = pd.DataFrame({"unique_id": ["s1"] * n, "ds": dates, "y": y}) - df = pd.DataFrame({ - 'unique_id': ['series1'] * n, - 'ds': dates, - 'y': y - }) + sf = StatsForecast(models=[UCM(season_length=12)], freq="MS", n_jobs=1) + sf.fit(df) + fcst = sf.predict(h=6) - sf = StatsForecast( - models=[LocalLevel(), LocalLinearTrend()], - freq='MS', - n_jobs=1, - ) + assert len(fcst) == 6 + assert "UCM" in fcst.columns - sf.fit(df) - forecast = sf.predict(h=6) - assert len(forecast) == 6 - assert 'LocalLevel' in forecast.columns - assert 'LocalLinearTrend' in forecast.columns +def test_statsforecast_cross_validation_fitted(): + import pandas as pd + from statsforecast import StatsForecast + from statsforecast.models import UCM + + np.random.seed(0) + n = 96 + dates = pd.date_range("2015-01-01", periods=n, freq="MS") + t = np.arange(n) + y = 30 + 0.2 * t + 8 * np.sin(2 * np.pi * t / 12) + np.random.randn(n) + df = pd.DataFrame({"unique_id": ["s1"] * n, "ds": dates, "y": y}) + + sf = StatsForecast(models=[UCM(season_length=12)], freq="MS", n_jobs=1) + cv = sf.cross_validation(df=df, h=6, n_windows=2, fitted=True) + + assert "UCM" in cv.columns + # Fitted values from cross-validation should be accessible. + fitted_cv = sf.cross_validation_fitted_values() + assert "UCM" in fitted_cv.columns + + +# --------------------------------------------------------------------------- +# Numerical correctness vs statsmodels +# --------------------------------------------------------------------------- +def test_ucm_matches_statsmodels(seasonal_series): + """Forecasts should agree with statsmodels' UnobservedComponents.""" + sm = pytest.importorskip("statsmodels.api") + + y = seasonal_series + h = 12 + + sf_mod = ucm_model(y, season_length=12) + sf_fcst = ucm_forecast(sf_mod, h)["mean"] + + sm_mod = sm.tsa.UnobservedComponents( + y, level="local linear trend", seasonal=12, stochastic_seasonal=True + ) + sm_res = sm_mod.fit(method="lbfgs", disp=False) + sm_fcst = np.asarray(sm_res.forecast(h)) + + # Point forecasts should be close in scale and direction. + np.testing.assert_allclose(sf_fcst, sm_fcst, rtol=0.15, atol=0.15 * np.std(y)) -if __name__ == '__main__': - pytest.main([__file__, '-v']) \ No newline at end of file +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 0523e41cf7a6c7dacf4d1d4512c73f199a505f0c Mon Sep 17 00:00:00 2001 From: MMenchero Date: Sun, 31 May 2026 23:29:22 -0600 Subject: [PATCH 2/8] add prediction intervals to UCM --- python/statsforecast/models.py | 50 ++++++++++++-------- python/statsforecast/ucm.py | 20 +++++--- tests/test_ucm.py | 84 ++++++++++++++++++++++++---------- 3 files changed, 105 insertions(+), 49 deletions(-) diff --git a/python/statsforecast/models.py b/python/statsforecast/models.py index 52413c207..7f9ec5960 100644 --- a/python/statsforecast/models.py +++ b/python/statsforecast/models.py @@ -6307,7 +6307,7 @@ class UCM(_TS): Args: season_length (int): Number of observations per seasonal cycle. Ex: 12 for monthly data. Use 1 for a trend-only model. alias (str): Custom name of the model. - prediction_intervals (Optional[ConformalIntervals]): Information to compute conformal prediction intervals. + prediction_intervals (Optional[ConformalIntervals]): Information to compute conformal prediction intervals. By default, the model computes parametric (Gaussian) prediction intervals. """ def __init__( @@ -6320,12 +6320,6 @@ def __init__( self.alias = alias self.prediction_intervals = prediction_intervals - def _no_pi(self, level: Optional[List[int]]): - if level is not None: - raise NotImplementedError( - "Prediction intervals are not yet supported for UCM." - ) - def fit(self, y: np.ndarray, X: Optional[np.ndarray] = None): r"""Fit the UCM model. @@ -6340,6 +6334,8 @@ def fit(self, y: np.ndarray, X: Optional[np.ndarray] = None): """ y = _ensure_float(y) self.model_ = ucm_model(y, season_length=self.season_length) + self.model_["residuals"] = y - self.model_["fitted"] + self._store_cs(y, X) return self def predict( @@ -6350,26 +6346,36 @@ def predict( Args: h (int): Forecast horizon. X (array-like): Optional exogenous of shape (h, n_x). Currently ignored. - level (List[float]): Confidence levels (0-100) for prediction intervals. Not yet supported. + level (List[float]): Confidence levels (0-100) for prediction intervals. Returns: - dict: Dictionary with entry `mean` for point predictions. + dict: Dictionary with entries `mean` for point predictions and `level_*` for probabilistic predictions. """ - self._no_pi(level) fcst = ucm_forecast(self.model_, h) - return {"mean": fcst["mean"]} + res = {"mean": fcst["mean"]} + if level is None: + return res + level = sorted(level) + if self.prediction_intervals is not None: + res = self._add_predict_conformal_intervals(res, level) + else: + res = {**res, **_calculate_intervals(res, level, h, fcst["sigma"])} + return res def predict_in_sample(self, level: Optional[List[int]] = None): r"""Access fitted UCM insample predictions. Args: - level (List[float]): Confidence levels (0-100) for prediction intervals. Not yet supported. + level (List[float]): Confidence levels (0-100) for prediction intervals. Returns: - dict: Dictionary with entry `fitted` for point predictions. + dict: Dictionary with entries `fitted` for point predictions and `level_*` for probabilistic predictions. """ - self._no_pi(level) - return {"fitted": self.model_["fitted"]} + res = {"fitted": self.model_["fitted"]} + if level is not None: + level = sorted(level) + res = _add_fitted_pi(res=res, se=self.model_["sigma"], level=level) + return res def forecast( self, @@ -6391,19 +6397,27 @@ def forecast( h (int): Forecast horizon. X (array-like): Optional insample exogenous of shape (t, n_x). Currently ignored. X_future (array-like): Optional exogenous of shape (h, n_x). Currently ignored. - level (List[float]): Confidence levels (0-100) for prediction intervals. Not yet supported. + level (List[float]): Confidence levels (0-100) for prediction intervals. fitted (bool): Whether or not to return insample predictions. Returns: - dict: Dictionary with entry `mean` for point predictions and optionally `fitted` insample predictions. + dict: Dictionary with entries `mean` for point predictions and `level_*` for probabilistic predictions. """ - self._no_pi(level) y = _ensure_float(y) mod = ucm_model(y, season_length=self.season_length) fcst = ucm_forecast(mod, h) res = {"mean": fcst["mean"]} if fitted: res["fitted"] = fcst["fitted"] + if level is not None: + level = sorted(level) + if self.prediction_intervals is not None: + res = self._add_conformal_intervals(fcst=res, y=y, X=X, level=level) + else: + res = {**res, **_calculate_intervals(res, level, h, fcst["sigma"])} + if fitted: + se = _calculate_sigma(y - fcst["fitted"], len(y)) + res = _add_fitted_pi(res=res, se=se, level=level) return res diff --git a/python/statsforecast/ucm.py b/python/statsforecast/ucm.py index 7d0cf9c66..0061d0769 100644 --- a/python/statsforecast/ucm.py +++ b/python/statsforecast/ucm.py @@ -56,7 +56,7 @@ def kalman_filter( H: float, a0: np.ndarray, P0: np.ndarray, -) -> Tuple[float, np.ndarray, np.ndarray]: +) -> Tuple[float, np.ndarray, np.ndarray, np.ndarray]: """Run a univariate Kalman filter.""" n = y.shape[0] z = Z[0] @@ -91,7 +91,7 @@ def kalman_filter( a_pred = T @ a_filt P_pred = T @ P_filt @ T.T + RQRt - return loglik, a_filt, one_step_pred + return loglik, a_filt, P_filt, one_step_pred def _make_QH(theta: np.ndarray, has_seasonal: bool) -> Tuple[np.ndarray, float]: @@ -116,7 +116,7 @@ def _nll( ) -> float: """Negative log-likelihood for a variance vector theta.""" Q, H = _make_QH(theta, has_seasonal) - loglik, _, _ = kalman_filter(y, Z, T, R, Q, H, a0, P0) + loglik, _, _, _ = kalman_filter(y, Z, T, R, Q, H, a0, P0) if not np.isfinite(loglik): return 1e10 return -loglik @@ -152,7 +152,7 @@ def ucm_model(y: np.ndarray, season_length: int = 1) -> Dict: params = result.x Q, H = _make_QH(params, has_seasonal) - _, a_n, fitted = kalman_filter(y, Z, T, R, Q, H, a0, P0) + _, a_n, P_n, fitted = kalman_filter(y, Z, T, R, Q, H, a0, P0) residuals = y - fitted sigma = np.sqrt(np.sum(residuals**2) / max(len(y) - 1, 1)) @@ -165,6 +165,7 @@ def ucm_model(y: np.ndarray, season_length: int = 1) -> Dict: "Q": Q, "H": H, "a_n": a_n, + "P_n": P_n, "fitted": fitted, "sigma": sigma, "season_length": s, @@ -173,15 +174,22 @@ def ucm_model(y: np.ndarray, season_length: int = 1) -> Dict: def ucm_forecast(mod: Dict, h: int) -> Dict: - """Produce h-step-ahead point forecasts from a fitted UCM.""" + """Produce h-step-ahead point forecasts and forecast-error std from a UCM.""" Z = mod["Z"] T = mod["T"] + H = mod["H"] z = Z[0] a = mod["a_n"].astype(np.float64).copy() + P = mod["P_n"].astype(np.float64).copy() + RQRt = mod["R"] @ mod["Q"] @ mod["R"].T mean = np.empty(h, dtype=np.float64) + sigma = np.empty(h, dtype=np.float64) for i in range(h): a = T @ a + P = T @ P @ T.T + RQRt mean[i] = z @ a + # Forecast error variance F = Z P Z' + H. + sigma[i] = np.sqrt(z @ (P @ z) + H) - return {"mean": mean, "fitted": mod["fitted"]} + return {"mean": mean, "sigma": sigma, "fitted": mod["fitted"]} diff --git a/tests/test_ucm.py b/tests/test_ucm.py index 582600ef0..4c4bf7bab 100644 --- a/tests/test_ucm.py +++ b/tests/test_ucm.py @@ -1,7 +1,6 @@ """Tests for the UCM (Unobserved Components Model).""" import warnings - import numpy as np import pytest @@ -14,7 +13,6 @@ warnings.simplefilter("ignore") - @pytest.fixture def trend_series(): """Local-linear-trend series (no seasonality).""" @@ -22,7 +20,6 @@ def trend_series(): n = 120 return 50.0 + 0.5 * np.arange(n) + np.cumsum(np.random.randn(n)) - @pytest.fixture def seasonal_series(): """Trend + seasonal (period 12) series.""" @@ -36,7 +33,6 @@ def seasonal_series(): + np.random.randn(n) ) - # --------------------------------------------------------------------------- # Computational layer # --------------------------------------------------------------------------- @@ -73,10 +69,11 @@ def test_kalman_filter_loglik_finite(trend_series): a0 = np.zeros(k) P0 = 1e6 * np.eye(k) - loglik, a_filt, one_step = kalman_filter(y, Z, T, R, Q, H, a0, P0) + loglik, a_filt, P_filt, one_step = kalman_filter(y, Z, T, R, Q, H, a0, P0) assert np.isfinite(loglik) assert a_filt.shape == (k,) + assert P_filt.shape == (k, k) assert one_step.shape == (len(y),) assert np.all(np.isfinite(one_step)) @@ -108,6 +105,10 @@ def test_ucm_forecast_shape(trend_series): assert "fitted" in fcst assert len(fcst["mean"]) == h assert np.all(np.isfinite(fcst["mean"])) + # Forecast-error std is positive and widens with the horizon. + assert len(fcst["sigma"]) == h + assert np.all(fcst["sigma"] > 0) + assert fcst["sigma"][-1] >= fcst["sigma"][0] def test_ucm_forecast_trend_direction(trend_series): @@ -116,7 +117,6 @@ def test_ucm_forecast_trend_direction(trend_series): fcst = ucm_forecast(mod, 12) assert fcst["mean"][-1] > fcst["mean"][0] - # --------------------------------------------------------------------------- # Model class API # --------------------------------------------------------------------------- @@ -125,7 +125,6 @@ def test_ucm_import(): assert UCM is not None - def test_ucm_fit_predict(trend_series): from statsforecast.models import UCM @@ -137,7 +136,6 @@ def test_ucm_fit_predict(trend_series): assert len(fcst["mean"]) == 10 assert not np.any(np.isnan(fcst["mean"])) - def test_ucm_seasonal_fit_predict(seasonal_series): from statsforecast.models import UCM @@ -147,7 +145,6 @@ def test_ucm_seasonal_fit_predict(seasonal_series): assert len(fcst["mean"]) == 12 - def test_ucm_forecast_method(trend_series): from statsforecast.models import UCM @@ -159,7 +156,6 @@ def test_ucm_forecast_method(trend_series): assert len(res["mean"]) == 10 assert len(res["fitted"]) == len(trend_series) - def test_ucm_predict_in_sample(trend_series): from statsforecast.models import UCM @@ -170,25 +166,61 @@ def test_ucm_predict_in_sample(trend_series): assert "fitted" in insample assert len(insample["fitted"]) == len(trend_series) +# --------------------------------------------------------------------------- +# Prediction intervals +# --------------------------------------------------------------------------- +def test_ucm_parametric_predict_intervals(trend_series): + from statsforecast.models import UCM -def test_ucm_level_not_implemented(trend_series): + model = UCM(season_length=1) + model.fit(trend_series) + fcst = model.predict(h=10, level=[80, 95]) + + for key in ("lo-80", "hi-80", "lo-95", "hi-95"): + assert key in fcst + # Wider level contains the narrower one, and the mean lies inside. + assert np.all(fcst["lo-95"] <= fcst["lo-80"]) + assert np.all(fcst["hi-95"] >= fcst["hi-80"]) + assert np.all(fcst["lo-80"] <= fcst["mean"]) + assert np.all(fcst["mean"] <= fcst["hi-80"]) + +def test_ucm_parametric_forecast_intervals(trend_series): + from statsforecast.models import UCM + + model = UCM(season_length=1) + res = model.forecast(y=trend_series, h=10, level=[90], fitted=True) + + for key in ("lo-90", "hi-90", "fitted-lo-90", "fitted-hi-90"): + assert key in res + assert len(res["lo-90"]) == 10 + assert len(res["fitted-lo-90"]) == len(trend_series) + +def test_ucm_predict_in_sample_intervals(trend_series): from statsforecast.models import UCM model = UCM(season_length=1) model.fit(trend_series) - with pytest.raises(NotImplementedError): - model.predict(h=5, level=[95]) - with pytest.raises(NotImplementedError): - model.forecast(y=trend_series, h=5, level=[95]) + insample = model.predict_in_sample(level=[95]) + assert "fitted-lo-95" in insample + assert "fitted-hi-95" in insample + assert np.all(insample["fitted-lo-95"] <= insample["fitted-hi-95"]) -def test_ucm_alias(): +def test_ucm_conformal_intervals(seasonal_series): from statsforecast.models import UCM + from statsforecast.utils import ConformalIntervals - model = UCM(season_length=12, alias="MyUCM") - assert str(model) == "MyUCM" - assert model.alias == "MyUCM" + pi = ConformalIntervals(h=12, n_windows=2) + model = UCM(season_length=12, prediction_intervals=pi) + model.fit(seasonal_series) + pred = model.predict(h=12, level=[80, 95]) + for key in ("lo-80", "hi-80", "lo-95", "hi-95"): + assert key in pred + + fcst = model.forecast(y=seasonal_series, h=12, level=[80, 95]) + for key in ("lo-80", "hi-80", "lo-95", "hi-95"): + assert key in fcst def test_ucm_new(): from statsforecast.models import UCM @@ -199,7 +231,6 @@ def test_ucm_new(): assert model2.season_length == model1.season_length assert model2 is not model1 - # --------------------------------------------------------------------------- # StatsForecast integration # --------------------------------------------------------------------------- @@ -221,7 +252,6 @@ def test_statsforecast_integration(): assert len(fcst) == 6 assert "UCM" in fcst.columns - def test_statsforecast_cross_validation_fitted(): import pandas as pd from statsforecast import StatsForecast @@ -242,7 +272,6 @@ def test_statsforecast_cross_validation_fitted(): fitted_cv = sf.cross_validation_fitted_values() assert "UCM" in fitted_cv.columns - # --------------------------------------------------------------------------- # Numerical correctness vs statsmodels # --------------------------------------------------------------------------- @@ -254,17 +283,22 @@ def test_ucm_matches_statsmodels(seasonal_series): h = 12 sf_mod = ucm_model(y, season_length=12) - sf_fcst = ucm_forecast(sf_mod, h)["mean"] + sf_out = ucm_forecast(sf_mod, h) + sf_fcst = sf_out["mean"] sm_mod = sm.tsa.UnobservedComponents( y, level="local linear trend", seasonal=12, stochastic_seasonal=True ) sm_res = sm_mod.fit(method="lbfgs", disp=False) - sm_fcst = np.asarray(sm_res.forecast(h)) + sm_pred = sm_res.get_forecast(h) + sm_fcst = np.asarray(sm_pred.predicted_mean) + sm_se = np.asarray(sm_pred.se_mean) # Point forecasts should be close in scale and direction. np.testing.assert_allclose(sf_fcst, sm_fcst, rtol=0.15, atol=0.15 * np.std(y)) - + # The native (parametric) forecast std should match the Kalman-filter + # forecast std exposed by statsmodels. + np.testing.assert_allclose(sf_out["sigma"], sm_se, rtol=0.2, atol=0.2 * np.std(y)) if __name__ == "__main__": pytest.main([__file__, "-v"]) From ec2da3add2fdf0d2f6c73f06fe4da988c3ce34a5 Mon Sep 17 00:00:00 2001 From: MMenchero Date: Tue, 2 Jun 2026 14:20:13 -0600 Subject: [PATCH 3/8] implement fixes to Copilot's valid issues --- python/statsforecast/models.py | 3 ++- python/statsforecast/ucm.py | 2 ++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/python/statsforecast/models.py b/python/statsforecast/models.py index 7f9ec5960..2cf72b93c 100644 --- a/python/statsforecast/models.py +++ b/python/statsforecast/models.py @@ -6374,7 +6374,8 @@ def predict_in_sample(self, level: Optional[List[int]] = None): res = {"fitted": self.model_["fitted"]} if level is not None: level = sorted(level) - res = _add_fitted_pi(res=res, se=self.model_["sigma"], level=level) + se = self.model_["sigma"] + res = _add_fitted_pi(res=res, se=se, level=level) return res def forecast( diff --git a/python/statsforecast/ucm.py b/python/statsforecast/ucm.py index 0061d0769..aa64fe6c9 100644 --- a/python/statsforecast/ucm.py +++ b/python/statsforecast/ucm.py @@ -126,6 +126,8 @@ def ucm_model(y: np.ndarray, season_length: int = 1) -> Dict: """Fit the minimal UCM by maximum likelihood.""" y = np.asarray(y, dtype=np.float64) s = int(season_length) + if s < 1: + raise ValueError(f"season_length must be >= 1, got {season_length}.") has_seasonal = s > 1 Z, T, R = _build_matrices(s) From 0c04a9e89fb96b37fcac87134cd0ee53374b74a2 Mon Sep 17 00:00:00 2001 From: MMenchero Date: Tue, 2 Jun 2026 14:50:20 -0600 Subject: [PATCH 4/8] fix test --- tests/test_ucm.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/tests/test_ucm.py b/tests/test_ucm.py index 4c4bf7bab..54ea4efd7 100644 --- a/tests/test_ucm.py +++ b/tests/test_ucm.py @@ -276,29 +276,23 @@ def test_statsforecast_cross_validation_fitted(): # Numerical correctness vs statsmodels # --------------------------------------------------------------------------- def test_ucm_matches_statsmodels(seasonal_series): - """Forecasts should agree with statsmodels' UnobservedComponents.""" + """Forecasts should agree with statsmodels' UCM.""" sm = pytest.importorskip("statsmodels.api") y = seasonal_series h = 12 sf_mod = ucm_model(y, season_length=12) - sf_out = ucm_forecast(sf_mod, h) - sf_fcst = sf_out["mean"] + sf_fcst = ucm_forecast(sf_mod, h)["mean"] sm_mod = sm.tsa.UnobservedComponents( y, level="local linear trend", seasonal=12, stochastic_seasonal=True ) sm_res = sm_mod.fit(method="lbfgs", disp=False) - sm_pred = sm_res.get_forecast(h) - sm_fcst = np.asarray(sm_pred.predicted_mean) - sm_se = np.asarray(sm_pred.se_mean) + sm_fcst = np.asarray(sm_res.forecast(h)) # Point forecasts should be close in scale and direction. np.testing.assert_allclose(sf_fcst, sm_fcst, rtol=0.15, atol=0.15 * np.std(y)) - # The native (parametric) forecast std should match the Kalman-filter - # forecast std exposed by statsmodels. - np.testing.assert_allclose(sf_out["sigma"], sm_se, rtol=0.2, atol=0.2 * np.std(y)) if __name__ == "__main__": pytest.main([__file__, "-v"]) From 31f62e273136c767a65c8cf9e74b9c75d137712e Mon Sep 17 00:00:00 2001 From: MMenchero Date: Thu, 4 Jun 2026 16:29:42 -0600 Subject: [PATCH 5/8] add warning in case optimization doesn't converge --- python/statsforecast/ucm.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/statsforecast/ucm.py b/python/statsforecast/ucm.py index aa64fe6c9..33f472b19 100644 --- a/python/statsforecast/ucm.py +++ b/python/statsforecast/ucm.py @@ -1,5 +1,6 @@ __all__ = ['ucm_model', 'ucm_forecast'] +import warnings import numpy as np from scipy.optimize import minimize from typing import Dict, Tuple @@ -151,6 +152,9 @@ def ucm_model(y: np.ndarray, season_length: int = 1) -> Dict: method="L-BFGS-B", bounds=bounds, ) + if not result.success: + warnings.warn(f"UCM optimization did not converge: {result.message}", UserWarning) + params = result.x Q, H = _make_QH(params, has_seasonal) From d229a76907c95097a6ab56014631a961dd6a0861 Mon Sep 17 00:00:00 2001 From: MMenchero Date: Fri, 5 Jun 2026 21:34:57 -0600 Subject: [PATCH 6/8] add C++ implementation for Kalman Filter and UCM experiment with M3 and M4 data --- experiments/ucm/data.py | 41 +++++++++++ experiments/ucm/evaluation.py | 58 +++++++++++++++ experiments/ucm/statsforecast_ucm.py | 46 ++++++++++++ experiments/ucm/statsmodels_ucm.py | 72 ++++++++++++++++++ python/statsforecast/ucm.py | 13 +++- src/statsforecast.cpp | 5 ++ src/ucm.cpp | 106 +++++++++++++++++++++++++++ tests/test_ucm.py | 28 +++++++ 8 files changed, 365 insertions(+), 4 deletions(-) create mode 100644 experiments/ucm/data.py create mode 100644 experiments/ucm/evaluation.py create mode 100644 experiments/ucm/statsforecast_ucm.py create mode 100644 experiments/ucm/statsmodels_ucm.py create mode 100644 src/ucm.cpp diff --git a/experiments/ucm/data.py b/experiments/ucm/data.py new file mode 100644 index 000000000..18a7e70a5 --- /dev/null +++ b/experiments/ucm/data.py @@ -0,0 +1,41 @@ +import fire +from datasetsforecast.m4 import M4, M4Info +from datasetsforecast.m3 import M3, M3Info + +dict_datasets = { + 'M4': (M4, M4Info), + 'M3': (M3, M3Info), +} + +def get_data(directory: str, dataset: str, group: str, train: bool = True): + if dataset not in dict_datasets.keys(): + raise Exception(f'dataset {dataset} not found') + + dataclass, datainfo = dict_datasets[dataset] + if group not in datainfo.groups: + raise Exception(f'group {group} not found for {dataset}') + + Y_df, *_ = dataclass.load(directory, group) + + horizon = datainfo[group].horizon + freq = datainfo[group].freq + seasonality = datainfo[group].seasonality + Y_df_test = Y_df.groupby('unique_id').tail(horizon) + Y_df = Y_df.drop(Y_df_test.index) + + if train: + return Y_df, horizon, freq, seasonality + + return Y_df_test, horizon, freq, seasonality + + +def save_data(dataset: str, group: str, train: bool = True): + df, *_ = get_data('data', dataset, group, train) + if train: + df.to_csv(f'data/{dataset}-{group}-train.csv', index=False) + else: + df.to_csv(f'data/{dataset}-{group}-test.csv', index=False) + + +if __name__ == "__main__": + fire.Fire(save_data) diff --git a/experiments/ucm/evaluation.py b/experiments/ucm/evaluation.py new file mode 100644 index 000000000..4b04cd529 --- /dev/null +++ b/experiments/ucm/evaluation.py @@ -0,0 +1,58 @@ +import fire +import pandas as pd +from utilsforecast.evaluation import evaluate +from utilsforecast.losses import mae, rmse, smape +from data import get_data + +models = ['UCM-sf', 'UCM-sm'] + +def accuracy(dataset: str, group: str): + y_test, _, _, _ = get_data('data/', dataset, group, False) + # Positional id to merge forecasts with the test set (robust to differing + # `ds` between datasets/libraries). + y_test['id'] = y_test.groupby('unique_id').cumcount() + 1 + y_test = y_test.drop(columns=['ds']) + + sf = pd.read_csv(f'data/UCM-sf-{dataset}-{group}.csv') + sf['id'] = sf.groupby('unique_id').cumcount() + 1 + sm = pd.read_csv(f'data/UCM-sm-{dataset}-{group}.csv') + sm['id'] = sm.groupby('unique_id').cumcount() + 1 + + forecasts = sf.merge(sm[['unique_id', 'id', 'UCM-sm']], on=['unique_id', 'id']) + predictions = forecasts.merge(y_test, on=['unique_id', 'id']) + predictions = predictions[['unique_id', 'ds', 'y'] + models] + + evaluation = evaluate(predictions, metrics=[mae, rmse, smape], models=models) + evals = evaluation.groupby('metric').mean(numeric_only=True).reset_index() + evals = pd.melt( + evals, id_vars=['metric'], value_vars=models, var_name='model', value_name='value' + ) + evals = evals.pivot(index='model', columns='metric', values='value').reset_index() + evals['dataset'] = f'{dataset}_{group}' + + times = pd.concat([ + pd.read_csv(f'data/UCM-sf-time-{dataset}-{group}.csv'), + pd.read_csv(f'data/UCM-sm-time-{dataset}-{group}.csv'), + ]) + evals = evals.merge(times, on='model', how='left') + return evals + +def main(dataset: str = 'M3'): + if dataset == 'M3': + groups = ['Yearly', 'Quarterly', 'Monthly', 'Other'] + elif dataset == 'M4': + groups = ['Yearly', 'Quarterly', 'Monthly', 'Weekly', 'Daily', 'Hourly'] + else: + raise ValueError(f'Dataset {dataset} not found') + + evaluation = pd.concat([accuracy(dataset, group) for group in groups]) + evaluation = evaluation[['dataset', 'model', 'mae', 'rmse', 'smape', 'time']] + evaluation['smape'] = evaluation['smape'] * 100 + evaluation['time'] = evaluation['time'] / 60 + evaluation[['mae', 'rmse', 'smape', 'time']] = evaluation[['mae', 'rmse', 'smape', 'time']].round(2) + evaluation = evaluation.sort_values(by=['dataset', 'model']) + evaluation.to_csv(f'data/evaluation-{dataset}.csv', index=False) + print(evaluation.to_markdown(index=False)) + +if __name__ == '__main__': + fire.Fire(main) diff --git a/experiments/ucm/statsforecast_ucm.py b/experiments/ucm/statsforecast_ucm.py new file mode 100644 index 000000000..a218e4292 --- /dev/null +++ b/experiments/ucm/statsforecast_ucm.py @@ -0,0 +1,46 @@ +import time +import fire +import pandas as pd +from multiprocessing import cpu_count +from statsforecast import StatsForecast +from statsforecast.models import UCM +from data import get_data + +def find_seasonality(group: str) -> int: + seasonality = { + 'Yearly': 1, + 'Quarterly': 4, + 'Monthly': 12, + 'Weekly': 1, + 'Daily': 7, + 'Hourly': 24, + 'Other': 1, + } + if group not in seasonality: + raise ValueError(f'Group {group} not found') + return seasonality[group] + +def main(dataset: str = 'M3', group: str = 'Other') -> None: + train, horizon, freq, _ = get_data('data/', dataset, group) + + if dataset == 'M4': + train['ds'] = train['ds'].astype(int) + freq = 1 # since values in ds column are integers + + season_length = find_seasonality(group) + models = [UCM(season_length=season_length, alias='UCM-sf')] + + start = time.time() + sf = StatsForecast(models=models, freq=freq, n_jobs=cpu_count()) + forecasts = sf.forecast(df=train, h=horizon) + end = time.time() + + forecasts.to_csv(f'data/UCM-sf-{dataset}-{group}.csv', index=False) + + time_df = pd.DataFrame({'time': [end - start], 'model': ['UCM-sf']}) + time_df.to_csv(f'data/UCM-sf-time-{dataset}-{group}.csv', index=False) + + print(f'Dataset: {dataset} - Group: {group} with statsforecast completed in {end - start} seconds') + +if __name__ == '__main__': + fire.Fire(main) diff --git a/experiments/ucm/statsmodels_ucm.py b/experiments/ucm/statsmodels_ucm.py new file mode 100644 index 000000000..a176954de --- /dev/null +++ b/experiments/ucm/statsmodels_ucm.py @@ -0,0 +1,72 @@ +import os + +os.environ["OMP_NUM_THREADS"] = "1" +os.environ["OPENBLAS_NUM_THREADS"] = "1" +os.environ["MKL_NUM_THREADS"] = "1" +os.environ["NUMEXPR_NUM_THREADS"] = "1" +os.environ["VECLIB_MAXIMUM_THREADS"] = "1" + +import time +import fire +import warnings +import concurrent.futures +import numpy as np +import pandas as pd +from multiprocessing import cpu_count +from statsmodels.tsa.statespace.structural import UnobservedComponents +from data import get_data +from statsforecast_ucm import find_seasonality + +warnings.simplefilter("ignore") + +def generate_fcst(vals): + unique_id, y, horizon, season_length = vals + # Equivalent to the statsforecast UCM: local linear trend + seasonal + irregular. + kwargs = dict(level="local linear trend", irregular=True) + if season_length > 1: + kwargs.update(seasonal=season_length, stochastic_seasonal=True) + try: + res = UnobservedComponents(y, **kwargs).fit(disp=False) + fcst = np.asarray(res.forecast(horizon)) + except Exception: + # Fallback is a naive model. + fcst = np.repeat(y[-1], horizon) + return unique_id, fcst + +def main(dataset: str = 'M4', group: str = 'Monthly', n_workers: int = None) -> None: + train, horizon, _, _ = get_data('data/', dataset, group) + season_length = find_seasonality(group) + + vals_list = [ + (uid, grp['y'].values, horizon, season_length) + for uid, grp in train.groupby('unique_id', sort=False) + ] + n_series = len(vals_list) + + if n_workers is None: + n_workers = cpu_count() + n_workers = min(n_workers, n_series) + + chunksize = min(256, max(1, n_series // (n_workers * 8))) + + ids = [] + fcsts = [] + start = time.time() + with concurrent.futures.ProcessPoolExecutor(max_workers=n_workers) as executor: + for uid, fcst in executor.map(generate_fcst, vals_list, chunksize=chunksize): + ids.append(uid) + fcsts.append(fcst) + end = time.time() + + forecasts = pd.DataFrame({ + 'unique_id': np.repeat(ids, horizon), + 'UCM-sm': np.concatenate(fcsts), + }) + + forecasts.to_csv(f'data/UCM-sm-{dataset}-{group}.csv', index=False) + + time_df = pd.DataFrame({'time': [end - start], 'model': ['UCM-sm']}) + time_df.to_csv(f'data/UCM-sm-time-{dataset}-{group}.csv', index=False) + +if __name__ == '__main__': + fire.Fire(main) \ No newline at end of file diff --git a/python/statsforecast/ucm.py b/python/statsforecast/ucm.py index 33f472b19..7f9243389 100644 --- a/python/statsforecast/ucm.py +++ b/python/statsforecast/ucm.py @@ -5,6 +5,8 @@ from scipy.optimize import minimize from typing import Dict, Tuple +from ._lib import ucm as _ucm + # Diffuse initialization: initial state mean is 0 and the initial state # covariance is a large multiple of the identity. _DIFFUSE_VARIANCE = 1e6 @@ -58,7 +60,9 @@ def kalman_filter( a0: np.ndarray, P0: np.ndarray, ) -> Tuple[float, np.ndarray, np.ndarray, np.ndarray]: - """Run a univariate Kalman filter.""" + """Run a univariate Kalman filter. + Pure-Python reference implementation, retained for validation. + """ n = y.shape[0] z = Z[0] RQRt = R @ Q @ R.T @@ -117,7 +121,7 @@ def _nll( ) -> float: """Negative log-likelihood for a variance vector theta.""" Q, H = _make_QH(theta, has_seasonal) - loglik, _, _, _ = kalman_filter(y, Z, T, R, Q, H, a0, P0) + loglik = _ucm.loglik(y, Z, T, R, Q, float(H), a0, P0) if not np.isfinite(loglik): return 1e10 return -loglik @@ -149,16 +153,17 @@ def ucm_model(y: np.ndarray, season_length: int = 1) -> Dict: _nll, x0, args=(y, Z, T, R, a0, P0, has_seasonal), - method="L-BFGS-B", + method="Nelder-Mead", bounds=bounds, ) + if not result.success: warnings.warn(f"UCM optimization did not converge: {result.message}", UserWarning) params = result.x Q, H = _make_QH(params, has_seasonal) - _, a_n, P_n, fitted = kalman_filter(y, Z, T, R, Q, H, a0, P0) + _, a_n, P_n, fitted = _ucm.filter(y, Z, T, R, Q, float(H), a0, P0) residuals = y - fitted sigma = np.sqrt(np.sum(residuals**2) / max(len(y) - 1, 1)) diff --git a/src/statsforecast.cpp b/src/statsforecast.cpp index 5a8b3ae2e..d51f36fd2 100644 --- a/src/statsforecast.cpp +++ b/src/statsforecast.cpp @@ -34,6 +34,10 @@ namespace mfles { void init(py::module_ &); } +namespace ucm { +void init(py::module_ &); +} + PYBIND11_MODULE(_lib, m) { arima::init(m); ets::init(m); @@ -43,4 +47,5 @@ PYBIND11_MODULE(_lib, m) { ces::init(m); tbats_ns::init(m); mfles::init(m); + ucm::init(m); } diff --git a/src/ucm.cpp b/src/ucm.cpp new file mode 100644 index 000000000..41241cee3 --- /dev/null +++ b/src/ucm.cpp @@ -0,0 +1,106 @@ +#define _USE_MATH_DEFINES +#include +#include +#include + +#include +#include +#include + +namespace ucm { +namespace py = pybind11; +using Eigen::MatrixXd; +using Eigen::VectorXd; + +double kalman_filter_core(const Eigen::Ref &y, + const Eigen::Ref &Z, + const Eigen::Ref &T, + const Eigen::Ref &R, + const Eigen::Ref &Q, double H, + const Eigen::Ref &a0, + const Eigen::Ref &P0, bool collect, + VectorXd &a_filt, MatrixXd &P_filt, + VectorXd &one_step_pred) { + const Eigen::Index n = y.size(); + const VectorXd z = Z.row(0).transpose(); + const MatrixXd RQRt = R * Q * R.transpose(); + + VectorXd a_pred = a0; + MatrixXd P_pred = P0; + + a_filt = a_pred; + P_filt = P_pred; + + double loglik = 0.0; + const double log_2pi = std::log(2.0 * M_PI); + + for (Eigen::Index t = 0; t < n; ++t) { + // Forecast for the current observation. + double y_hat = z.dot(a_pred); + if (collect) { + one_step_pred[t] = y_hat; + } + double v = y[t] - y_hat; + VectorXd Pz = P_pred * z; + double F = z.dot(Pz) + H; + + // Update step (skip if F is not usable). + if (std::isfinite(F) && F > 0.0) { + loglik += -0.5 * (log_2pi + std::log(F) + v * v / F); + VectorXd K = Pz / F; + a_filt = a_pred + K * v; + P_filt = P_pred - K * Pz.transpose(); + } else { + a_filt = a_pred; + P_filt = P_pred; + } + + // Predict the next state. + a_pred = T * a_filt; + P_pred = T * P_filt * T.transpose() + RQRt; + } + + return loglik; +} + +// Scalar log-likelihood for the optimizer hot path. +double loglik(const Eigen::Ref &y, + const Eigen::Ref &Z, + const Eigen::Ref &T, + const Eigen::Ref &R, + const Eigen::Ref &Q, double H, + const Eigen::Ref &a0, + const Eigen::Ref &P0) { + VectorXd a_filt; + MatrixXd P_filt; + VectorXd dummy; + return kalman_filter_core(y, Z, T, R, Q, H, a0, P0, false, a_filt, P_filt, + dummy); +} + +// Full filter returning (loglik, a_filt, P_filt, one_step_pred) for the +// final fit. +std::tuple +filter(const Eigen::Ref &y, + const Eigen::Ref &Z, + const Eigen::Ref &T, + const Eigen::Ref &R, + const Eigen::Ref &Q, double H, + const Eigen::Ref &a0, + const Eigen::Ref &P0) { + VectorXd a_filt; + MatrixXd P_filt; + VectorXd one_step_pred(y.size()); + double ll = + kalman_filter_core(y, Z, T, R, Q, H, a0, P0, true, a_filt, P_filt, + one_step_pred); + return {ll, a_filt, P_filt, one_step_pred}; +} + +void init(py::module_ &m) { + py::module_ ucm_mod = m.def_submodule("ucm"); + ucm_mod.def("loglik", &loglik, py::call_guard()); + ucm_mod.def("filter", &filter, py::call_guard()); +} + +} // namespace ucm diff --git a/tests/test_ucm.py b/tests/test_ucm.py index 54ea4efd7..f1097dbe0 100644 --- a/tests/test_ucm.py +++ b/tests/test_ucm.py @@ -78,6 +78,34 @@ def test_kalman_filter_loglik_finite(trend_series): assert np.all(np.isfinite(one_step)) +@pytest.mark.parametrize("season_length", [1, 12]) +def test_kalman_filter_cpp_matches_python( + season_length, trend_series, seasonal_series +): + """The compiled Kalman filter must match the Python reference.""" + from statsforecast._lib import ucm as _ucm + + y = trend_series if season_length == 1 else seasonal_series + Z, T, R = _build_matrices(season_length=season_length) + k = T.shape[0] + n_shocks = R.shape[1] + Q = np.diag(np.linspace(1.0, 0.1, n_shocks)) + H = 1.5 + a0 = np.zeros(k) + P0 = 1e6 * np.eye(k) + + py_ll, py_a, py_P, py_pred = kalman_filter(y, Z, T, R, Q, H, a0, P0) + cpp_ll, cpp_a, cpp_P, cpp_pred = _ucm.filter(y, Z, T, R, Q, H, a0, P0) + + np.testing.assert_allclose(cpp_ll, py_ll, rtol=1e-7, atol=1e-7) + np.testing.assert_allclose(cpp_a, py_a, rtol=1e-7, atol=1e-7) + np.testing.assert_allclose(cpp_P, py_P, rtol=1e-6, atol=1e-6) + np.testing.assert_allclose(cpp_pred, py_pred, rtol=1e-7, atol=1e-7) + np.testing.assert_allclose( + _ucm.loglik(y, Z, T, R, Q, H, a0, P0), py_ll, rtol=1e-7, atol=1e-7 + ) + + def test_ucm_model_params_trend_only(trend_series): mod = ucm_model(trend_series, season_length=1) # [sigma2_eps, sigma2_eta, sigma2_zeta] From e703382bb7ed2c18c8fe5795c31cdcd039aa9947 Mon Sep 17 00:00:00 2001 From: MMenchero Date: Mon, 8 Jun 2026 20:32:29 -0600 Subject: [PATCH 7/8] improve statsmodels eval --- experiments/ucm/evaluation.py | 2 +- experiments/ucm/statsmodels_ucm.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/experiments/ucm/evaluation.py b/experiments/ucm/evaluation.py index 4b04cd529..02091e4c0 100644 --- a/experiments/ucm/evaluation.py +++ b/experiments/ucm/evaluation.py @@ -49,7 +49,7 @@ def main(dataset: str = 'M3'): evaluation = evaluation[['dataset', 'model', 'mae', 'rmse', 'smape', 'time']] evaluation['smape'] = evaluation['smape'] * 100 evaluation['time'] = evaluation['time'] / 60 - evaluation[['mae', 'rmse', 'smape', 'time']] = evaluation[['mae', 'rmse', 'smape', 'time']].round(2) + evaluation[['mae', 'rmse', 'smape', 'time']] = evaluation[['mae', 'rmse', 'smape', 'time']].round(3) evaluation = evaluation.sort_values(by=['dataset', 'model']) evaluation.to_csv(f'data/evaluation-{dataset}.csv', index=False) print(evaluation.to_markdown(index=False)) diff --git a/experiments/ucm/statsmodels_ucm.py b/experiments/ucm/statsmodels_ucm.py index a176954de..e1896bf83 100644 --- a/experiments/ucm/statsmodels_ucm.py +++ b/experiments/ucm/statsmodels_ucm.py @@ -68,5 +68,7 @@ def main(dataset: str = 'M4', group: str = 'Monthly', n_workers: int = None) -> time_df = pd.DataFrame({'time': [end - start], 'model': ['UCM-sm']}) time_df.to_csv(f'data/UCM-sm-time-{dataset}-{group}.csv', index=False) + print(f'Dataset: {dataset} - Group: {group} with statsmodels completed in {end - start} seconds') + if __name__ == '__main__': fire.Fire(main) \ No newline at end of file From e1c12121b4cf3ee6025e8f05552ee96d809e5c00 Mon Sep 17 00:00:00 2001 From: MMenchero Date: Sat, 20 Jun 2026 17:09:17 -0600 Subject: [PATCH 8/8] add UCM experiment with M3 and M4 data for prediction intervals --- .gitignore | 8 +- experiments/ucm_intervals/data.py | 41 ++++++++++ experiments/ucm_intervals/evaluation.py | 81 +++++++++++++++++++ .../ucm_intervals/statsforecast_ucm.py | 46 +++++++++++ experiments/ucm_intervals/statsmodels_ucm.py | 80 ++++++++++++++++++ 5 files changed, 255 insertions(+), 1 deletion(-) create mode 100644 experiments/ucm_intervals/data.py create mode 100644 experiments/ucm_intervals/evaluation.py create mode 100644 experiments/ucm_intervals/statsforecast_ucm.py create mode 100644 experiments/ucm_intervals/statsmodels_ucm.py diff --git a/.gitignore b/.gitignore index 793b94cde..601097377 100644 --- a/.gitignore +++ b/.gitignore @@ -39,4 +39,10 @@ nbs/_extensions .quarto *.png -.coverage* \ No newline at end of file +.coverage* + + +# local +.cursor/ +instructions_chronax_comparison.md +my_nbs/ \ No newline at end of file diff --git a/experiments/ucm_intervals/data.py b/experiments/ucm_intervals/data.py new file mode 100644 index 000000000..18a7e70a5 --- /dev/null +++ b/experiments/ucm_intervals/data.py @@ -0,0 +1,41 @@ +import fire +from datasetsforecast.m4 import M4, M4Info +from datasetsforecast.m3 import M3, M3Info + +dict_datasets = { + 'M4': (M4, M4Info), + 'M3': (M3, M3Info), +} + +def get_data(directory: str, dataset: str, group: str, train: bool = True): + if dataset not in dict_datasets.keys(): + raise Exception(f'dataset {dataset} not found') + + dataclass, datainfo = dict_datasets[dataset] + if group not in datainfo.groups: + raise Exception(f'group {group} not found for {dataset}') + + Y_df, *_ = dataclass.load(directory, group) + + horizon = datainfo[group].horizon + freq = datainfo[group].freq + seasonality = datainfo[group].seasonality + Y_df_test = Y_df.groupby('unique_id').tail(horizon) + Y_df = Y_df.drop(Y_df_test.index) + + if train: + return Y_df, horizon, freq, seasonality + + return Y_df_test, horizon, freq, seasonality + + +def save_data(dataset: str, group: str, train: bool = True): + df, *_ = get_data('data', dataset, group, train) + if train: + df.to_csv(f'data/{dataset}-{group}-train.csv', index=False) + else: + df.to_csv(f'data/{dataset}-{group}-test.csv', index=False) + + +if __name__ == "__main__": + fire.Fire(save_data) diff --git a/experiments/ucm_intervals/evaluation.py b/experiments/ucm_intervals/evaluation.py new file mode 100644 index 000000000..57fcadf4e --- /dev/null +++ b/experiments/ucm_intervals/evaluation.py @@ -0,0 +1,81 @@ +import fire +import pandas as pd +from typing import List +from utilsforecast.evaluation import evaluate +from utilsforecast.losses import scaled_crps, coverage +from data import get_data + +models = ['UCM-sf', 'UCM-sm'] +level = [80, 95] +interval_cols = [f'{m}-{side}-{lvl}' for m in models for side in ('lo', 'hi') for lvl in level] + +def winkler( + df, + models: List[str], + level: int, + id_col: str = 'unique_id', + target_col: str = 'y', + cutoff_col: str = 'cutoff', +): + alpha = 1 - level / 100 + y = df[target_col].to_numpy() + out = df[[id_col]].copy() + for m in models: + lo = df[f'{m}-lo-{level}'].to_numpy() + hi = df[f'{m}-hi-{level}'].to_numpy() + out[m] = (hi - lo) + (2 / alpha) * (lo - y) * (y < lo) + (2 / alpha) * (y - hi) * (y > hi) + return out.groupby(id_col, observed=True).mean().reset_index() + +def accuracy(dataset: str, group: str): + y_test, _, _, _ = get_data('data/', dataset, group, False) + # Positional id to merge forecasts with the test set (robust to differing + # `ds` between datasets/libraries). + y_test['id'] = y_test.groupby('unique_id').cumcount() + 1 + y_test = y_test.drop(columns=['ds']) + + sf = pd.read_csv(f'data/UCM-sf-{dataset}-{group}.csv') + sf['id'] = sf.groupby('unique_id').cumcount() + 1 + sm = pd.read_csv(f'data/UCM-sm-{dataset}-{group}.csv') + sm['id'] = sm.groupby('unique_id').cumcount() + 1 + + sm_cols = ['unique_id', 'id', 'UCM-sm'] + [c for c in interval_cols if c.startswith('UCM-sm')] + forecasts = sf.merge(sm[sm_cols], on=['unique_id', 'id']) + predictions = forecasts.merge(y_test, on=['unique_id', 'id']) + predictions = predictions[['unique_id', 'ds', 'y'] + models + interval_cols] + + evaluation = evaluate( + predictions, metrics=[scaled_crps, winkler, coverage], models=models, level=level + ) + evals = evaluation.groupby('metric').mean(numeric_only=True).reset_index() + evals = pd.melt( + evals, id_vars=['metric'], value_vars=models, var_name='model', value_name='value' + ) + evals = evals.pivot(index='model', columns='metric', values='value').reset_index() + evals['dataset'] = f'{dataset}_{group}' + + times = pd.concat([ + pd.read_csv(f'data/UCM-sf-time-{dataset}-{group}.csv'), + pd.read_csv(f'data/UCM-sm-time-{dataset}-{group}.csv'), + ]) + evals = evals.merge(times, on='model', how='left') + return evals + +def main(dataset: str = 'M3'): + if dataset == 'M3': + groups = ['Yearly', 'Quarterly', 'Monthly', 'Other'] + elif dataset == 'M4': + groups = ['Yearly', 'Quarterly', 'Monthly', 'Weekly', 'Daily', 'Hourly'] + else: + raise ValueError(f'Dataset {dataset} not found') + + metrics = ['scaled_crps', 'winkler_level80', 'winkler_level95', 'coverage_level80', 'coverage_level95'] + evaluation = pd.concat([accuracy(dataset, group) for group in groups]) + evaluation = evaluation[['dataset', 'model'] + metrics + ['time']] + evaluation['time'] = evaluation['time'] / 60 + evaluation[metrics + ['time']] = evaluation[metrics + ['time']].round(4) + evaluation = evaluation.sort_values(by=['dataset', 'model']) + evaluation.to_csv(f'data/evaluation-{dataset}.csv', index=False) + print(evaluation.to_markdown(index=False)) + +if __name__ == '__main__': + fire.Fire(main) diff --git a/experiments/ucm_intervals/statsforecast_ucm.py b/experiments/ucm_intervals/statsforecast_ucm.py new file mode 100644 index 000000000..334b27e4a --- /dev/null +++ b/experiments/ucm_intervals/statsforecast_ucm.py @@ -0,0 +1,46 @@ +import time +import fire +import pandas as pd +from multiprocessing import cpu_count +from statsforecast import StatsForecast +from statsforecast.models import UCM +from data import get_data + +def find_seasonality(group: str) -> int: + seasonality = { + 'Yearly': 1, + 'Quarterly': 4, + 'Monthly': 12, + 'Weekly': 1, + 'Daily': 7, + 'Hourly': 24, + 'Other': 1, + } + if group not in seasonality: + raise ValueError(f'Group {group} not found') + return seasonality[group] + +def main(dataset: str = 'M3', group: str = 'Other') -> None: + train, horizon, freq, _ = get_data('data/', dataset, group) + + if dataset == 'M4': + train['ds'] = train['ds'].astype(int) + freq = 1 # since values in ds column are integers + + season_length = find_seasonality(group) + models = [UCM(season_length=season_length, alias='UCM-sf')] + + start = time.time() + sf = StatsForecast(models=models, freq=freq, n_jobs=cpu_count()) + forecasts = sf.forecast(df=train, h=horizon, level=[80, 95]) + end = time.time() + + forecasts.to_csv(f'data/UCM-sf-{dataset}-{group}.csv', index=False) + + time_df = pd.DataFrame({'time': [end - start], 'model': ['UCM-sf']}) + time_df.to_csv(f'data/UCM-sf-time-{dataset}-{group}.csv', index=False) + + print(f'Dataset: {dataset} - Group: {group} with statsforecast completed in {end - start} seconds') + +if __name__ == '__main__': + fire.Fire(main) diff --git a/experiments/ucm_intervals/statsmodels_ucm.py b/experiments/ucm_intervals/statsmodels_ucm.py new file mode 100644 index 000000000..ed97e084c --- /dev/null +++ b/experiments/ucm_intervals/statsmodels_ucm.py @@ -0,0 +1,80 @@ +import os + +os.environ["OMP_NUM_THREADS"] = "1" +os.environ["OPENBLAS_NUM_THREADS"] = "1" +os.environ["MKL_NUM_THREADS"] = "1" +os.environ["NUMEXPR_NUM_THREADS"] = "1" +os.environ["VECLIB_MAXIMUM_THREADS"] = "1" + +import time +import fire +import warnings +import concurrent.futures +import numpy as np +import pandas as pd +from multiprocessing import cpu_count +from statsmodels.tsa.statespace.structural import UnobservedComponents +from data import get_data +from statsforecast_ucm import find_seasonality + +warnings.simplefilter("ignore") + +def generate_fcst(vals): + unique_id, y, horizon, season_length = vals + # Equivalent to the statsforecast UCM: local linear trend + seasonal + irregular. + kwargs = dict(level="local linear trend", irregular=True) + if season_length > 1: + kwargs.update(seasonal=season_length, stochastic_seasonal=True) + res = UnobservedComponents(y, **kwargs).fit(disp=False) + pred = res.get_forecast(horizon) + mean = np.asarray(pred.predicted_mean) + lo80, hi80 = np.asarray(pred.conf_int(alpha=0.20)).T + lo95, hi95 = np.asarray(pred.conf_int(alpha=0.05)).T + return unique_id, mean, lo80, hi80, lo95, hi95 + +def main(dataset: str = 'M4', group: str = 'Monthly', n_workers: int = None) -> None: + train, horizon, _, _ = get_data('data/', dataset, group) + season_length = find_seasonality(group) + + vals_list = [ + (uid, grp['y'].values, horizon, season_length) + for uid, grp in train.groupby('unique_id', sort=False) + ] + n_series = len(vals_list) + + if n_workers is None: + n_workers = cpu_count() + n_workers = min(n_workers, n_series) + + chunksize = min(256, max(1, n_series // (n_workers * 8))) + + ids, means, lo80s, hi80s, lo95s, hi95s = [], [], [], [], [], [] + start = time.time() + with concurrent.futures.ProcessPoolExecutor(max_workers=n_workers) as executor: + for uid, mean, lo80, hi80, lo95, hi95 in executor.map(generate_fcst, vals_list, chunksize=chunksize): + ids.append(uid) + means.append(mean) + lo80s.append(lo80) + hi80s.append(hi80) + lo95s.append(lo95) + hi95s.append(hi95) + end = time.time() + + forecasts = pd.DataFrame({ + 'unique_id': np.repeat(ids, horizon), + 'UCM-sm': np.concatenate(means), + 'UCM-sm-lo-80': np.concatenate(lo80s), + 'UCM-sm-hi-80': np.concatenate(hi80s), + 'UCM-sm-lo-95': np.concatenate(lo95s), + 'UCM-sm-hi-95': np.concatenate(hi95s), + }) + + forecasts.to_csv(f'data/UCM-sm-{dataset}-{group}.csv', index=False) + + time_df = pd.DataFrame({'time': [end - start], 'model': ['UCM-sm']}) + time_df.to_csv(f'data/UCM-sm-time-{dataset}-{group}.csv', index=False) + + print(f'Dataset: {dataset} - Group: {group} with statsmodels completed in {end - start} seconds') + +if __name__ == '__main__': + fire.Fire(main)