scripts/compute_ece.py

#!/usr/bin/env python3
"""
Compute ECE (Expected Calibration Error) for each ablation variant
using Temperature Scaling calibrated on the val set.
Plots:
  - Reliability diagram (calibration curve) per variant
  - ECE bar chart across all variants
Data source: outputs/hfm_run1/clara_hfm.pt  (best run = hfm_run1_full_up)
"""
from __future__ import annotations
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

import numpy as np
import torch
from scipy.optimize import minimize_scalar
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from transformers import CLIPProcessor, DebertaV2Tokenizer

from src.hfm_pipeline import (
    HFMLoader,
    create_dataloaders,
    estimate_max_length,
    gather_logits_variant,
    load_checkpoint,
    resolve_device,
)

# ── Config ────────────────────────────────────────────────────────────────────
CHECKPOINT  = str(PROJECT_ROOT / "outputs/hfm_run1/clara_hfm.pt")
DATA_ROOT   = str(PROJECT_ROOT / "data/HFM")
TEXT_DIR    = str(PROJECT_ROOT / "data/HFM/text")
OUT_DIR     = PROJECT_ROOT / "results/hfm_run1_full_up"
N_BINS      = 15   # calibration bins

VARIANT_MAP = [
    ("Full",             "full"),
    ("w/o Verification", "w/o_verification"),
    ("w/o Feedback",     "w/o_feedback"),
    ("w/o Co-Attention", "w/o_coattn"),
    ("Text-only",        "text_only"),
    ("Vision-only",      "vision_only"),
    ("w/o Text",         "w/o_text"),
    ("w/o Image",        "w/o_image"),
]

# ── ECE helpers ───────────────────────────────────────────────────────────────
def compute_ece(probs: np.ndarray, y_true: np.ndarray, n_bins: int = 15):
    confidence = probs.max(axis=1)
    y_pred     = probs.argmax(axis=1)
    correct    = (y_pred == y_true).astype(float)

    bins = np.linspace(0.0, 1.0, n_bins + 1)
    bin_conf  = np.zeros(n_bins)
    bin_acc   = np.zeros(n_bins)
    bin_count = np.zeros(n_bins, dtype=int)

    for i in range(n_bins):
        lo, hi = bins[i], bins[i + 1]
        mask = (confidence >= lo) & (confidence <= hi) if i == 0 else \
               (confidence > lo) & (confidence <= hi)
        cnt = mask.sum()
        if cnt > 0:
            bin_conf[i]  = confidence[mask].mean()
            bin_acc[i]   = correct[mask].mean()
            bin_count[i] = cnt

    N   = len(y_true)
    ece = float(np.sum(bin_count / N * np.abs(bin_conf - bin_acc)))
    return ece, bin_conf, bin_acc, bin_count


def temperature_scale(logits: np.ndarray, y_true: np.ndarray) -> float:
    """Find optimal temperature T on given logits/labels via NLL minimization."""
    def nll(log_T):
        T = np.exp(log_T)
        scaled = logits / T
        # log-sum-exp trick
        shifted = scaled - scaled.max(axis=1, keepdims=True)
        log_probs = shifted - np.log(np.exp(shifted).sum(axis=1, keepdims=True))
        return -log_probs[np.arange(len(y_true)), y_true].mean()

    result = minimize_scalar(nll, bounds=(-3.0, 3.0), method="bounded")
    return float(np.exp(result.x))


