scripts/figure_tokens.py

#!/usr/bin/env python
"""Figure: Token cost vs performance (Claim 4, REPORTED).

Reproduces Table 3's token/performance trade-off. Highlights L1+Tool
(12.614M tokens) vs L1 CoT (4.016M tokens) = ~3.14x tokens for equal-or-worse
mean gap. Clearly labelled REPORTED (LLM slice not run here).
"""
from __future__ import annotations

import json
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

ROOT = Path("/home/ubuntu/samuel/dynasched-repro")
SRC = ROOT / "outputs" / "tokens_reported.json"
OUT = ROOT / "figs" / "token_cost.png"


def main():
    j = json.loads(SRC.read_text())
    d = j["table3_reported"]
    ratio = j["token_ratio_tool_over_cot"]
    configs = list(d.keys())
    toks = [d[c]["tokens_M"] for c in configs]

    colors = []
    for c in configs:
        if c == "L1+Tool":
            colors.append("#b91c1c")
        elif c == "L1 CoT":
            colors.append("#0f766e")
        else:
            colors.append("#94a3b8")

    x = np.arange(len(configs))
    fig, ax = plt.subplots(figsize=(9.5, 5.5))
    bars = ax.bar(x, toks, color=colors)
    ax.set_xticks(x)
    ax.set_xticklabels(configs, rotation=30, ha="right", fontsize=9)
    ax.set_ylabel("Total tokens (millions)")
    ax.set_title(
        "Token cost by configuration (Table 3, REPORTED — LLM slice not run)\n"
        f"L1+Tool / L1 CoT = 12.614M / 4.016M = {ratio:.2f}x tokens for equal-or-worse gap"
    )
    for c, bar in zip(configs, bars):
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            bar.get_height() + 0.15,
            f"{d[c]['tokens_M']:.2f}M\n{d[c]['mean_gap_pct']:.1f}%",
            ha="center",
            va="bottom",
            fontsize=8,
        )
    ax.text(
        0.99,
        0.97,
        "REPORTED — not reproduced (no LLM rollouts)",
        transform=ax.transAxes,
        ha="right",
        va="top",
        fontsize=9,
        color="#b91c1c",
        style="italic",
    )
    plt.tight_layout()
    OUT.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(OUT, dpi=130)
    print(f"Wrote {OUT}")


if __name__ == "__main__":
    main()