src/ml/compare_models.py

"""Reusable utility to compare regression models side by side.

Imported by both ``train.py`` and the EDA notebook to avoid duplication.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable

import numpy as np
import pandas as pd
from sklearn.base import RegressorMixin
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import KFold, cross_val_score
from sklearn.pipeline import Pipeline


@dataclass
class ModelResult:
    name: str
    mae: float
    rmse: float
    r2: float
    cv_mae_mean: float | None = None
    cv_mae_std: float | None = None


def evaluate_pipeline(
    name: str,
    pipe: Pipeline,
    X_train: pd.DataFrame,
    y_train: pd.Series,
    X_test: pd.DataFrame,
    y_test: pd.Series,
    cv_folds: int | None = 5,
) -> ModelResult:
    pipe.fit(X_train, y_train)
    pred = pipe.predict(X_test)
    mae = float(mean_absolute_error(y_test, pred))
    rmse = float(np.sqrt(mean_squared_error(y_test, pred)))
    r2 = float(r2_score(y_test, pred))

    cv_mean = cv_std = None
    if cv_folds and cv_folds > 1:
        kf = KFold(n_splits=cv_folds, shuffle=True, random_state=42)
        scores = -cross_val_score(
            pipe, X_train, y_train, scoring="neg_mean_absolute_error", cv=kf, n_jobs=-1
        )
        cv_mean = float(scores.mean())
        cv_std = float(scores.std())

    return ModelResult(name, mae, rmse, r2, cv_mean, cv_std)


def results_to_dataframe(results: Iterable[ModelResult]) -> pd.DataFrame:
    return pd.DataFrame(
        [
            {
                "model": r.name,
                "MAE": r.mae,
                "RMSE": r.rmse,
                "R2": r.r2,
                "CV_MAE_mean": r.cv_mae_mean,
                "CV_MAE_std": r.cv_mae_std,
            }
            for r in results
        ]
    ).sort_values("MAE")