Create calibration_utils.py

src/calibration_utils.py

@@ -0,0 +1,316 @@
+# src/calibration_utils.py
+"""
+Calibration assessment + bootstrap confidence intervals for binary
+classifiers and Cox survival models.
+ 
+Designed to slot into the existing manuscript pipeline:
+  - call signatures mirror inference_utils.compute_metrics()
+  - returns dicts that downstream code can merge into existing metric dicts
+  - matplotlib figures use the same style as the existing ROC/PR plots
+"""
+ 
+from __future__ import annotations
+ 
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+from sklearn.metrics import roc_auc_score, brier_score_loss
+from sklearn.calibration import calibration_curve
+from scipy.special import logit
+import statsmodels.api as sm
+from lifelines.utils import concordance_index
+ 
+ 
+# ----------------------------------------------------------------------
+# Helpers
+# ----------------------------------------------------------------------
+ 
+def _clean_binary_inputs(y_true, y_prob):
+    """Mirror the NaN-handling pattern used in inference_utils.compute_metrics()."""
+    y_true = pd.to_numeric(pd.Series(y_true).astype(str).str.strip(), errors="coerce")
+    y_prob = pd.to_numeric(pd.Series(y_prob), errors="coerce")
+    valid = y_true.notna() & y_prob.notna()
+    y_true = y_true.loc[valid].astype(int).to_numpy()
+    y_prob = y_prob.loc[valid].astype(float).to_numpy()
+    return y_true, y_prob
+ 
+ 
+def _safe_logit(p, eps=1e-6):
+    return logit(np.clip(p, eps, 1.0 - eps))
+ 
+ 
+# ----------------------------------------------------------------------
+# Bootstrap CI helpers
+# ----------------------------------------------------------------------
+ 
+def bootstrap_auroc_ci(y_true, y_prob, n_bootstraps=1000, seed=42):
+    """
+    Returns (auroc_point, ci_low, ci_high).
+    Stratified bootstrap (preserves event prevalence).
+    """
+    y_true, y_prob = _clean_binary_inputs(y_true, y_prob)
+    if len(y_true) == 0 or len(np.unique(y_true)) < 2:
+        return (np.nan, np.nan, np.nan)
+ 
+    point = float(roc_auc_score(y_true, y_prob))
+    pos_idx = np.where(y_true == 1)[0]
+    neg_idx = np.where(y_true == 0)[0]
+    rng = np.random.default_rng(seed)
+ 
+    boot = []
+    for _ in range(n_bootstraps):
+        pos_b = rng.choice(pos_idx, size=len(pos_idx), replace=True)
+        neg_b = rng.choice(neg_idx, size=len(neg_idx), replace=True)
+        idx = np.concatenate([pos_b, neg_b])
+        try:
+            boot.append(roc_auc_score(y_true[idx], y_prob[idx]))
+        except ValueError:
+            continue
+ 
+    if len(boot) == 0:
+        return (point, np.nan, np.nan)
+    lo, hi = np.percentile(boot, [2.5, 97.5])
+    return (point, float(lo), float(hi))
+ 
+ 
+def bootstrap_c_index_ci(durations, events, risk_scores, n_bootstraps=1000, seed=42):
+    """
+    Bootstrap C-index for a Cox-style risk score (higher score = higher risk).
+    Returns (c_point, ci_low, ci_high).
+    """
+    durations = np.asarray(durations, dtype=float).ravel()
+    events = np.asarray(events, dtype=int).ravel()
+    risk_scores = np.asarray(risk_scores, dtype=float).ravel()
+ 
+    valid = ~(np.isnan(durations) | np.isnan(risk_scores)) & (durations > 0)
+    durations = durations[valid]
+    events = events[valid]
+    risk_scores = risk_scores[valid]
+ 
+    if len(durations) < 10 or events.sum() < 5:
+        return (np.nan, np.nan, np.nan)
+ 
+    try:
+        point = float(concordance_index(durations, -risk_scores, events))
+    except Exception:
+        return (np.nan, np.nan, np.nan)
+ 
+    rng = np.random.default_rng(seed)
+    n = len(durations)
+    boot = []
+    for _ in range(n_bootstraps):
+        idx = rng.choice(n, size=n, replace=True)
+        if np.asarray(events)[idx].sum() < 2:
+            continue
+        try:
+            boot.append(concordance_index(durations[idx], -risk_scores[idx], events[idx]))
+        except Exception:
+            continue
+ 
+    if len(boot) == 0:
+        return (point, np.nan, np.nan)
+    lo, hi = np.percentile(boot, [2.5, 97.5])
+    return (point, float(lo), float(hi))
+ 
+ 
+# ----------------------------------------------------------------------
+# Core calibration stats
+# ----------------------------------------------------------------------
+ 
+def compute_calibration_stats(y_true, y_prob, n_bins=10):
+    """
+    Compute calibration intercept, slope, decile-level points, and Brier score.
+ 
+    Returns a dict with the same flat-key style as compute_metrics():
+        {
+            'N': int, 'Events': int, 'Prevalence': float,
+            'Brier': float, 'CalibrationInTheLarge': float,
+            'CalibrationIntercept': float,
+            'CalibrationIntercept_CI_low': float,
+            'CalibrationIntercept_CI_high': float,
+            'CalibrationSlope': float,
+            'CalibrationSlope_CI_low': float,
+            'CalibrationSlope_CI_high': float,
+            'BinProbTrue': np.ndarray,   # observed event rate per bin
+            'BinProbPred': np.ndarray,   # mean predicted prob per bin
+            'BinCounts':   np.ndarray,   # N per bin
+        }
+    """
+    y_true, y_prob = _clean_binary_inputs(y_true, y_prob)
+ 
+    out = {
+        "N": int(len(y_true)),
+        "Events": int(y_true.sum()) if len(y_true) else 0,
+        "Prevalence": float(y_true.mean()) if len(y_true) else np.nan,
+        "Brier": np.nan,
+        "CalibrationInTheLarge": np.nan,
+        "CalibrationIntercept": np.nan,
+        "CalibrationIntercept_CI_low": np.nan,
+        "CalibrationIntercept_CI_high": np.nan,
+        "CalibrationSlope": np.nan,
+        "CalibrationSlope_CI_low": np.nan,
+        "CalibrationSlope_CI_high": np.nan,
+        "BinProbTrue": np.array([]),
+        "BinProbPred": np.array([]),
+        "BinCounts": np.array([]),
+    }
+ 
+    if len(y_true) < 20 or len(np.unique(y_true)) < 2:
+        return out
+ 
+    # Brier
+    out["Brier"] = float(brier_score_loss(y_true, np.clip(y_prob, 1e-15, 1 - 1e-15)))
+ 
+    # Calibration-in-the-large
+    out["CalibrationInTheLarge"] = float(y_true.mean() - y_prob.mean())
+ 
+    # Decile points
+    try:
+        prob_true, prob_pred = calibration_curve(y_true, y_prob, n_bins=n_bins, strategy="quantile")
+        # Bin counts via quantile cut on predicted probabilities
+        bin_edges = np.quantile(y_prob, np.linspace(0, 1, n_bins + 1))
+        bin_edges[0] -= 1e-9
+        bin_edges[-1] += 1e-9
+        bin_idx = np.digitize(y_prob, bin_edges, right=True) - 1
+        bin_idx = np.clip(bin_idx, 0, n_bins - 1)
+        bin_counts = np.bincount(bin_idx, minlength=n_bins)[: len(prob_pred)]
+ 
+        out["BinProbTrue"] = prob_true
+        out["BinProbPred"] = prob_pred
+        out["BinCounts"] = bin_counts
+    except Exception:
+        pass
+ 
+    # Calibration intercept (offset) and slope via logistic recalibration:
+    #   logit(p_obs) = intercept + slope * logit(p_pred)
+    try:
+        logits = _safe_logit(y_prob)
+        X = sm.add_constant(logits)
+        result = sm.Logit(y_true, X).fit(disp=0)
+        intercept = float(result.params[0])
+        slope = float(result.params[1])
+        ci = result.conf_int()
+        out["CalibrationIntercept"] = intercept
+        out["CalibrationIntercept_CI_low"] = float(ci.iloc[0, 0])
+        out["CalibrationIntercept_CI_high"] = float(ci.iloc[0, 1])
+        out["CalibrationSlope"] = slope
+        out["CalibrationSlope_CI_low"] = float(ci.iloc[1, 0])
+        out["CalibrationSlope_CI_high"] = float(ci.iloc[1, 1])
+    except Exception:
+        pass
+ 
+    return out
+ 
+ 
+# ----------------------------------------------------------------------
+# Plotting
+# ----------------------------------------------------------------------
+ 
+def plot_calibration_curve(y_true, y_prob, title="Calibration", n_bins=10,
+                           ax=None, return_stats=False):
+    """
+    Calibration plot with:
+      - decile points (observed vs predicted), marker size proportional to N per bin
+      - diagonal reference line (perfect calibration)
+      - histogram of predicted probabilities along the bottom
+      - intercept and slope (with 95% CIs) annotated in the title
+      - Brier score annotated in the legend
+ 
+    If ax is provided, draws into that axis. Otherwise creates a new fig/ax.
+    Returns the matplotlib figure (and stats dict if return_stats=True).
+    """
+    stats = compute_calibration_stats(y_true, y_prob, n_bins=n_bins)
+ 
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(6.5, 6.0))
+    else:
+        fig = ax.figure
+ 
+    # Diagonal reference
+    ax.plot([0, 1], [0, 1], linestyle="--", color="0.4", label="Perfect calibration")
+ 
+    # Decile points with marker size proportional to bin count
+    prob_true = stats["BinProbTrue"]
+    prob_pred = stats["BinProbPred"]
+    bin_counts = stats["BinCounts"]
+ 
+    if len(prob_true) and len(prob_pred):
+        if len(bin_counts) == len(prob_pred) and bin_counts.sum() > 0:
+            sizes = 40 + 200 * (bin_counts / bin_counts.max())
+        else:
+            sizes = np.full(len(prob_pred), 80.0)
+ 
+        ax.scatter(prob_pred, prob_true, s=sizes, color="#1F77B4",
+                   edgecolor="white", linewidth=0.8, zorder=3, label="Model")
+        ax.plot(prob_pred, prob_true, color="#1F77B4", alpha=0.5, zorder=2)
+ 
+    # Annotation
+    icpt = stats["CalibrationIntercept"]
+    icpt_lo = stats["CalibrationIntercept_CI_low"]
+    icpt_hi = stats["CalibrationIntercept_CI_high"]
+    slope = stats["CalibrationSlope"]
+    slope_lo = stats["CalibrationSlope_CI_low"]
+    slope_hi = stats["CalibrationSlope_CI_high"]
+    brier = stats["Brier"]
+ 
+    def fmt(v):
+        return "NA" if (v is None or np.isnan(v)) else f"{v:.2f}"
+ 
+    full_title = (
+        f"{title}\n"
+        f"Intercept = {fmt(icpt)} ({fmt(icpt_lo)} to {fmt(icpt_hi)})    "
+        f"Slope = {fmt(slope)} ({fmt(slope_lo)} to {fmt(slope_hi)})    "
+        f"Brier = {fmt(brier)}"
+    )
+    ax.set_title(full_title, fontsize=10)
+    ax.set_xlabel("Mean predicted probability")
+    ax.set_ylabel("Observed event rate")
+    ax.set_xlim(-0.02, 1.02)
+    ax.set_ylim(-0.02, 1.02)
+    ax.grid(alpha=0.3, linestyle=":")
+    ax.legend(loc="upper left", fontsize=9, framealpha=0.85)
+ 
+    # Histogram of predicted probabilities at the bottom (twin axis)
+    y_true_arr, y_prob_arr = _clean_binary_inputs(y_true, y_prob)
+    if len(y_prob_arr) > 0:
+        ax2 = ax.twinx()
+        ax2.hist(y_prob_arr, bins=30, range=(0, 1), color="0.7", alpha=0.5,
+                 edgecolor="white", linewidth=0.3)
+        ax2.set_ylabel("Count (predicted prob)", fontsize=9, color="0.5")
+        ax2.tick_params(axis="y", labelsize=8, colors="0.5")
+        # Keep histogram in the bottom third
+        hist_max = ax2.get_ylim()[1]
+        ax2.set_ylim(0, hist_max * 3)
+        ax2.set_zorder(0)
+        ax.set_zorder(1)
+        ax.patch.set_alpha(0)
+ 
+    fig.tight_layout()
+    if return_stats:
+        return fig, stats
+    return fig
+ 
+ 
+def plot_multi_cohort_calibration(cohort_results, ncols=3, figsize=None):
+    """
+    Composite calibration figure across cohorts. Use this for the
+    supplementary figure referenced in the manuscript revision roadmap.
+ 
+    cohort_results: dict of {cohort_name: (y_true, y_prob)}
+    """
+    n = len(cohort_results)
+    if n == 0:
+        return None
+    nrows = int(np.ceil(n / ncols))
+    if figsize is None:
+        figsize = (5 * ncols, 4.5 * nrows)
+    fig, axes = plt.subplots(nrows, ncols, figsize=figsize, squeeze=False)
+    for i, (cohort_name, (y_true, y_prob)) in enumerate(cohort_results.items()):
+        r, c = divmod(i, ncols)
+        plot_calibration_curve(y_true, y_prob, title=cohort_name, ax=axes[r][c])
+    # Hide unused panels
+    for j in range(n, nrows * ncols):
+        r, c = divmod(j, ncols)
+        axes[r][c].axis("off")
+    fig.tight_layout()
+    return fig