# ── Main ──────────────────────────────────────────────────────────────────────
def main():
    device = resolve_device("auto")
    print(f"Device: {device}")

    model, cfg, _ = load_checkpoint(CHECKPOINT, device)
    cfg["image_root"] = DATA_ROOT
    cfg["text_dir"]   = TEXT_DIR

    clip_proc = CLIPProcessor.from_pretrained(cfg["vision_model_id"])
    tokenizer = DebertaV2Tokenizer.from_pretrained(cfg["text_model_id"])

    loader_obj   = HFMLoader(TEXT_DIR, DATA_ROOT)
    all_samples  = loader_obj.load()
    train_s      = loader_obj.get_split("train")
    val_s        = loader_obj.get_split("val")
    test_s       = loader_obj.get_split("test")

    max_len = cfg.get("max_length") or int(estimate_max_length(all_samples))
    cfg["max_length"] = max_len

    _, val_loader, test_loader = create_dataloaders(
        train_samples=train_s,  val_samples=val_s, test_samples=test_s,
        clip_processor=clip_proc, tokenizer=tokenizer,
        batch_size=cfg.get("batch_size", 64),
        max_length=max_len,
        num_workers=cfg.get("num_workers", 0),
        pin_memory=False,
        weighted_train_sampler=False,
    )

    # ── Compute ECE with Temperature Scaling per variant ─────────────────────
    results = {}
    for display, key in VARIANT_MAP:
        print(f"  [{key}] calibrating on val ...")
        val_logits,  val_labels  = gather_logits_variant(model, val_loader,  device, key)
        print(f"  [{key}] inferring on test ...")
        test_logits, test_labels = gather_logits_variant(model, test_loader, device, key)

        T = temperature_scale(val_logits, val_labels)
        cal_probs = np.exp(test_logits / T - np.log(
            np.exp(test_logits / T).sum(axis=1, keepdims=True)))

        ece, bin_conf, bin_acc, bin_cnt = compute_ece(cal_probs, test_labels, N_BINS)
        results[display] = dict(ece=ece, bin_conf=bin_conf,
                                bin_acc=bin_acc, bin_cnt=bin_cnt, T=T)
        print(f"    T={T:.4f}  ECE = {ece:.4f}")

    OUT_DIR.mkdir(parents=True, exist_ok=True)

    # ── Plot 1: Reliability diagrams (2×3 grid) ───────────────────────────────
    COLORS = {
        "Full":             "#2c6e9e",
        "w/o Verification": "#e07b39",
        "w/o Feedback":     "#3aaa5e",
        "w/o Co-Attention": "#9b59b6",
        "Text-only":        "#c0392b",
        "Vision-only":      "#7f8c8d",
        "w/o Text":         "#8e44ad",
        "w/o Image":        "#16a085",
    }

    n_cols = 3
    n_rows = int(np.ceil(len(VARIANT_MAP) / n_cols))
    fig, axes = plt.subplots(n_rows, n_cols, figsize=(12, 3.8 * n_rows))
    axes = np.atleast_1d(axes).flatten()
    bin_edges = np.linspace(0, 1, N_BINS + 1)
    bin_centers = 0.5 * (bin_edges[:-1] + bin_edges[1:])
    bin_width = bin_edges[1] - bin_edges[0]

    for ax, (display, _) in zip(axes, VARIANT_MAP):
        r = results[display]
        color = COLORS[display]

        # perfect calibration diagonal
        ax.plot([0, 1], [0, 1], "k--", lw=1.2, label="Perfect calibration")

        # gap (overconfidence / underconfidence)
        for i in range(N_BINS):
            if r["bin_cnt"][i] > 0:
                lo = min(r["bin_conf"][i], r["bin_acc"][i])
                hi = max(r["bin_conf"][i], r["bin_acc"][i])
                ax.bar(bin_centers[i], hi - lo, bottom=lo,
                       width=bin_width * 0.9, color="tomato", alpha=0.35)

        # actual accuracy bars
        mask = r["bin_cnt"] > 0
        ax.bar(bin_centers[mask], r["bin_acc"][mask],
               width=bin_width * 0.9, color=color, alpha=0.85, label="Accuracy")

        ax.set_xlim(0, 1); ax.set_ylim(0, 1)
        ax.set_xticks([0, 0.25, 0.5, 0.75, 1.0])
        ax.set_yticks([0, 0.25, 0.5, 0.75, 1.0])
        ax.set_xlabel("Confidence", fontsize=9)
        ax.set_ylabel("Accuracy",   fontsize=9)
        ax.set_title(f"{display}\nECE = {r['ece']:.4f}  (T={r['T']:.3f})", fontsize=10, fontweight="bold")
        ax.legend(fontsize=7, loc="upper left")

    for ax in axes[len(VARIANT_MAP) :]:
        ax.axis("off")

    fig.suptitle("Reliability Diagrams after Temperature Scaling — HFM Ablation", fontsize=13, y=1.01)
    fig.tight_layout()
    p_diag = OUT_DIR / "figure_ece_reliability.png"
    fig.savefig(p_diag, dpi=150, bbox_inches="tight")
    print(f"Saved → {p_diag}")
    plt.close(fig)

    # ── Plot 2: ECE bar chart ─────────────────────────────────────────────────
    names = [d for d, _ in VARIANT_MAP]
    eces  = [results[d]["ece"] for d in names]
    colors_bar = [COLORS[d] for d in names]

    fig2, ax2 = plt.subplots(figsize=(8, 4.5))
    bars = ax2.bar(names, eces, color=colors_bar, edgecolor="white", width=0.55)

    # value labels on bars
    for bar, val in zip(bars, eces):
        ax2.text(bar.get_x() + bar.get_width() / 2,
                 bar.get_height() + 0.0012,
                 f"{val:.4f}", ha="center", va="bottom", fontsize=9.5, fontweight="bold")

    ax2.set_ylabel("ECE ↓", fontsize=11)
    ax2.set_ylim(0, max(eces) * 1.22)
    ax2.set_xticks(range(len(names)))
    ax2.set_xticklabels(names, rotation=20, ha="right", fontsize=10)
    ax2.set_title("ECE after Temperature Scaling — HFM Ablation", fontsize=12, pad=10)
    ax2.axhline(eces[0], color=COLORS["Full"], lw=1.3, ls="--", alpha=0.7,
                label=f"Full ECE (cal) = {eces[0]:.4f}")
    ax2.legend(fontsize=9)
    ax2.spines["top"].set_visible(False)
    ax2.spines["right"].set_visible(False)
    fig2.tight_layout()

    p_bar = OUT_DIR / "figure_ece_bar.png"
    fig2.savefig(p_bar, dpi=150, bbox_inches="tight")
    print(f"Saved → {p_bar}")
    plt.close(fig2)

    print("\n── ECE Summary (after Temperature Scaling) ──────────────────────────────")
    for d, _ in VARIANT_MAP:
        print(f"  {d:<20} T={results[d]['T']:.4f}  ECE = {results[d]['ece']:.4f}")


if __name__ == "__main__":
    main()