src/models/baseline.py

import statsmodels.api as sm
import pandas as pd
import numpy as np

class OLSBaseline:
    """
    Ordinary Least Squares (OLS) Multiple Regression baseline model.
    Encapsulates statsmodels OLS with automatic intercept handling.
    """
    def __init__(self):
        self.model = None
        self.results = None
        self.params = None

    def fit(self, X: pd.DataFrame, y: pd.Series):
        """Fits the OLS model on the training features X and target y."""
        X_sm = sm.add_constant(X)
        self.model = sm.OLS(y, X_sm)
        self.results = self.model.fit()
        self.params = self.results.params
        return self

    def predict(self, X: pd.DataFrame) -> np.ndarray:
        """Generates predictions for the given features X."""
        if self.results is None:
            raise ValueError("Model is not fitted yet. Call fit() first.")
        X_sm = sm.add_constant(X, has_constant='add')
        return self.results.predict(X_sm).values

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