chrismbryant / classification-curves Goto Github PK

from typing import Optional, List

import pandas as pd
from ??? import BinomialCI
import sklearn.metrics
import matplotlib as mpl


import numpy as np
import matplotlib.pyplot as plt

print("Imported functions:")

print(" >> assess_prob_calibration(...)")
def assess_prob_calibration(
    pred_df: pd.DataFrame,
    probability_col: str = "prob",
    label_col: str = "label",
) -> pd.DataFrame:

    pred_df["quantile"] = pred_df[probability_col].rank() / len(pred_df)
    pred_df["quantile_bucket"] = pred_df["quantile"].round(2)
    pred_df["score_bucket"] = pred_df[probability_col].round(2)

    score_prob_calibration = pred_df.groupby("score_bucket").agg(
        avg_prediction=(probability_col, "mean"),
        num_actual_pos=(label_col, "sum"),
        num_examples=(label_col, "count"),
    ).reset_index()

    quantile_prob_calibration = pred_df.groupby("quantile_bucket").agg(
        avg_prediction=(probability_col, "mean"),
        num_actual_pos=(label_col, "sum"),
        num_examples=(label_col, "count"),
    ).reset_index()

    prob_calibration = pd.concat([score_prob_calibration, quantile_prob_calibration])

    prob_calibration["proportion"], prob_calibration["lower"], prob_calibration["upper"] = (
        BinomialCI().get_ci(
            prob_calibration["num_actual_pos"],
            prob_calibration["num_examples"],
        )
    )

    return prob_calibration.sort_values("avg_prediction")


print(" >> plot_probability_calibration(...)")
def plot_probability_calibration(
    prob_calibration: pd.DataFrame,
    plot_confidence_band: bool = False,
    dpi: Optional[int] = None,
    color: str = "tab:blue",
    title: str = "Probability Calibration Curve",
    x_label: str = "Mean Predicted Value",
    y_label: str = "Proportion Labeled Positive",
    return_fig: bool = False,
) -> None:

    fig = plt.figure(figsize = (10, 7), dpi = dpi)
    ax = fig.add_subplot(1, 1, 1, aspect = "equal")
    ax.grid(True)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)

    # Add line of perfect calibration
    ax.plot([0, 1], [0, 1], "k-")

    if plot_confidence_band:

        # Shade confidence band
        ax.fill_between(
            x = prob_calibration["avg_prediction"],
            y1 = prob_calibration["lower"],
            y2 = prob_calibration["upper"],
            alpha = 0.3,
            color = color)
        ax.plot(
            prob_calibration["avg_prediction"],
            prob_calibration["proportion"],
            color = color)

    else:

        # Plot points
        ax.scatter(
            x = prob_calibration["avg_prediction"],
            y = prob_calibration["proportion"],
            s = 30,
            color = color,
            marker = ".",
            zorder = 100)

        # Plot error bars on points
        ax.errorbar(
            x = prob_calibration["avg_prediction"],
            y = prob_calibration["proportion"],
            yerr = [
                prob_calibration["proportion"] - prob_calibration["lower"], 
                prob_calibration["upper"] - prob_calibration["proportion"]],
            alpha = 0.6,
            color = color,
            fmt = ".",
            markersize = 0,
            zorder = 100)

        # Set labels
        ax.set_xlabel(x_label)
        ax.set_ylabel(y_label)
        ax.set_title(title)

    if return_fig:
        return fig, ax


print(" >> plot_predictions(...)")
def plot_predictions(
    df: pd.DataFrame,
    pred_col: str = "pred",
    label_col: str = "label",
    dpi: Optional[int] = None,
    color: str = "tab:blue",
    title: str = "Quantity Predictions",
    x_label: str = "Predicted Quantity",
    y_label: str = "Actual Quantity",
    x_rng: Optional[List[int]] = None,
    x_jitter: float = 0,
    y_jitter: float = 0,
    log_scale: bool = False,
    alpha: float = 0.5,
    size: float = 20,
    vmax: float = 10,
    scatter_density: bool = False,
    cmap: str = "viridis",
    return_fig: bool = False,
) -> None:

    fig = plt.figure(figsize=(10, 7), dpi=dpi)
    ax = fig.add_subplot(1, 1, 1, aspect="equal")
    ax.grid(True, alpha=0.5)

    # Gaussian jitter
    x_jitter = np.random.normal(0, x_jitter, len(df))
    y_jitter = np.random.normal(0, y_jitter, len(df))

    if log_scale:
        x = df[pred_col] * (1 + x_jitter)
        y = df[label_col] * (1 + y_jitter)
    else:
        x = df[pred_col] + x_jitter
        y = df[label_col] + y_jitter

    if scatter_density:
        import mpl_scatter_density
        norm = mpl.colors.Normalize(0, vmax)
        ax = fig.add_subplot(1, 1, 1, projection="scatter_density")
        ax.scatter_density(x, y, cmap=cmap, norm=norm, dpi=dpi)
    else:
        ax.scatter(
            x = x,
            y = y,
            s = size,
            color = color,
            marker = ".",
            ec="none",
            alpha=alpha
        )

    ax.set_xlabel(x_label)
    ax.set_ylabel(y_label)
    ax.set_title(title)
    ax.set_aspect("equal")

    if x_rng is not None:
        ax.set_xlim(x_rng[0], x_rng[1])
        ax.set_ylim(x_rng[0], x_rng[1])
        
        # Add line of perfect calibration
        ax.plot(
            [x_rng[0], x_rng[1]],
            [x_rng[0], x_rng[1]],
            "k-")

    # Include R2, RMSE, and QI as text in plot
    df["lit_0"] = 0
    r2 = sklearn.metrics.r2_score(df[label_col], df[pred_col])
    rmse = np.sqrt(sklearn.metrics.mean_squared_error(df[label_col], df[pred_col]))
    rmsle = np.sqrt(sklearn.metrics.mean_squared_log_error(
        df[[label_col, "lit_0"]].max(axis=1),
        df[[pred_col, "lit_0"]].max(axis=1),
    ))
    qi = df[pred_col].sum() / df[label_col].sum()
    
    ax.text(
        x=0.05,
        y=0.95,
        s=f"      $R^2$: {r2:.3f}\n RMSE: {rmse:.3f}\nRMSLE: {rmsle:.3f}\n      QI: {qi:.3f}",
        horizontalalignment="left",
        verticalalignment="top",
        transform=ax.transAxes,
        fontsize=12
    )

    if log_scale:
        ax.set_yscale("log")
        ax.set_xscale("log")
    
    if return_fig:
        return fig, ax

show_confidence_band: Union[bool, float, List[float]] = False