src/models/sarima.py

import pandas as pd
import numpy as np
from statsmodels.tsa.statespace.sarimax import SARIMAX

class SARIMAModel:
    """
    Seasonal Autoregressive Integrated Moving Average (SARIMA) model.
    Encapsulates statsmodels SARIMAX for modeling time series structure.
    """
    def __init__(self, order=(1, 0, 0), seasonal_order=(1, 0, 0, 24)):
        self.order = order
        self.seasonal_order = seasonal_order
        self.model = None
        self.results = None

    def fit(self, y: pd.Series):
        """Fits the SARIMA model on time series target y."""
        # Suppress potential warnings from statespace initialization
        self.model = SARIMAX(
            y,
            order=self.order,
            seasonal_order=self.seasonal_order,
            enforce_stationarity=False,
            enforce_invertibility=False
        )
        self.results = self.model.fit(disp=False)
        return self

    def predict(self, y_new: pd.Series) -> pd.Series:
        """
        Generates one-step-ahead predictions for the new series y_new
        by extending the fitted statespace model.
        """
        if self.results is None:
            raise ValueError("Model is not fitted yet. Call fit() first.")
        extended = self.results.extend(y_new)
        return extended.predict(start=y_new.index[0], end=y_new.index[-1])

    def summary(self) -> str:
        """Returns the summary of the fitted model."""
        if self.results is None:
            return "Model not fitted yet."
        return self.results.summary().as_text()