scripts/build_ctt_rollout_comparison.py

#!/usr/bin/env python
from __future__ import annotations

import argparse
import json
import math
import re
import shutil
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))

from scripts.eval_metrics import main as eval_metrics_main  # noqa: E402


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Aggregate measured CTT rollout runs.")
    parser.add_argument(
        "--run-glob",
        default="runs/ctt_residual_rollout_val69_seed*",
        help="Glob for run directories containing measured_candidates.json.",
    )
    parser.add_argument("--out-dir", type=Path, default=Path("runs/ctt_val_rollout_comparison"))
    parser.add_argument("--k", type=int, default=8)
    parser.add_argument("--bootstrap-samples", type=int, default=1000)
    args = parser.parse_args(argv)

    run_dirs = [
        path
        for path in sorted(Path().glob(args.run_glob))
        if (path / "measured_candidates.json").exists()
    ]
    if not run_dirs:
        raise SystemExit(f"no measured rollout runs found for glob {args.run_glob!r}")

    combined_rows: list[dict[str, Any]] = []
    payloads = []
    for run_dir in run_dirs:
        payload = json.loads((run_dir / "measured_candidates.json").read_text())
        payloads.append(payload)
        train_seed = _seed_from_path(run_dir)
        for row in payload.get("rows", []):
            item = dict(row)
            item["train_seed"] = train_seed
            item["seed"] = f"{row.get('seed', 'unknown')}/train{train_seed}"
            combined_rows.append(item)

    out_dir = args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    combined = {
        "report_type": "ctt_measured_rollout_comparison",
        "schema_version": 1,
        "k": args.k,
        "run_dirs": [str(path) for path in run_dirs],
        "train_seeds": [_seed_from_path(path) for path in run_dirs],
        "num_rows": len(combined_rows),
        "data_hash": _first(payloads, "target_content_hash"),
        "split_hash": _first(payloads, "target_split_hash"),
        "source_content_hash": _first(payloads, "source_content_hash"),
        "target_content_hash": _first(payloads, "target_content_hash"),
        "target_split_hash": _first(payloads, "target_split_hash"),
        "rows": combined_rows,
    }
    combined_path = out_dir / "combined_measured_candidates.json"
    combined_path.write_text(json.dumps(combined, indent=2, sort_keys=True) + "\n")
    eval_metrics_main(
        [
            "--input",
            str(combined_path),
            "--out-dir",
            str(out_dir / "measured_metrics"),
            "--mode",
            "measured",
            "--k",
            str(args.k),
            "--bootstrap-samples",
            str(args.bootstrap_samples),
        ]
    )
    metrics = json.loads((out_dir / "measured_metrics" / "metrics.json").read_text())
    (out_dir / "metrics.json").write_text(json.dumps(_summary_payload(combined, metrics), indent=2, sort_keys=True) + "\n")
    (out_dir / "table.tex").write_text((out_dir / "measured_metrics" / "table.tex").read_text())
    (out_dir / "report.md").write_text(_report(combined, metrics) + "\n")
    for filename in ("metrics_by_task.json", "metrics_by_seed.json"):
        shutil.copyfile(out_dir / "measured_metrics" / filename, out_dir / filename)
    (out_dir / "command.txt").write_text(
        "python scripts/build_ctt_rollout_comparison.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(combined.get("target_content_hash") or "") + "\n")
    (out_dir / "split_hash.txt").write_text(str(combined.get("target_split_hash") or "") + "\n")
    (out_dir / "source_data_hash.txt").write_text(str(combined.get("source_content_hash") or "") + "\n")
    (out_dir / "train.log").write_text("comparison artifact; source runs trained/evaluated separately\n")
    (out_dir / "eval.log").write_text(
        "\n".join(f"{run_dir}: measured_candidates.json" for run_dir in combined["run_dirs"]) + "\n"
    )
    print(json.dumps({"out_dir": str(out_dir), "runs": len(run_dirs), "rows": len(combined_rows)}, indent=2))
    return 0


def _summary_payload(combined: dict[str, Any], metrics: dict[str, Any]) -> dict[str, Any]:
    return {
        "report_type": "ctt_measured_rollout_comparison",
        "k": combined["k"],
        "run_dirs": combined["run_dirs"],
        "train_seeds": combined["train_seeds"],
        "num_rows": combined["num_rows"],
        "summary": metrics.get("summary", {}),
        "success_summary": _success_summary(combined.get("rows", []), k=int(combined["k"])),
        "source_content_hash": combined.get("source_content_hash"),
        "target_content_hash": combined.get("target_content_hash"),
        "target_split_hash": combined.get("target_split_hash"),
    }


def _report(combined: dict[str, Any], metrics: dict[str, Any]) -> str:
    summary = metrics.get("summary", {})
    split = _split_name(combined)
    lines = [
        f"# CTT {split.title()} Measured Rollout Comparison",
        "",
        f"Runs: `{len(combined['run_dirs'])}`",
        f"Rows: `{combined['num_rows']}`",
        f"K: `{combined['k']}`",
        "",
        "| Metric | N | Micro mean | 95% CI |",
        "| --- | ---: | ---: | ---: |",
    ]
    for name, payload in sorted(summary.items()):
        micro = payload.get("micro", {})
        lines.append(
            f"| {name} | {micro.get('n', 0)} | {_fmt(micro.get('mean'))} | "
            f"[{_fmt(micro.get('low'))}, {_fmt(micro.get('high'))}] |"
        )
    lines.append("")
    success = _success_summary(combined.get("rows", []), k=int(combined["k"]))
    lines.extend(
        [
            "| Success/Utility | Mean |",
            "| --- | ---: |",
            f"| base_success_rate | {_fmt(success.get('base_success_rate'))} |",
            f"| selected_success_rate | {_fmt(success.get('selected_success_rate'))} |",
            f"| proposal_oracle_success_rate | {_fmt(success.get('proposal_oracle_success_rate'))} |",
            f"| hidden_chart_oracle_success_rate | {_fmt(success.get('hidden_chart_oracle_success_rate'))} |",
            f"| selected_success_gain_over_base | {_fmt(success.get('selected_success_gain_over_base'))} |",
            f"| proposal_oracle_success_gain_over_base | {_fmt(success.get('proposal_oracle_success_gain_over_base'))} |",
            f"| success_support_gap | {_fmt(success.get('success_support_gap'))} |",
            f"| success_selector_gap | {_fmt(success.get('success_selector_gap'))} |",
            f"| base_utility_mean | {_fmt(success.get('base_utility_mean'))} |",
            f"| selected_utility_mean | {_fmt(success.get('selected_utility_mean'))} |",
            f"| proposal_oracle_utility_mean | {_fmt(success.get('proposal_oracle_utility_mean'))} |",
            f"| hidden_chart_oracle_utility_mean | {_fmt(success.get('hidden_chart_oracle_utility_mean'))} |",
            "",
        ]
    )
    lines.append("These are measured generated-candidate rollouts, not PPTC proxies.")
    lines.append("")
    lines.append("Run dirs:")
    for run_dir in combined["run_dirs"]:
        lines.append(f"- `{run_dir}`")
    return "\n".join(lines)


def _seed_from_path(path: Path) -> str:
    match = re.search(r"seed(\d+)", path.name)
    return match.group(1) if match else path.name


def _split_name(combined: dict[str, Any]) -> str:
    text = " ".join(str(item) for item in combined.get("run_dirs", []))
    if "test" in text:
        return "test"
    if "val" in text or "validation" in text:
        return "validation"
    return "measured"


def _first(payloads: list[dict[str, Any]], key: str) -> Any:
    for payload in payloads:
        value = payload.get(key)
        if value:
            return value
    return None


def _success_summary(rows: list[dict[str, Any]], *, k: int) -> dict[str, Any]:
    base_success = []
    selected_success = []
    oracle_success = []
    base_utility = []
    selected_utility = []
    oracle_utility = []
    hidden_oracle_utility = []
    hidden_oracle_success = []
    success_support_gap = []
    success_selector_gap = []
    selected_success_gain = []
    proposal_oracle_success_gain = []
    restore_errors = []
    for row in rows:
        generated_utilities = [float(value) for value in row.get("generated_utilities", [])[:k]]
        generated_success = [bool(value) for value in row.get("candidate_success", [])[:k]]
        selected_index = int(row.get("selected_index", 0))
        selected_success_value: float | None = None
        proposal_oracle_success_value: float | None = None
        base_success_value: float | None = None
        if "base_success" in row:
            base_success_value = float(bool(row["base_success"]))
            base_success.append(base_success_value)
        if selected_index < len(generated_success):
            selected_success_value = float(generated_success[selected_index])
            selected_success.append(selected_success_value)
        if generated_success:
            proposal_oracle_success_value = float(any(generated_success))
            oracle_success.append(proposal_oracle_success_value)
        if "base_utility" in row:
            base_utility.append(float(row["base_utility"]))
        if selected_index < len(generated_utilities):
            selected_utility.append(generated_utilities[selected_index])
        if generated_utilities:
            oracle_utility.append(max(generated_utilities))
        hidden = [float(value) for value in row.get("hidden_chart_utilities", [])]
        if hidden:
            hidden_oracle_utility.append(max(hidden))
            hidden_success_value = float(any(value >= 1.0 for value in hidden))
            hidden_oracle_success.append(hidden_success_value)
            if proposal_oracle_success_value is not None:
                success_support_gap.append(
                    max(0.0, hidden_success_value - proposal_oracle_success_value)
                )
            if (
                proposal_oracle_success_value is not None
                and selected_success_value is not None
            ):
                success_selector_gap.append(
                    max(0.0, proposal_oracle_success_value - selected_success_value)
                )
        if base_success_value is not None and selected_success_value is not None:
            selected_success_gain.append(selected_success_value - base_success_value)
        if base_success_value is not None and proposal_oracle_success_value is not None:
            proposal_oracle_success_gain.append(
                proposal_oracle_success_value - base_success_value
            )
        if "restore_error" in row:
            restore_errors.append(float(row["restore_error"]))
    return {
        "base_success_rate": _mean(base_success),
        "selected_success_rate": _mean(selected_success),
        "proposal_oracle_success_rate": _mean(oracle_success),
        "hidden_chart_oracle_success_rate": _mean(hidden_oracle_success),
        "selected_success_gain_over_base": _mean(selected_success_gain),
        "proposal_oracle_success_gain_over_base": _mean(proposal_oracle_success_gain),
        "success_support_gap": _mean(success_support_gap),
        "success_selector_gap": _mean(success_selector_gap),
        "base_utility_mean": _mean(base_utility),
        "selected_utility_mean": _mean(selected_utility),
        "proposal_oracle_utility_mean": _mean(oracle_utility),
        "hidden_chart_oracle_utility_mean": _mean(hidden_oracle_utility),
        "max_restore_error": max(restore_errors) if restore_errors else None,
    }


def _mean(values: list[float]) -> float | None:
    clean = [float(value) for value in values if math.isfinite(float(value))]
    return sum(clean) / len(clean) if clean else None


def _fmt(value: Any) -> str:
    if not isinstance(value, (int, float)) or not math.isfinite(float(value)):
        return "n/a"
    return f"{float(value):.4f}"


def _run(command: list[str]) -> str:
    try:
        return subprocess.check_output(command, cwd=PROJECT_ROOT, text=True).strip()
    except (subprocess.CalledProcessError, FileNotFoundError):
        return ""


if __name__ == "__main__":
    raise SystemExit(main())