#!/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 candidate_diversity, collapse_rate, macro_micro_summary, mean_nearest_distance_to_set, 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) parser.add_argument( "--no-markdown-report", action="store_true", help="Do not write report.md; the persistent prose summary lives in README.md.", ) 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"]) chart_feature_mode = str(checkpoint.get("chart_feature_mode", "base")) 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, chart_feature_mode=chart_feature_mode, ) target_charts, target_index = load_charts( args.target_index, max_charts=args.max_target_charts, chart_feature_mode=chart_feature_mode, ) _validate_indexes(args.source_index, source_index, args.target_index, target_index) 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}", f"chart_feature_mode={chart_feature_mode}", ] 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": target.seed, "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 positive_distance = mean_nearest_distance_to_set(proposals, positives, k=args.k) row[f"mean_positive_distance_at_{args.k}"] = positive_distance negative_distance = mean_nearest_distance_to_set(proposals, negatives, k=args.k) row[f"mean_negative_distance_at_{args.k}"] = negative_distance closer = positives_closer_than_negatives(proposals, positives, negatives, k=args.k) row[f"pos_closer_than_neg_at_{args.k}"] = closer row[f"candidate_diversity_at_{args.k}"] = candidate_diversity(proposals, k=args.k) row[f"collapse_rate_at_{args.k}"] = collapse_rate(proposals, k=args.k) 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"), "chart_feature_mode": chart_feature_mode, } (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(json.dumps(_by_group(rows, metric_names, "seed"), indent=2, sort_keys=True) + "\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") _write_report_artifact(out_dir, summary, k=args.k, no_markdown_report=args.no_markdown_report) print(json.dumps({"out_dir": str(out_dir), "num_rows": len(rows)}, indent=2)) return 0 def _validate_indexes( source_path: Path, source_index: dict[str, Any], target_path: Path, target_index: dict[str, Any], ) -> None: if source_index.get("split") != "train" or not source_index.get("retrieval_index_allowed"): raise SystemExit( f"{source_path} is not a train-only retrieval index; CTT proxy eval must " "retrieve source positives from train split only" ) if not source_index.get("include_outcomes"): raise SystemExit(f"{source_path} must include train outcomes for source positives") if not target_index.get("include_outcomes"): raise SystemExit( f"{target_path} does not expose evaluator outcomes/labels. " "Proxy support evaluation needs an evaluator-only target chart DB; " "do not substitute hidden labels or distance proxies from train self-target." ) if target_index.get("split") != "train" and target_index.get("retrieval_index_allowed"): raise SystemExit(f"{target_path} is non-train but marked retrieval_index_allowed") 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]]: return _by_group(rows, metric_names, "task_id") def _by_group( rows: list[dict[str, Any]], metric_names: list[str], group_key: str, ) -> dict[str, dict[str, float]]: output: dict[str, dict[str, float]] = {} for row in rows: group = str(row[group_key]) output.setdefault(group, {}) for group in output: group_rows = [row for row in rows if str(row[group_key]) == group] for metric in metric_names: values = [float(row[metric]) for row in group_rows if isinstance(row.get(metric), (int, float))] if values: output[group][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) def _write_report_artifact( out_dir: Path, summary: dict[str, Any], *, k: int, no_markdown_report: bool, ) -> None: report_path = out_dir / "report.md" if no_markdown_report: if report_path.exists(): report_path.unlink() return report_path.write_text(_report(summary, k) + "\n") if __name__ == "__main__": raise SystemExit(main())