ctt artifacts 2026-07-02 workspace/scripts/eval_ctt_proxy.py

workspace/scripts/eval_ctt_proxy.py

@@ -0,0 +1,223 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+import torch  # noqa: E402
+
+from cil.metrics import (  # noqa: E402
+    macro_micro_summary,
+    negative_near_at_threshold,
+    positives_closer_than_negatives,
+    proxy_positive_tangent_coverage_at_k,
+    proxy_support_distance,
+)
+from cil.models import CTTConfig, CausalTangentTransport, ChartEncoder, TangentNormalizer  # noqa: E402
+from scripts.train_ctt import load_charts  # noqa: E402
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Evaluate CTT support geometry with proxy metrics.")
+    parser.add_argument("--checkpoint", type=Path, required=True)
+    parser.add_argument("--source-index", type=Path, default=Path("data/cil_charts/train/index.json"))
+    parser.add_argument("--target-index", type=Path, default=Path("data/cil_charts/train/index.json"))
+    parser.add_argument("--out-dir", type=Path, default=Path("runs/ctt_residual_smoke_proxy"))
+    parser.add_argument("--k", type=int, default=16)
+    parser.add_argument("--thresholds", default="0.20,0.40")
+    parser.add_argument("--max-target-charts", type=int, default=64)
+    parser.add_argument("--neighbors", type=int, default=8)
+    args = parser.parse_args(argv)
+
+    thresholds = [float(item) for item in args.thresholds.split(",") if item.strip()]
+    checkpoint = torch.load(args.checkpoint, map_location="cpu")
+    config = CTTConfig(**checkpoint["config"])
+    encoder = ChartEncoder(config.chart_feature_dim, output_dim=config.chart_dim)
+    ctt = CausalTangentTransport(config)
+    encoder.load_state_dict(checkpoint["chart_encoder"])
+    ctt.load_state_dict(checkpoint["ctt"])
+    encoder.eval()
+    ctt.eval()
+    normalizer = TangentNormalizer.from_dict(checkpoint["normalizer"])
+
+    source_charts, source_index = load_charts(args.source_index, max_charts=None)
+    target_charts, target_index = load_charts(args.target_index, max_charts=args.max_target_charts)
+    rows = []
+    log_lines = [
+        f"source_charts={len(source_charts)} target_charts={len(target_charts)} k={args.k}",
+        f"source_index={args.source_index}",
+        f"target_index={args.target_index}",
+    ]
+    source_by_task: dict[str, list[Any]] = {}
+    for chart in source_charts:
+        source_by_task.setdefault(chart.task_id, []).append(chart)
+
+    with torch.no_grad():
+        for target in target_charts:
+            pool = source_by_task.get(target.task_id, source_charts)
+            neighbors = sorted(
+                pool,
+                key=lambda source: torch.linalg.vector_norm(source.feature - target.feature).item(),
+            )[: args.neighbors]
+            proposals = []
+            z_target = encoder(target.feature.unsqueeze(0))
+            for source in neighbors:
+                z_source = encoder(source.feature.unsqueeze(0))
+                for xi_source in source.positives[: max(1, args.k // max(1, len(neighbors)) + 1)]:
+                    if len(proposals) >= args.k:
+                        break
+                    xi_norm = normalizer.transform(xi_source.unsqueeze(0))
+                    xi_hat_norm = ctt(z_source, z_target, xi_norm)
+                    proposals.append(normalizer.inverse_transform(xi_hat_norm).squeeze(0).cpu().tolist())
+                if len(proposals) >= args.k:
+                    break
+            positives = target.positives.cpu().tolist()
+            negatives = target.negatives.cpu().tolist()
+            row: dict[str, Any] = {
+                "chart_id": target.chart_id,
+                "task_id": target.task_id,
+                "seed": "unknown",
+                "num_proposals": len(proposals),
+            }
+            for threshold in thresholds:
+                suffix = f"{threshold:.2f}".replace(".", "p")
+                row[f"pptc_at_{args.k}_thr_{suffix}"] = proxy_positive_tangent_coverage_at_k(
+                    proposals,
+                    positives,
+                    threshold=threshold,
+                    k=args.k,
+                )
+                row[f"negative_near_at_{args.k}_thr_{suffix}"] = negative_near_at_threshold(
+                    proposals,
+                    negatives,
+                    threshold=threshold,
+                    k=args.k,
+                )
+            distance = proxy_support_distance(proposals, positives, k=args.k)
+            row[f"proxy_support_distance_at_{args.k}"] = distance
+            closer = positives_closer_than_negatives(proposals, positives, negatives, k=args.k)
+            row[f"pos_closer_than_neg_at_{args.k}"] = closer
+            rows.append(row)
+
+    metric_names = sorted(
+        {
+            key
+            for row in rows
+            for key, value in row.items()
+            if isinstance(value, (int, float)) and math.isfinite(float(value))
+        }
+        - {"num_proposals"}
+    )
+    summary = {name: macro_micro_summary(rows, name, bootstrap_samples=200) for name in metric_names}
+    out_dir = args.out_dir
+    out_dir.mkdir(parents=True, exist_ok=True)
+    _write_run_provenance(out_dir, args, source_index, target_index)
+    metrics = {
+        "report_type": "ctt_proxy_eval",
+        "k": args.k,
+        "thresholds": thresholds,
+        "num_rows": len(rows),
+        "rows": rows,
+        "summary": summary,
+        "data_hash": source_index.get("content_hash"),
+        "split_hash": source_index.get("split_hash"),
+        "target_split_hash": target_index.get("split_hash"),
+    }
+    (out_dir / "metrics.json").write_text(json.dumps(metrics, indent=2, sort_keys=True) + "\n")
+    (out_dir / "metrics_by_task.json").write_text(json.dumps(_by_task(rows, metric_names), indent=2, sort_keys=True) + "\n")
+    (out_dir / "metrics_by_seed.json").write_text("{}\n")
+    (out_dir / "eval.log").write_text("\n".join(log_lines) + "\n")
+    (out_dir / "train.log").write_text("see checkpoint run\n")
+    (out_dir / "table.tex").write_text(_table(summary) + "\n")
+    (out_dir / "report.md").write_text(_report(summary, args.k) + "\n")
+    print(json.dumps({"out_dir": str(out_dir), "num_rows": len(rows)}, indent=2))
+    return 0
+
+
+def _write_run_provenance(
+    out_dir: Path,
+    args: argparse.Namespace,
+    source_index: dict[str, Any],
+    target_index: dict[str, Any],
+) -> None:
+    (out_dir / "config.yaml").write_text("\n".join(f"{k}: {v}" for k, v in sorted(vars(args).items())) + "\n")
+    (out_dir / "command.txt").write_text("python scripts/eval_ctt_proxy.py " + " ".join(sys.argv[1:]) + "\n")
+    (out_dir / "git_hash.txt").write_text(_run(["git", "rev-parse", "HEAD"]) + "\n")
+    (out_dir / "data_hash.txt").write_text(str(source_index.get("content_hash", "")) + "\n")
+    (out_dir / "split_hash.txt").write_text(str(target_index.get("split_hash", "")) + "\n")
+
+
+def _run(command: list[str]) -> str:
+    try:
+        return subprocess.check_output(command, cwd=PROJECT_ROOT, text=True).strip()
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        return ""
+
+
+def _by_task(rows: list[dict[str, Any]], metric_names: list[str]) -> dict[str, dict[str, float]]:
+    output: dict[str, dict[str, float]] = {}
+    for row in rows:
+        task = str(row["task_id"])
+        output.setdefault(task, {})
+    for task in output:
+        task_rows = [row for row in rows if row["task_id"] == task]
+        for metric in metric_names:
+            values = [float(row[metric]) for row in task_rows if isinstance(row.get(metric), (int, float))]
+            if values:
+                output[task][metric] = sum(values) / len(values)
+    return output
+
+
+def _table(summary: dict[str, Any]) -> str:
+    lines = [
+        "% Auto-generated by scripts/eval_ctt_proxy.py",
+        "\\begin{tabular}{lrrr}",
+        "\\toprule",
+        "Metric & N & Mean & CI high \\\\",
+        "\\midrule",
+    ]
+    for name, payload in sorted(summary.items()):
+        micro = payload["micro"]
+        lines.append(
+            f"{_latex_escape(name)} & {micro['n']} & {micro['mean']:.4f} & "
+            f"{micro['high']:.4f} \\\\"
+        )
+    lines.extend(["\\bottomrule", "\\end{tabular}"])
+    return "\n".join(lines)
+
+
+def _latex_escape(value: str) -> str:
+    return value.replace("_", "\\_").replace("%", "\\%").replace("&", "\\&")
+
+
+def _report(summary: dict[str, Any], k: int) -> str:
+    lines = [
+        "# CTT Proxy Evaluation",
+        "",
+        f"K: `{k}`",
+        "",
+        "| Metric | N | Mean | 95% CI |",
+        "| --- | ---: | ---: | ---: |",
+    ]
+    for name, payload in sorted(summary.items()):
+        micro = payload["micro"]
+        lines.append(
+            f"| {name} | {micro['n']} | {micro['mean']:.4f} | "
+            f"[{micro['low']:.4f}, {micro['high']:.4f}] |"
+        )
+    lines.append("")
+    lines.append("This is PPTC/proxy support geometry, not OutcomePTR or rollout success.")
+    return "\n".join(lines)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())