bundle/repro_agentselect/build_figures.py

#!/usr/bin/env python3
"""Render only data-backed figures for the local AgentSelect review poster."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--audit", type=Path, required=True)
    parser.add_argument("--claim4", type=Path, required=True)
    parser.add_argument("--out-dir", type=Path, required=True)
    args = parser.parse_args()
    audit = json.loads(args.audit.read_text())
    claim4 = json.loads(args.claim4.read_text())
    args.out_dir.mkdir(parents=True, exist_ok=True)

    claimed = np.array([111179, 107721, 251103])
    observed = np.array([
        audit["current_release_totals_with_paper_claim_counting_convention"]["queries"],
        audit["current_release_totals_with_paper_claim_counting_convention"]["agents"],
        audit["current_release_totals_with_paper_claim_counting_convention"]["positive_pairs"],
    ])
    names = ["Queries", "Agents", "Positive pairs"]
    x = np.arange(len(names))
    fig, ax = plt.subplots(figsize=(11, 5.5), layout="constrained")
    width = 0.37
    ax.bar(x - width / 2, claimed, width, label="Official Claim 2", color="#516A8B")
    ax.bar(x + width / 2, observed, width, label="Current linked release", color="#C66A4A")
    ax.set_ylabel("Count")
    ax.set_xticks(x, names)
    ax.ticklabel_format(style="plain", axis="y")
    ax.legend(frameon=False, loc="upper left")
    ax.set_title("Current official release is smaller than the stated benchmark totals")
    for index, value in enumerate(claimed):
        ax.text(index - width / 2, value + 3000, f"{value:,}", ha="center", va="bottom", fontsize=9)
    for index, value in enumerate(observed):
        ax.text(index + width / 2, value + 3000, f"{value:,}", ha="center", va="bottom", fontsize=9)
    fig.savefig(args.out_dir / "claim2_release_counts.png", dpi=220)
    plt.close(fig)

    parts = ["PartII", "PartIII"]
    content = [claim4["parts"][part]["mean_content_metrics"]["ndcg_at_10"] for part in parts]
    baseline = [claim4["parts"][part]["mean_id_metrics"]["ndcg_at_10"] for part in parts]
    fig, (ax, diff_ax) = plt.subplots(1, 2, figsize=(12, 5.3), gridspec_kw={"width_ratios": [1.15, 1]}, layout="constrained")
    x = np.arange(len(parts))
    ax.bar(x - width / 2, baseline, width, label="ID-KNN + popularity", color="#8E9AAF")
    ax.bar(x + width / 2, content, width, label="Content TwoTowerTFIDF", color="#2F7F7A")
    ax.set_ylabel("Mean nDCG@10")
    ax.set_xticks(x, ["Part II", "Part III"])
    ax.set_ylim(0, max(content + baseline) * 1.25)
    ax.legend(frameon=False, fontsize=8)
    ax.set_title("Full-release GPU evaluation")
    for part_index, part in enumerate(parts):
        values = claim4["parts"][part]["nDCG_at_10_content_minus_ID_by_seed"]
        seeds = list(values)
        means = np.array([values[seed]["mean"] for seed in seeds])
        low = np.array([values[seed]["ci95_low"] for seed in seeds])
        high = np.array([values[seed]["ci95_high"] for seed in seeds])
        offset = (part_index - 0.5) * 0.10
        diff_ax.errorbar(np.arange(len(seeds)) + offset, means, yerr=[means - low, high - means], fmt="o", capsize=3, label=f"Part {'II' if part == 'PartII' else 'III'}")
    diff_ax.axhline(0, color="#333333", linewidth=1)
    diff_ax.set_xticks(np.arange(3), ["101", "202", "303"])
    diff_ax.set_xlabel("Preregistered seed")
    diff_ax.set_ylabel("Content minus ID nDCG@10 (95% paired bootstrap CI)")
    diff_ax.set_title("Matched effect by seed")
    diff_ax.legend(frameon=False, fontsize=8)
    fig.savefig(args.out_dir / "claim4_content_vs_id.png", dpi=220)
    plt.close(fig)

    labels = ["Claim 1", "Claim 2", "Claim 3", "Claim 4", "Claim 5"]
    verdicts = [
        "Partial\n(schema only)",
        "Mismatch\n(current release)",
        "Verified\n(current release)",
        "Reproduced\nin scope",
        "Untested\n(no checkpoint)",
    ]
    colors = ["#C9A24A", "#C66A4A", "#2F7F7A", "#2F7F7A", "#8E9AAF"]
    fig, ax = plt.subplots(figsize=(12, 3.1), layout="constrained")
    for index, (label, verdict, color) in enumerate(zip(labels, verdicts, colors, strict=True)):
        ax.barh(index, 1, color=color, height=0.72)
        ax.text(0.03, index, label, va="center", ha="left", color="white", weight="bold", fontsize=11)
        ax.text(0.97, index, verdict, va="center", ha="right", color="white", weight="bold", fontsize=10)
    ax.set_xlim(0, 1)
    ax.set_ylim(len(labels) - 0.45, -0.55)
    ax.axis("off")
    ax.set_title("Outcome by official claim", loc="left", weight="bold")
    fig.savefig(args.out_dir / "claim_verdicts.png", dpi=220)
    plt.close(fig)


if __name__ == "__main__":
    main()