scripts/eval_ctt_proxy.py

#!/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())