Log pairwise calibration ECE for learned selector

workspace/scripts/eval_learned_dominance_selector.py

@@ -19,7 +19,7 @@ if str(PROJECT_ROOT) not in sys.path:
 import numpy as np  # noqa: E402
 
 from cil.chart_features import CHART_FEATURE_MODES, OBJECT_LAYOUT_EMBED_DIM, OBSERVATION_EMBED_DIM  # noqa: E402
-from cil.metrics import macro_micro_summary  # noqa: E402
+from cil.metrics import macro_micro_summary, pairwise_causal_dominance_ece  # noqa: E402
 from scripts.eval_dominance_selector import _DominanceScorer, _chart_map, _first_train_seed, _rows  # noqa: E402
 from scripts.eval_ctt_generated_rollout import load_chart_items  # noqa: E402
 
@@ -256,13 +256,28 @@ def main(argv: list[str] | None = None) -> int:
         fit_objective=args.fit_objective,
         pairwise_weight=args.pairwise_weight,
     )
-    eval_cases = _evaluate_dataset(eval_dataset, best["weights"], best["mean"], best["std"], tau=best["tau"])
-    calibration_cases = _evaluate_dataset(
+    eval_predictions = _linear_predictions(eval_dataset, best["weights"], best["mean"], best["std"])
+    calibration_predictions = _linear_predictions(
         calibration_dataset,
         best["weights"],
         best["mean"],
         best["std"],
+    )
+    eval_pairwise = _pairwise_calibration_summary(eval_dataset, eval_predictions)
+    calibration_pairwise = _pairwise_calibration_summary(calibration_dataset, calibration_predictions)
+    eval_cases = _evaluate_predictions(
+        eval_dataset,
+        eval_predictions,
+        tau=best["tau"],
+        include_pairwise_calibration=True,
+        pairwise_calibration=eval_pairwise,
+    )
+    calibration_cases = _evaluate_predictions(
+        calibration_dataset,
+        calibration_predictions,
         tau=best["tau"],
+        include_pairwise_calibration=True,
+        pairwise_calibration=calibration_pairwise,
     )
     metric_names = sorted(
         {
@@ -326,8 +341,12 @@ def main(argv: list[str] | None = None) -> int:
         "num_calibration_candidates": len(calibration_dataset["samples"]),
         "num_eval_candidates": len(eval_dataset["samples"]),
         "calibration_model_selection": best["selection"],
-        "calibration_summary": _simple_summary(calibration_cases),
-        "eval_summary": _simple_summary(eval_cases),
+        "calibration_summary": _summary_with_pairwise(calibration_cases, calibration_pairwise),
+        "eval_summary": _summary_with_pairwise(eval_cases, eval_pairwise),
+        "pairwise_causal_calibration": {
+            "calibration": _pairwise_calibration_global(calibration_pairwise),
+            "eval": _pairwise_calibration_global(eval_pairwise),
+        },
         "summary": summary,
         "rows": eval_cases,
     }
@@ -1079,11 +1098,27 @@ def _evaluate_dataset(
     std: np.ndarray,
     *,
     tau: float | dict[str, float],
+    include_pairwise_calibration: bool = False,
 ) -> list[dict[str, Any]]:
+    predictions = _linear_predictions(dataset, weights, mean, std)
+    return _evaluate_predictions(
+        dataset,
+        predictions,
+        tau=tau,
+        include_pairwise_calibration=include_pairwise_calibration,
+    )
+
+
+def _linear_predictions(
+    dataset: dict[str, Any],
+    weights: np.ndarray,
+    mean: np.ndarray,
+    std: np.ndarray,
+) -> np.ndarray:
     x = np.stack([sample["feature"] for sample in dataset["samples"]], axis=0)
     x_norm = (x - mean) / std
     x_norm[:, 0] = 1.0
-    return _evaluate_predictions(dataset, x_norm @ weights, tau=tau)
+    return x_norm @ weights
 
 
 def _evaluate_predictions(
@@ -1091,8 +1126,13 @@ def _evaluate_predictions(
     predictions: np.ndarray,
     *,
     tau: float | dict[str, float],
+    include_pairwise_calibration: bool = False,
+    pairwise_calibration: dict[str, Any] | None = None,
 ) -> list[dict[str, Any]]:
     samples = dataset["samples"]
+    if include_pairwise_calibration and pairwise_calibration is None:
+        pairwise_calibration = _pairwise_calibration_summary(dataset, predictions)
+    pairwise_rows = (pairwise_calibration or {}).get("rows", {})
     rows: list[dict[str, Any]] = []
     for row_index, sample_indices in sorted(dataset["by_row"].items()):
         best_index = max(sample_indices, key=lambda index: float(predictions[index]))
@@ -1108,40 +1148,138 @@ def _evaluate_predictions(
         )
         hidden_utility = float(sample["hidden_chart_oracle_utility"])
         hidden_success = float(sample["hidden_chart_oracle_success"])
-        rows.append(
+        row = {
+            "chart_id": sample["chart_id"],
+            "task_id": sample["task_id"],
+            "seed": sample["seed"],
+            "train_seed": sample["train_seed"],
+            "selected_candidate_index": int(sample["candidate_index"]),
+            "predicted_margin": predicted_margin,
+            "tau": row_tau,
+            "execute_generated": float(execute),
+            "coverage": float(execute),
+            "fallback_rate": float(not execute),
+            "base_utility": float(sample["base_utility"]),
+            "base_success": float(sample["base_success"]),
+            "selected_utility": selected_utility,
+            "selected_success": selected_success,
+            "selected_utility_gain_over_base": selected_utility - float(sample["base_utility"]),
+            "selected_success_gain_over_base": selected_success - float(sample["base_success"]),
+            "proposal_oracle_utility": float(sample["proposal_oracle_utility"]),
+            "proposal_oracle_success": float(sample["proposal_oracle_success"]),
+            "hidden_chart_oracle_utility": hidden_utility,
+            "hidden_chart_oracle_success": hidden_success,
+            "outcome_ptr": float(sample["outcome_ptr"]),
+            "selector_regret": max(0.0, float(sample["proposal_oracle_utility"]) - selected_utility),
+            "success_selector_gap": max(0.0, float(sample["proposal_oracle_success"]) - selected_success),
+            "support_gap": max(0.0, hidden_utility - float(sample["proposal_oracle_utility"]))
+            if math.isfinite(hidden_utility)
+            else math.nan,
+            "success_support_gap": max(0.0, hidden_success - float(sample["proposal_oracle_success"]))
+            if math.isfinite(hidden_success)
+            else math.nan,
+        }
+        if include_pairwise_calibration:
+            calibration = pairwise_rows.get(row_index) or pairwise_rows.get(str(row_index), {})
+            row.update(_pairwise_calibration_scalars(calibration))
+        rows.append(row)
+    return rows
+
+
+def _pairwise_calibration_summary(
+    dataset: dict[str, Any],
+    predictions: np.ndarray,
+    *,
+    n_bins: int = 10,
+) -> dict[str, Any]:
+    if len(predictions) != len(dataset["samples"]):
+        raise ValueError("predictions must align with dataset samples")
+    bins = [
+        {
+            "count": 0,
+            "accuracy_sum": 0.0,
+            "confidence_sum": 0.0,
+            "lower": index / n_bins,
+            "upper": (index + 1) / n_bins,
+        }
+        for index in range(n_bins)
+    ]
+    by_row: dict[int, dict[str, Any]] = {}
+    total_pairs = 0
+    correct_sum = 0.0
+    confidence_sum = 0.0
+    samples = dataset["samples"]
+    for row_index, sample_indices in sorted(dataset["by_row"].items()):
+        row_scores = [float(predictions[index]) for index in sample_indices]
+        row_utilities = [float(samples[index]["candidate_utility"]) for index in sample_indices]
+        row_metrics = pairwise_causal_dominance_ece(row_scores, row_utilities, n_bins=n_bins)
+        by_row[int(row_index)] = row_metrics
+        row_pairs = int(row_metrics.get("num_pairs") or 0)
+        if row_pairs <= 0:
+            continue
+        total_pairs += row_pairs
+        correct_sum += float(row_metrics.get("accuracy") or 0.0) * row_pairs
+        confidence_sum += float(row_metrics.get("mean_confidence") or 0.0) * row_pairs
+        for index, row_bin in enumerate(row_metrics.get("bins", [])):
+            if index >= len(bins):
+                break
+            count = int(row_bin.get("count") or 0)
+            bins[index]["count"] += count
+            bins[index]["accuracy_sum"] += float(row_bin.get("accuracy") or 0.0) * count
+            bins[index]["confidence_sum"] += float(row_bin.get("confidence") or 0.0) * count
+
+    ece = 0.0
+    rendered_bins: list[dict[str, float | int]] = []
+    for bucket in bins:
+        count = int(bucket["count"])
+        accuracy = bucket["accuracy_sum"] / count if count else 0.0
+        confidence = bucket["confidence_sum"] / count if count else 0.0
+        if total_pairs:
+            ece += (count / total_pairs) * abs(accuracy - confidence)
+        rendered_bins.append(
             {
-                "chart_id": sample["chart_id"],
-                "task_id": sample["task_id"],
-                "seed": sample["seed"],
-                "train_seed": sample["train_seed"],
-                "selected_candidate_index": int(sample["candidate_index"]),
-                "predicted_margin": predicted_margin,
-                "tau": row_tau,
-                "execute_generated": float(execute),
-                "coverage": float(execute),
-                "fallback_rate": float(not execute),
-                "base_utility": float(sample["base_utility"]),
-                "base_success": float(sample["base_success"]),
-                "selected_utility": selected_utility,
-                "selected_success": selected_success,
-                "selected_utility_gain_over_base": selected_utility - float(sample["base_utility"]),
-                "selected_success_gain_over_base": selected_success - float(sample["base_success"]),
-                "proposal_oracle_utility": float(sample["proposal_oracle_utility"]),
-                "proposal_oracle_success": float(sample["proposal_oracle_success"]),
-                "hidden_chart_oracle_utility": hidden_utility,
-                "hidden_chart_oracle_success": hidden_success,
-                "outcome_ptr": float(sample["outcome_ptr"]),
-                "selector_regret": max(0.0, float(sample["proposal_oracle_utility"]) - selected_utility),
-                "success_selector_gap": max(0.0, float(sample["proposal_oracle_success"]) - selected_success),
-                "support_gap": max(0.0, hidden_utility - float(sample["proposal_oracle_utility"]))
-                if math.isfinite(hidden_utility)
-                else math.nan,
-                "success_support_gap": max(0.0, hidden_success - float(sample["proposal_oracle_success"]))
-                if math.isfinite(hidden_success)
-                else math.nan,
+                "lower": float(bucket["lower"]),
+                "upper": float(bucket["upper"]),
+                "count": count,
+                "accuracy": accuracy,
+                "confidence": confidence,
+                "abs_gap": abs(accuracy - confidence),
             }
         )
-    return rows
+    return {
+        "n_bins": int(n_bins),
+        "num_rows": len(dataset["by_row"]),
+        "ece": ece if total_pairs else math.nan,
+        "num_pairs": int(total_pairs),
+        "accuracy": correct_sum / total_pairs if total_pairs else math.nan,
+        "mean_confidence": confidence_sum / total_pairs if total_pairs else math.nan,
+        "bins": rendered_bins,
+        "rows": by_row,
+    }
+
+
+def _pairwise_calibration_scalars(calibration: dict[str, Any]) -> dict[str, float]:
+    return {
+        "pairwise_causal_calibration_ece": _finite_or_nan(calibration.get("ece")),
+        "pairwise_causal_calibration_pairs": float(calibration.get("num_pairs") or 0),
+        "pairwise_causal_calibration_accuracy": _finite_or_nan(calibration.get("accuracy")),
+        "pairwise_causal_calibration_confidence": _finite_or_nan(
+            calibration.get("mean_confidence")
+        ),
+    }
+
+
+def _pairwise_calibration_global(calibration: dict[str, Any]) -> dict[str, Any]:
+    return {key: value for key, value in calibration.items() if key != "rows"}
+
+
+def _summary_with_pairwise(
+    rows: list[dict[str, Any]],
+    pairwise_calibration: dict[str, Any],
+) -> dict[str, float | None]:
+    summary = _simple_summary(rows)
+    summary.update(_pairwise_calibration_scalars(pairwise_calibration))
+    return summary
 
 
 def _tau_for_sample(sample: dict[str, Any], tau: float | dict[str, float]) -> float:
@@ -1244,6 +1382,10 @@ def _mean(values: list[Any]) -> float | None:
     return sum(clean) / len(clean) if clean else None
 
 
+def _finite_or_nan(value: Any) -> float:
+    return float(value) if isinstance(value, (int, float)) and math.isfinite(float(value)) else math.nan
+
+
 def _source_rank(value: Any) -> float:
     match = re.search(r"rank(\d+)", str(value))
     return float(match.group(1)) if match else 0.0
@@ -1253,15 +1395,16 @@ def _table(metrics: dict[str, Any]) -> str:
     summary = metrics["eval_summary"]
     lines = [
         "% Auto-generated by scripts/eval_learned_dominance_selector.py",
-        "\\begin{tabular}{lrrrrrrrr}",
+        "\\begin{tabular}{lrrrrrrrrr}",
         "\\toprule",
-        "Rows & Coverage & Fallback & Base succ. & Selected succ. & Oracle succ. & OutcomePTR & Succ. support gap & Succ. selector gap \\\\",
+        "Rows & Coverage & Fallback & Base succ. & Selected succ. & Oracle succ. & OutcomePTR & Succ. support gap & Succ. selector gap & Cal. ECE \\\\",
         "\\midrule",
         f"{metrics['num_eval_rows']} & {_fmt(summary.get('coverage'))} & "
         f"{_fmt(summary.get('fallback_rate'))} & {_fmt(summary.get('base_success'))} & "
         f"{_fmt(summary.get('selected_success'))} & {_fmt(summary.get('proposal_oracle_success'))} & "
         f"{_fmt(summary.get('outcome_ptr'))} & {_fmt(summary.get('success_support_gap'))} & "
-        f"{_fmt(summary.get('success_selector_gap'))} \\\\",
+        f"{_fmt(summary.get('success_selector_gap'))} & "
+        f"{_fmt(summary.get('pairwise_causal_calibration_ece'))} \\\\",
         "\\bottomrule",
         "\\end{tabular}",
     ]
@@ -1285,16 +1428,18 @@ def _report(metrics: dict[str, Any]) -> str:
         "",
         "The ridge calibrator and threshold are fit on calibration measured rows only. Eval outcomes are used only for reporting.",
         "",
-        "| Split | Coverage | Fallback | Base success | Selected success | Proposal oracle | OutcomePTR | Success support gap | Success selector gap |",
-        "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |",
+        "| Split | Coverage | Fallback | Base success | Selected success | Proposal oracle | OutcomePTR | Success support gap | Success selector gap | Calibration ECE |",
+        "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |",
         f"| calibration | {_fmt(calibration.get('coverage'))} | {_fmt(calibration.get('fallback_rate'))} | "
         f"{_fmt(calibration.get('base_success'))} | {_fmt(calibration.get('selected_success'))} | "
         f"{_fmt(calibration.get('proposal_oracle_success'))} | {_fmt(calibration.get('outcome_ptr'))} | "
-        f"{_fmt(calibration.get('success_support_gap'))} | {_fmt(calibration.get('success_selector_gap'))} |",
+        f"{_fmt(calibration.get('success_support_gap'))} | {_fmt(calibration.get('success_selector_gap'))} | "
+        f"{_fmt(calibration.get('pairwise_causal_calibration_ece'))} |",
         f"| eval | {_fmt(summary.get('coverage'))} | {_fmt(summary.get('fallback_rate'))} | "
         f"{_fmt(summary.get('base_success'))} | {_fmt(summary.get('selected_success'))} | "
         f"{_fmt(summary.get('proposal_oracle_success'))} | {_fmt(summary.get('outcome_ptr'))} | "
-        f"{_fmt(summary.get('success_support_gap'))} | {_fmt(summary.get('success_selector_gap'))} |",
+        f"{_fmt(summary.get('success_support_gap'))} | {_fmt(summary.get('success_selector_gap'))} | "
+        f"{_fmt(summary.get('pairwise_causal_calibration_ece'))} |",
         "",
         "This is a selector diagnostic over already measured candidates, not a new rollout.",
     ]