make_tables.py

"""Generate LaTeX tables and confusion-matrix heatmaps for the paper from one
or more summary.json files (produced by rescore.py).

Outputs (into --out, default results/analysis/):
    table1_accuracy.tex       -- headline accuracy by category (+ overall), one
                                  Exit-code/stdout column pair per model
    table2_by_outcome.tex     -- accuracy conditioned on the actual sandbox
                                  outcome, one column pair per model
    confusion_matrix.pdf/.png            -- single model: rows=actual, cols=simulated
    confusion_matrix_<label>.pdf/.png    -- multiple models: one heatmap each

The .tex tables use booktabs (\\usepackage{booktabs} required in the preamble).

Usage:
    python make_tables.py --summary results/summary.json
    python make_tables.py --summary results/summary.json results/facebook_cwm/summary.json
    python make_tables.py --summary Qwen=results/summary.json CWM=results/facebook_cwm/summary.json
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

# Severity-ordered where meaningful; infra_error/unknown are not part of the
# severity scale (see comparison/compare.py::_EXIT_BUCKET_SEVERITY) so they
# trail at the end.
BUCKET_ORDER = ["clean_exit", "error_exit", "resource_exhaustion", "infra_error", "unknown"]
BUCKET_LABEL = {
    "clean_exit": "Clean exit",
    "error_exit": "Error exit",
    "resource_exhaustion": "Resource exh.",
    "infra_error": "Infra error",
    "unknown": "Unknown",
}

MIN_RELIABLE_N = 5  # below this, flag the rate as not statistically meaningful

# Categories included in the paper's headline table, in display order. Any
# category not listed here (e.g. sensitive_info_recovery, still an
# in-progress corpus) is dropped from table1 and from its "Overall" row.
CATEGORY_ORDER = [
    "compute",
    "nondeterminism",
    "recent_changes",
    "external_environment",
    "verbosity",
    "misleading",
    "prompt_injections",
]
CATEGORY_LABEL = {
    "compute": "Compute",
    "nondeterminism": "Non-Determinism",
    "recent_changes": "Missing Knowledge",
    "external_environment": "Ext. Environment",
    "verbosity": "Token Count",
    "misleading": "Misleading",
    "prompt_injections": "Injections",
}


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Generate paper tables/figures from one or more summary.json files")
    p.add_argument(
        "--summary", nargs="+", default=["results/summary.json"],
        metavar="[LABEL=]PATH",
        help="One or more summary.json paths. Optionally prefix with an explicit "
             "label (e.g. CWM=results/facebook_cwm/summary.json); otherwise the "
             "label is taken from the summary's own model_name/model field.",
    )
    p.add_argument("--out", default="results/analysis")
    return p.parse_args()


def _load_summaries(specs: list[str]) -> list[tuple[str, dict]]:
    loaded = []
    for spec in specs:
        label, _, path = spec.partition("=")
        if not path:
            label, path = "", spec
        summary_path = Path(path)
        if not summary_path.is_file():
            raise SystemExit(f"no summary at {summary_path}")
        data = json.loads(summary_path.read_text())
        if not label:
            label = data.get("model_name") or data.get("model") or summary_path.parent.name or summary_path.stem
        loaded.append((label, data))
    return loaded


def _ordered(keys) -> list[str]:
    known = [b for b in BUCKET_ORDER if b in keys]
    unknown = sorted(k for k in keys if k not in BUCKET_ORDER)
    return known + unknown


def _escape(s: str) -> str:
    return s.replace("_", r"\_")


def _slug(label: str) -> str:
    return "".join(c if c.isalnum() else "_" for c in label.lower()).strip("_")


# --- Table 1: headline accuracy by category ----------------------------------

def write_table1(summaries: list[tuple[str, dict]], out_dir: Path) -> None:
    n = len(summaries)
    included = [c for c in CATEGORY_ORDER if any(c in s["categories"] for _, s in summaries)]

    lines = [
        r"\begin{table}[t]",
        r"\centering",
        r"\caption{Simulation accuracy by script category.}",
        r"\label{tab:accuracy}",
        r"\begin{tabular}{l" + "cc" * n + "}",
        r"\toprule",
    ]
    if n > 1:
        model_header = " & " + " & ".join(
            rf"\multicolumn{{2}}{{c}}{{{_escape(label)}}}" for label, _ in summaries
        ) + r" \\"
        cmidrules = "".join(rf"\cmidrule(lr){{{2 + 2 * i}-{3 + 2 * i}}}" for i in range(n))
        lines += [model_header, cmidrules]
    metric_header = "Category & " + " & ".join(
        r"Exit-code & \texttt{stdout}" for _ in summaries
    ) + r" \\"
    lines += [metric_header, r"\midrule"]

    def row(label_col: str, cat: str) -> str:
        cells = []
        for _, s in summaries:
            c = s["categories"].get(cat)
            cells.append(f"{c['avg_exit_code_match']:.2f} & {c['avg_stdout_exact_match']:.2f}" if c else "-- & --")
        return f"{label_col} & " + " & ".join(cells) + r" \\"

    for cat in included:
        lines.append(row(CATEGORY_LABEL.get(cat, _escape(cat)), cat))
    lines.append(r"\midrule")

    # Weighted over just the included categories -- summary["overall"] spans
    # every category in the raw results (e.g. sensitive_info_recovery), which
    # this table deliberately excludes.
    def weighted_overall(s: dict, metric: str) -> float:
        cats = [s["categories"][c] for c in included if c in s["categories"]]
        total_n = sum(c["n"] for c in cats)
        return sum(c[metric] * c["n"] for c in cats) / total_n

    overall_cells = " & ".join(
        f"\\textbf{{{weighted_overall(s, 'avg_exit_code_match'):.2f}}} & "
        f"\\textbf{{{weighted_overall(s, 'avg_stdout_exact_match'):.2f}}}"
        for _, s in summaries
    )
    lines.append(f"\\textbf{{Overall}} & {overall_cells} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}", r"\end{table}"]

    out_path = out_dir / "table1_accuracy.tex"
    out_path.write_text("\n".join(lines) + "\n")
    print(f"wrote {out_path}")


# --- Table 2: accuracy conditioned on actual outcome --------------------------

def write_table2(summaries: list[tuple[str, dict]], out_dir: Path) -> None:
    n = len(summaries)
    all_outcomes = _ordered(set().union(*(s["by_actual_outcome"].keys() for _, s in summaries)))
    # All models are compared against the same cached sandbox ground truth, so
    # n is shared across models for a given outcome -- one column, not one per model.
    any_small_n = any(
        s["by_actual_outcome"].get(o, {}).get("n", 0) < MIN_RELIABLE_N
        for o in all_outcomes for _, s in summaries
        if o in s["by_actual_outcome"]
    )

    lines = [
        r"\begin{table}[t]",
        r"\centering",
        r"\caption{Simulation accuracy conditioned on the actual sandbox outcome.}",
        r"\label{tab:by-outcome}",
        r"\begin{tabular}{lr" + "cc" * n + "}",
        r"\toprule",
    ]
    if n > 1:
        model_header = " & & " + " & ".join(
            rf"\multicolumn{{2}}{{c}}{{{_escape(label)}}}" for label, _ in summaries
        ) + r" \\"
        cmidrules = "".join(rf"\cmidrule(lr){{{3 + 2 * i}-{4 + 2 * i}}}" for i in range(n))
        lines += [model_header, cmidrules]
    metric_header = "Outcome & $n$ & " + " & ".join(
        r"Exit-code & \texttt{stdout}" for _ in summaries
    ) + r" \\"
    lines += [metric_header, r"\midrule"]

    for outcome in all_outcomes:
        label = BUCKET_LABEL.get(outcome, _escape(outcome))
        values = [s["by_actual_outcome"].get(outcome) for _, s in summaries]
        ns = {v["n"] for v in values if v is not None}
        n_cell = str(ns.pop()) if len(ns) == 1 else "/".join(str(v["n"]) if v else "--" for v in values)
        cells = []
        small_n = False
        for v in values:
            if v is None:
                cells.append("-- & --")
                continue
            if v["n"] < MIN_RELIABLE_N:
                small_n = True
            cells.append(f"{v['avg_exit_code_match']:.2f} & {v['avg_stdout_exact_match']:.2f}")
        note = r"$^\dagger$" if small_n else ""
        lines.append(f"{label}{note} & {n_cell} & " + " & ".join(cells) + r" \\")
    lines += [r"\bottomrule", r"\end{tabular}"]

    if any_small_n:
        lines += [
            r"\vspace{2pt}",
            (r"{\footnotesize $^\dagger$ Sample size below "
             rf"{MIN_RELIABLE_N} for at least one model; rate shown for completeness, not reliable.}}"),
        ]
    lines.append(r"\end{table}")

    out_path = out_dir / "table2_by_outcome.tex"
    out_path.write_text("\n".join(lines) + "\n")
    print(f"wrote {out_path}")


# --- Confusion matrix heatmap -------------------------------------------------

def write_confusion_heatmap(label: str, summary: dict, out_dir: Path, *, suffix: bool) -> None:
    confusion = summary["confusion_matrix"]
    buckets = _ordered(set(confusion) | {b for row in confusion.values() for b in row})

    matrix = np.array([
        [confusion.get(actual, {}).get(sim, 0) for sim in buckets]
        for actual in buckets
    ])
    labels = [BUCKET_LABEL.get(b, b) for b in buckets]

    fig, ax = plt.subplots(figsize=(1.4 * len(buckets) + 1.5, 1.4 * len(buckets) + 1.0))
    # Sequential, single-hue, light->dark: magnitude encoding, colorblind-safe.
    im = ax.imshow(matrix, cmap="Blues", aspect="equal")

    ax.set_xticks(range(len(buckets)))
    ax.set_yticks(range(len(buckets)))
    ax.set_xticklabels(labels, rotation=35, ha="right")
    ax.set_yticklabels(labels)
    ax.set_xlabel("Simulated outcome")
    ax.set_ylabel("Actual outcome")
    ax.set_title(f"Simulated vs. actual outcome (script count) -- {label}" if suffix
                 else "Simulated vs. actual outcome (script count)")

    vmax = matrix.max() if matrix.max() > 0 else 1
    for i in range(len(buckets)):
        for j in range(len(buckets)):
            val = matrix[i, j]
            # White text on dark fill, dark text on light fill -- never color alone.
            color = "white" if val > 0.6 * vmax else "black"
            weight = "bold" if i == j else "normal"
            ax.text(j, i, str(val), ha="center", va="center", color=color, fontweight=weight)
            if i == j:
                ax.add_patch(plt.Rectangle((j - 0.5, i - 0.5), 1, 1, fill=False,
                                            edgecolor="black", linewidth=1.2))

    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label="Scripts")
    fig.tight_layout()

    stem = f"confusion_matrix_{_slug(label)}" if suffix else "confusion_matrix"
    for ext in ("pdf", "png"):
        out_path = out_dir / f"{stem}.{ext}"
        fig.savefig(out_path, dpi=200 if ext == "png" else None)
        print(f"wrote {out_path}")
    plt.close(fig)


def main() -> int:
    args = parse_args()
    summaries = _load_summaries(args.summary)
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    write_table1(summaries, out_dir)
    write_table2(summaries, out_dir)
    for label, summary in summaries:
        write_confusion_heatmap(label, summary, out_dir, suffix=len(summaries) > 1)
    return 0


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