model_evaluation.py

# ---------------------------------------------------------------------------
#  Model evaluation and visualization for NYC taxi regression
# ---------------------------------------------------------------------------
# Provides helpers to extract LightGBM from pipeline output, plot feature
# importance, and compare baseline vs tuned validation metrics. Does not
# train or tune models; works on fitted models and metric dicts from
# regression_models.evaluate_regression_metrics.
# ---------------------------------------------------------------------------

import numpy as np
import matplotlib.pyplot as plt


def get_lightgbm_from_fitted_list(fitted_list):
    """
    Return the LightGBM model from a list of (model_name, fitted_model) tuples.

    Used when the pipeline returns one entry per model type (e.g. from
    train_and_evaluate_model_loop). Identifies LightGBM by name.

    Args:
        fitted_list: List of (name, model) tuples (e.g. models_dur["base+all"]).

    Returns:
        Fitted LightGBM model, or None if not found.
    """
    if not fitted_list:
        return None
    for name, model in fitted_list:
        if "LightGBM" in name or "LGBM" in name:
            return model
    return None


def plot_feature_importance(model, feature_names, title, top_n=20, figsize=(10, 6)):
    """
    Plot horizontal bar chart of feature importances for a tree-based model.

    Uses model.feature_importances_ (e.g. LightGBM, scikit-learn tree models).
    Sorts by importance and shows the top_n features.

    Args:
        model: Fitted model with .feature_importances_ (1d array).
        feature_names: Sequence of feature names (same order as model features).
        title: Plot title string.
        top_n: Number of top features to show (default 20).
        figsize: Figure size (default (10, 6)).
    """
    imp = np.asarray(model.feature_importances_).squeeze()
    if hasattr(feature_names, "tolist"):
        feature_names = feature_names.tolist()
    feature_names = list(feature_names)
    order = np.argsort(imp)[::-1][:top_n]
    fig, ax = plt.subplots(figsize=figsize)
    ax.barh(range(len(order)), imp[order], align="center")
    ax.set_yticks(range(len(order)))
    ax.set_yticklabels([feature_names[i] for i in order], fontsize=9)
    ax.invert_yaxis()
    ax.set_xlabel("Feature importance")
    ax.set_title(title)
    plt.tight_layout()
    plt.show()


def plot_baseline_vs_tuned(metrics_baseline, metrics_tuned, target_name):
    """
    Plot bar charts comparing baseline and tuned model validation metrics.

    Uses two subplots so RMSE and R^2 each have an appropriate scale (RMSE can be
    hundreds, R^2 is 0-1). Expects dicts with keys "RMSE" and "R2" (e.g. from
    regression_models.evaluate_regression_metrics).

    Args:
        metrics_baseline: Dict with at least "RMSE" and "R2".
        metrics_tuned: Dict with at least "RMSE" and "R2".
        target_name: Label for the plot title (e.g. "duration_seconds" or "fare_amount").
    """
    labels = ["Baseline", "Tuned"]
    rmse_vals = [metrics_baseline["RMSE"], metrics_tuned["RMSE"]]
    r2_vals = [metrics_baseline["R2"], metrics_tuned["R2"]]
    x = np.arange(len(labels))
    w = 0.5

    fig, (ax_rmse, ax_r2) = plt.subplots(1, 2, figsize=(10, 4))

    # RMSE: lower is better
    bars_rmse = ax_rmse.bar(x, rmse_vals, width=w, color=["#1f77b4", "#1f77b4"], edgecolor="black", linewidth=0.8)
    ax_rmse.bar_label(bars_rmse, fmt="%.1f", padding=4)
    ax_rmse.set_xticks(x)
    ax_rmse.set_xticklabels(labels)
    ax_rmse.set_ylabel("RMSE")
    ax_rmse.set_title("RMSE (lower is better)")
    ax_rmse.spines["top"].set_visible(False)
    ax_rmse.spines["right"].set_visible(False)

    # R^2: higher is better (scale 0-1, or slightly below 0 if R^2 is negative)
    r2_min = min(r2_vals)
    r2_low = min(0, r2_min) - 0.05 if r2_min < 0 else 0
    r2_high = 1.05
    bars_r2 = ax_r2.bar(x, r2_vals, width=w, color=["#ff7f0e", "#ff7f0e"], edgecolor="black", linewidth=0.8)
    ax_r2.bar_label(bars_r2, fmt="%.3f", padding=4)
    ax_r2.set_xticks(x)
    ax_r2.set_xticklabels(labels)
    ax_r2.set_ylabel("R2")
    ax_r2.set_title("R2 (higher is better)")
    ax_r2.set_ylim(r2_low, r2_high)
    ax_r2.spines["top"].set_visible(False)
    ax_r2.spines["right"].set_visible(False)

    fig.suptitle(f"{target_name} - Baseline vs Tuned (validation)", fontsize=12, y=1.02)
    plt.tight_layout()
    plt.show()