scripts/generate_paper_assets.py

#!/usr/bin/env python3
"""
Derived from Andrej Karpathy's nanochat project.

MIT License

Copyright (c) 2025 Andrej Karpathy

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
"""

from __future__ import annotations

import csv
import json
import math
from pathlib import Path
import random
import statistics


ROOT = Path(__file__).resolve().parents[1]
PAPER_DIR = ROOT / "paper"
FIG_DIR = PAPER_DIR / "figures"
TABLE_DIR = PAPER_DIR / "tables"


REGIMES = {
    "openweb": {
        "label": "OpenWebText10K",
        "short": "OWT10K",
        "coeff": ROOT
        / "runs/coefficient_calibration/cross_regime_backtest/openwebtext10k_main_interaction/coefficients.json",
        "cells": ROOT
        / "runs/coefficient_calibration/cross_regime_backtest/openwebtext10k_main_interaction/calibration_cells.csv",
        "curve_run": ROOT / "runs/screen_static/20260525-133008",
        "curve_model": "L12_H8_D320",
        "curve_prefix": 1_000_000,
        "curve_title": "OpenWebText10K, L12, 1M",
        "small_stream": ROOT
        / "runs/openwebtext10k_l16_updated_formula_clean_5seed/locked_stream/20260530-174525",
        "small_condition": "openwebtext10k_interaction",
        "small_model": "L16_H8_D384",
        "small_protocol": "4M final prefix / 5 seeds",
        "l20_stream": ROOT
        / "runs/openwebtext10k_l20_8m_extrapolation_3seed/locked_stream/20260601-191150",
        "l20_condition": "openwebtext10k_l20_interaction_8m",
    },
    "tinystories": {
        "label": "TinyStories",
        "short": "TinyStories",
        "coeff": ROOT
        / "runs/coefficient_calibration/tinystories_combined_plus_all_holdouts_interaction/coefficients.json",
        "cells": ROOT
        / "runs/coefficient_calibration/tinystories_combined_plus_all_holdouts_interaction/calibration_cells.csv",
        "curve_run": ROOT
        / "runs/regime_calibration_tinystories_coarse/screen_static/20260529-161726",
        "curve_model": "L12_H8_D320",
        "curve_prefix": 1_000_000,
        "curve_title": "TinyStories, L12, 1M",
        "small_stream": [
            ROOT
            / "runs/streaming_tinystories_interaction_schedule_l12/locked_stream/20260530-053831",
            ROOT
            / "runs/streaming_tinystories_multiseed_validation_l12/locked_stream/20260530-111523",
            ROOT
            / "runs/streaming_tinystories_multiseed_validation_l12/locked_stream/20260530-141335",
        ],
        "small_condition": "interaction",
        "small_model": "L12_H8_D320",
        "small_protocol": "4M final prefix / 5 seeds",
        "l20_stream": ROOT
        / "runs/tinystories_l20_8m_extrapolation_3seed/locked_stream/20260602-112249",
        "l20_condition": "tinystories_l20_interaction_8m",
    },
    "wikitext": {
        "label": "WikiText-103",
        "short": "WikiText",
        "coeff": ROOT / "runs/coefficient_calibration/wikitext103_interaction/coefficients.json",
        "cells": ROOT / "runs/coefficient_calibration/wikitext103_interaction/calibration_cells.csv",
        "curve_run": ROOT
        / "runs/wikitext103_interaction_static_screen/screen_static/20260603-153314",
        "curve_model": "L12_H8_D320",
        "curve_prefix": 250_000,
        "curve_title": "WikiText-103, L12, 250k",
        "small_stream": ROOT
        / "runs/wikitext103_l12_streaming_validation_5seed/locked_stream/20260531-093525",
        "small_condition": "wikitext103_formula_l12",
        "small_model": "L12_H8_D320",
        "small_protocol": "4M final prefix / 5 seeds",
        "l20_stream": ROOT
        / "runs/wikitext103_l20_8m_extrapolation_3seed/locked_stream/20260603-231952",
        "l20_condition": "wikitext103_l20_interaction_8m",
    },
}


def read_json(path: Path) -> dict | list:
    return json.loads(path.read_text(encoding="utf-8"))


def read_jsonl(path: Path) -> list[dict]:
    rows = []
    for line in path.read_text(encoding="utf-8").splitlines():
        if line.strip():
            rows.append(json.loads(line))
    return rows


def read_csv(path: Path) -> list[dict]:
    with path.open(newline="", encoding="utf-8") as handle:
        return list(csv.DictReader(handle))


def mean_std(values: list[float]) -> tuple[float, float]:
    if not values:
        return float("nan"), float("nan")
    mean = statistics.fmean(values)
    std = statistics.stdev(values) if len(values) > 1 else 0.0
    return mean, std


def bootstrap_ci(
    values: list[float], *, confidence: float = 0.95, iterations: int = 20_000, seed: int = 1337
) -> tuple[float, float]:
    if not values:
        return float("nan"), float("nan")
    if len(values) == 1:
        return values[0], values[0]
    rng = random.Random(seed)
    n = len(values)
    means = []
    for _ in range(iterations):
        sample = [values[rng.randrange(n)] for _ in range(n)]
        means.append(statistics.fmean(sample))
    means.sort()
    lower_q = (1.0 - confidence) / 2.0
    upper_q = 1.0 - lower_q
    lower_idx = min(iterations - 1, max(0, int(lower_q * iterations)))
    upper_idx = min(iterations - 1, max(0, int(upper_q * iterations)))
    return means[lower_idx], means[upper_idx]


def fmt(x: float, digits: int = 4) -> str:
    return f"{x:.{digits}f}"


def tex_escape(value: object) -> str:
    text = str(value)
    return (
        text.replace("\\", "\\textbackslash{}")
        .replace("_", "\\_")
        .replace("&", "\\&")
        .replace("%", "\\%")
        .replace("#", "\\#")
    )


def model_label(model_name: str) -> str:
    return tex_escape(model_name.replace("_", "-"))


def token_label(value: int) -> str:
    if value >= 1_000_000:
        whole = value / 1_000_000
        return f"{whole:g}M"
    return f"{value // 1000}k"


def protocol_budget(condition_rows: list[dict], stages: list[int]) -> dict:
    rows_by_seed: dict[int, list[dict]] = {}
    for row in condition_rows:
        rows_by_seed.setdefault(int(row["seed"]), []).append(row)
    seed = sorted(rows_by_seed)[0]
    seed_rows = sorted(rows_by_seed[seed], key=lambda item: int(item["stage"]))
    sampled_deltas = []
    previous_tokens_seen = 0
    for row in seed_rows:
        tokens_seen = int(row["tokens_seen"])
        sampled_deltas.append(tokens_seen - previous_tokens_seen)
        previous_tokens_seen = tokens_seen
    first = seed_rows[0]
    steps_per_stage = sorted({int(row["steps"]) for row in seed_rows})
    block_size = int(first["model_config"]["block_size"])
    sampled_per_stage = sorted(set(sampled_deltas))
    batch_sizes = sorted(
        {
            round(delta / (int(row["steps"]) * block_size))
            for delta, row in zip(sampled_deltas, seed_rows, strict=True)
        }
    )
    return {
        "steps_per_stage": steps_per_stage,
        "batch_sizes": batch_sizes,
        "block_size": block_size,
        "sampled_per_stage": sampled_per_stage,
        "final_sampled_tokens": int(seed_rows[-1]["tokens_seen"]),
        "prefix_path": [int(row["token_limit"]) for row in seed_rows],
        "stage_count": len(stages),
    }


def metric_stage_summary(run_dir: Path | list[Path], condition_name: str) -> dict:
    run_dirs = run_dir if isinstance(run_dir, list) else [run_dir]
    rows = []
    for path in run_dirs:
        rows.extend(read_jsonl(path / "metrics.jsonl"))
    stages = sorted({int(row["stage"]) for row in rows if row.get("stage") is not None})
    final_stage = max(stages)
    by_condition_stage: dict[tuple[str, int], list[dict]] = {}
    for row in rows:
        if row.get("stage") is None:
            continue
        key = (row["condition"], int(row["stage"]))
        by_condition_stage.setdefault(key, []).append(row)

    condition_rows_by_stage = [
        row
        for (condition, _stage), group in by_condition_stage.items()
        if condition == condition_name
        for row in group
    ]
    condition_rows = by_condition_stage[(condition_name, final_stage)]
    final_vals = [float(row["val_eval_loss"]) for row in condition_rows]
    condition_final_mean, condition_final_std = mean_std(final_vals)

    final_static_by_condition: dict[str, list[dict]] = {}
    for row in rows:
        if int(row.get("stage", -1)) != final_stage:
            continue
        if row["condition_kind"] == "static":
            final_static_by_condition.setdefault(row["condition"], []).append(row)
    static_means = {
        condition: statistics.fmean(float(row["val_eval_loss"]) for row in group)
        for condition, group in final_static_by_condition.items()
    }
    best_static_condition = min(static_means, key=static_means.get)
    best_static_rows = final_static_by_condition[best_static_condition]
    best_static_vals = [float(row["val_eval_loss"]) for row in best_static_rows]
    best_static_mean, best_static_std = mean_std(best_static_vals)

    final_static_by_seed: dict[int, tuple[str, float]] = {}
    for row in rows:
        if int(row.get("stage", -1)) != final_stage or row["condition_kind"] != "static":
            continue
        seed = int(row["seed"])
        value = float(row["val_eval_loss"])
        if seed not in final_static_by_seed or value < final_static_by_seed[seed][1]:
            final_static_by_seed[seed] = (row["condition"], value)

    paired_deltas = []
    paired_records = []
    for row in sorted(condition_rows, key=lambda item: int(item["seed"])):
        seed = int(row["seed"])
        decay_loss = float(row["val_eval_loss"])
        static_condition, static_loss = final_static_by_seed[seed]
        delta = decay_loss - static_loss
        paired_deltas.append(delta)
        paired_records.append(
            {
                "seed": seed,
                "decay_loss": decay_loss,
                "best_static_condition": static_condition,
                "best_static_loss": static_loss,
                "delta": delta,
                "gain": -delta,
            }
        )
    paired_mean, paired_std = mean_std(paired_deltas)
    paired_gains = [-delta for delta in paired_deltas]
    paired_gain_mean, paired_gain_std = mean_std(paired_gains)
    paired_gain_ci_low, paired_gain_ci_high = bootstrap_ci(paired_gains)

    stage_deltas = []
    for stage in stages:
        condition_group = by_condition_stage[(condition_name, stage)]
        condition_mean = statistics.fmean(
            float(row["val_eval_loss"]) for row in condition_group
        )
        static_candidates = []
        for (condition, candidate_stage), group in by_condition_stage.items():
            if candidate_stage != stage:
                continue
            if group[0]["condition_kind"] != "static":
                continue
            static_candidates.append(
                (
                    condition,
                    statistics.fmean(float(row["val_eval_loss"]) for row in group),
                    statistics.stdev(float(row["val_eval_loss"]) for row in group)
                    if len(group) > 1
                    else 0.0,
                )
            )
        best_static_stage = min(static_candidates, key=lambda item: item[1])
        first = condition_group[0]
        stage_deltas.append(
            {
                "stage": stage,
                "prefix": int(first["token_limit"]),
                "condition_dropout": statistics.fmean(
                    float(row["dropout_active_final"]) for row in condition_group
                ),
                "condition_mean": condition_mean,
                "condition_std": statistics.stdev(
                    float(row["val_eval_loss"]) for row in condition_group
                )
                if len(condition_group) > 1
                else 0.0,
                "best_static_condition": best_static_stage[0],
                "best_static_mean": best_static_stage[1],
                "best_static_std": best_static_stage[2],
                "delta": condition_mean - best_static_stage[1],
            }
        )

    return {
        "n": len(condition_rows),
        "final_stage": final_stage,
        "final_prefix": int(condition_rows[0]["token_limit"]),
        "condition_final_mean": condition_final_mean,
        "condition_final_std": condition_final_std,
        "best_static_condition": best_static_condition,
        "best_static_mean": best_static_mean,
        "best_static_std": best_static_std,
        "improvement": best_static_mean - condition_final_mean,
        "paired_deltas": paired_deltas,
        "paired_records": paired_records,
        "paired_wins": sum(1 for delta in paired_deltas if delta < 0),
        "paired_mean_delta": paired_mean,
        "paired_std_delta": paired_std,
        "paired_gain_mean": paired_gain_mean,
        "paired_gain_std": paired_gain_std,
        "paired_gain_ci_low": paired_gain_ci_low,
        "paired_gain_ci_high": paired_gain_ci_high,
        "budget": protocol_budget(condition_rows_by_stage, stages),
        "stage_deltas": stage_deltas,
    }


def static_curve_points(run_dir: Path, model: str, prefix: int) -> list[tuple[float, float]]:
    rows = read_json(run_dir / "summary.json")
    curve = []
    for row in rows:
        if row["run_mode"] != "screen_static":
            continue
        if row["condition_kind"] != "static":
            continue
        if row["model_name"] != model:
            continue
        if int(row["token_limit"]) != prefix:
            continue
        curve.append((float(row["dropout_initial"]), float(row["mean_val_eval_loss"])))
    return sorted(curve)


def calibration_points(cells_path: Path) -> dict[str, list[tuple[int, float]]]:
    rows = read_csv(cells_path)
    grouped: dict[str, dict[int, float]] = {}
    for row in rows:
        sampled = int(float(row["sampled_tokens"]))
        unique = int(float(row["unique_tokens"]))
        source = row["source"]
        if "sample_pressure_low" in source:
            continue
        if "heldout_model" in source:
            continue
        model = row["model_name"]
        target = max(0.0, float(row["target_dropout"]))
        grouped.setdefault(model, {})[unique] = target
    return {
        model: sorted(points.items())
        for model, points in sorted(grouped.items())
        if len(points) >= 2
    }


def coeff_summary(regime_key: str, paths: dict) -> dict:
    coeff = read_json(paths["coeff"])
    return {
        "regime": paths["label"],
        "cells": int(coeff["metrics"]["n"]),
        "A": float(coeff["coefficients"]["A"]),
        "B": float(coeff["coefficients"]["B"]),
        "D": float(coeff["coefficients"]["D"]),
        "C0": float(coeff["coefficients"]["C0"]),
        "mae": float(coeff["metrics"]["mae"]),
        "rmse": float(coeff["metrics"]["rmse"]),
        "leave_model_mae": float(coeff["cv"]["leave_model"]["mae"]),
        "leave_prefix_mae": float(coeff["cv"]["leave_prefix"]["mae"]),
        "formula": coeff["formula"],
    }


def tikz_coordinates(points: list[tuple[float, float]], digits: int = 4) -> str:
    return " ".join(f"({x:.4g},{y:.{digits}f})" for x, y in points)


def write_calibration_curve_figure() -> None:
    axes = []
    for idx, (key, paths) in enumerate(REGIMES.items()):
        raw_curve = static_curve_points(
            paths["curve_run"], paths["curve_model"], paths["curve_prefix"]
        )
        minimum = min(loss for _, loss in raw_curve)
        points = [(dropout, loss - minimum) for dropout, loss in raw_curve]
        best = min(raw_curve, key=lambda item: item[1])
        axis_options = [
            f"    title={{{paths['curve_title']}}},",
            "    xlabel={Dropout},",
        ]
        if idx == 0:
            axis_options.append("    ylabel={Validation loss above best},")
        axis_options.extend(
            [
                "    ymin=0,",
                "    grid=both,",
            ]
        )
        axes.append(
            rf"""
\nextgroupplot[
{chr(10).join(axis_options)}
]
\addplot+[thick, mark=*] coordinates {{{tikz_coordinates(points, digits=5)}}};
\addlegendentry{{static sweep}}
\addplot+[only marks, mark=star, mark size=3pt] coordinates {{({best[0]:.4g},0)}};
\addlegendentry{{target $p^\star={best[0]:.2f}$}}
"""
        )
    content = (
        "\\begin{tikzpicture}\n"
        "\\begin{groupplot}[\n"
        "    group style={group size=3 by 1, horizontal sep=1.0cm},\n"
        "    width=0.32\\textwidth,\n"
        "    height=0.26\\textwidth,\n"
        "    legend style={font=\\scriptsize},\n"
        "    tick label style={font=\\scriptsize},\n"
        "    label style={font=\\scriptsize},\n"
        "    title style={font=\\scriptsize},\n"
        "]\n"
        + "\n".join(axes)
        + "\n\\end{groupplot}\n\\end{tikzpicture}\n"
    )
    (FIG_DIR / "calibration_cell_curves.tex").write_text(content, encoding="utf-8")


def write_static_optima_figure() -> None:
    axes = []
    for key, paths in REGIMES.items():
        grouped = calibration_points(paths["cells"])
        labels = sorted({prefix for points in grouped.values() for prefix, _ in points})
        label_to_x = {prefix: idx for idx, prefix in enumerate(labels)}
        plots = []
        for model, points in grouped.items():
            coords = [
                (label_to_x[prefix], target)
                for prefix, target in points
                if prefix in label_to_x
            ]
            plots.append(
                "\\addplot+[thick, mark=*] coordinates {"
                + tikz_coordinates(coords)
                + "};\n"
                + f"\\addlegendentry{{{model_label(model)}}}"
            )
        xticks = ",".join(str(i) for i in range(len(labels)))
        xticklabels = ",".join(token_label(prefix) for prefix in labels)
        axes.append(
            rf"""
\nextgroupplot[
    title={{{paths["label"]}}},
    xlabel={{Unique prefix $U$}},
    ylabel={{Fitted target dropout}},
    xtick={{{xticks}}},
    xticklabels={{{xticklabels}}},
    ymin=0,
    ymax=0.65,
    grid=both,
]
{chr(10).join(plots)}
"""
        )
    content = (
        "\\begin{tikzpicture}\n"
        "\\begin{groupplot}[\n"
        "    group style={group size=3 by 1, horizontal sep=1.0cm},\n"
        "    width=0.32\\textwidth,\n"
        "    height=0.28\\textwidth,\n"
        "    legend style={font=\\tiny},\n"
        "    tick label style={font=\\scriptsize},\n"
        "    label style={font=\\scriptsize},\n"
        "    title style={font=\\scriptsize},\n"
        "]\n"
        + "\n".join(axes)
        + "\n\\end{groupplot}\n\\end{tikzpicture}\n"
    )
    (FIG_DIR / "static_optima_by_regime.tex").write_text(content, encoding="utf-8")


def write_l20_stage_delta_figure(summaries: dict[str, dict]) -> None:
    axes = []
    for key, paths in REGIMES.items():
        stage_rows = summaries[key]["l20"]["stage_deltas"]
        labels = [token_label(row["prefix"]) for row in stage_rows]
        coords = [(idx, row["delta"]) for idx, row in enumerate(stage_rows)]
        zero_coords = [(0, 0.0), (len(stage_rows) - 1, 0.0)]
        xticks = ",".join(str(i) for i in range(len(labels)))
        xticklabels = ",".join(labels)
        axes.append(
            rf"""
\nextgroupplot[
    title={{{paths["label"]}}},
    xlabel={{Prefix}},
    ylabel={{Decay $-$ best static}},
    xtick={{{xticks}}},
    xticklabels={{{xticklabels}}},
    grid=both,
]
\addplot+[thick, mark=*] coordinates {{{tikz_coordinates(coords, digits=5)}}};
\addplot+[black, dashed, mark=none] coordinates {{{tikz_coordinates(zero_coords, digits=5)}}};
"""
        )
    content = (
        "\\begin{tikzpicture}\n"
        "\\begin{groupplot}[\n"
        "    group style={group size=3 by 1, horizontal sep=1.0cm},\n"
        "    width=0.32\\textwidth,\n"
        "    height=0.27\\textwidth,\n"
        "    tick label style={font=\\scriptsize},\n"
        "    label style={font=\\scriptsize},\n"
        "    title style={font=\\scriptsize},\n"
        "]\n"
        + "\n".join(axes)
        + "\n\\end{groupplot}\n\\end{tikzpicture}\n"
    )
    (FIG_DIR / "l20_stage_deltas.tex").write_text(content, encoding="utf-8")


def write_final_improvement_figure(summaries: dict[str, dict]) -> None:
    labels = []
    coords_small = []
    coords_l20 = []
    for idx, (key, paths) in enumerate(REGIMES.items()):
        labels.append(paths["short"])
        coords_small.append((idx, summaries[key]["small"]["improvement"]))
        coords_l20.append((idx, summaries[key]["l20"]["improvement"]))
    xticks = ",".join(str(i) for i in range(len(labels)))
    xticklabels = ",".join(labels)
    content = rf"""
\begin{{tikzpicture}}
\begin{{axis}}[
    ybar,
    bar width=8pt,
    width=0.72\textwidth,
    height=0.34\textwidth,
    ymin=0,
    ylabel={{Best static final loss $-$ decay final loss}},
    symbolic x coords={{{xticklabels}}},
    xtick=data,
    xtick={{{xticklabels}}},
    grid=both,
    legend style={{font=\scriptsize, at={{(0.5,1.04)}}, anchor=south, legend columns=2}},
    tick label style={{font=\scriptsize}},
    label style={{font=\scriptsize}},
]
\addplot coordinates {{{" ".join(f"({label},{y:.5f})" for label, (_, y) in zip(labels, coords_small))}}};
\addlegendentry{{4M validation}}
\addplot coordinates {{{" ".join(f"({label},{y:.5f})" for label, (_, y) in zip(labels, coords_l20))}}};
\addlegendentry{{L20 8M extrapolation}}
\end{{axis}}
\end{{tikzpicture}}
"""
    (FIG_DIR / "final_improvements.tex").write_text(content, encoding="utf-8")


def write_protocol_schematic() -> None:
    content = r"""
\begin{tikzpicture}[font=\scriptsize]
\tikzstyle{box}=[
    draw,
    rounded corners=2pt,
    thick,
    align=center,
    text width=0.19\textwidth,
    minimum height=1.5cm,
    inner sep=5pt
]
\node[box, fill=blue!7] at (0,0) (sweep) {Static calibration cells\\
\footnotesize model $P$, prefix $U$, sampled tokens $C$, dropout grid};
\node[box, fill=green!8] at (3.7,0) (target) {Extract $p^\star$\\
\footnotesize grid best or bracketed quadratic minimum};
\node[box, fill=yellow!12] at (7.4,0) (fit) {Fit pressure law\\
\footnotesize $p^\star \approx A x + B y + Dxy + C_0$};
\node[box, fill=red!7] at (11.1,0) (stream) {Frozen streaming test\\
\footnotesize compare decay path to best static dropout grid};
\draw[->, thick] (sweep.east) -- (target.west);
\draw[->, thick] (target.east) -- (fit.west);
\draw[->, thick] (fit.east) -- (stream.west);
\node[align=center, text width=0.9\textwidth] at (5.55,-1.45) {
Inputs to the fitted law are absolute quantities available before or during a run:
$x=\log_{10}(P/U)$ measures capacity pressure and
$y=\log_{10}(C/U)$ measures repeated-sampling pressure.
};
\end{tikzpicture}
"""
    (FIG_DIR / "calibration_to_streaming_protocol.tex").write_text(
        content.strip() + "\n", encoding="utf-8"
    )


def write_tables(summaries: dict[str, dict], coeffs: dict[str, dict]) -> None:
    symbol_lines = [
        "\\begin{tabular}{ll}",
        "\\toprule",
        "Symbol & Meaning \\\\",
        "\\midrule",
        "$P$ & Trainable parameter count \\\\",
        "$U_t$ & Unique-token prefix currently revealed by stream stage $t$ \\\\",
        "$C_t$ & Cumulative sampled training tokens consumed by the optimizer \\\\",
        "$x_t$ & $\\log_{10}(P/U_t)$, capacity pressure \\\\",
        "$y_t$ & $\\log_{10}(C_t/U_t)$, repeated-sampling pressure \\\\",
        "$p_t$ & Dropout rate used at stream stage $t$ \\\\",
        "$p^\\star$ & Static-sweep target dropout for one calibration cell \\\\",
        "$A,B,D,C_0$ & Regime-specific fitted pressure-law coefficients \\\\",
        "$p_{\\min},p_{\\max}$ & Clamp bounds applied to generated dropout schedules \\\\",
        "\\bottomrule",
        "\\end{tabular}",
    ]
    (TABLE_DIR / "symbol_table.tex").write_text(
        "\n".join(symbol_lines) + "\n", encoding="utf-8"
    )

    regime_lines = [
        "\\begin{tabular}{llll}",
        "\\toprule",
        "Regime & Corpus source & Tokenizer / vocab & Streaming prefixes \\\\",
        "\\midrule",
        (
            "OpenWebText10K & Cached OpenWebText-style subset & BPE / 4096 & "
            "250k, 500k, 1M, 2M, 4M; plus 8M L20 tier \\\\"
        ),
        (
            "TinyStories & Cached TinyStories subset & BPE / 4096 & "
            "500k, 1M, 2M, 4M; plus 8M L20 tier \\\\"
        ),
        (
            "WikiText-103 & Cached WikiText-103 subset & BPE / 4096 & "
            "250k, 500k, 1M, 2M, 4M; plus 8M L20 tier \\\\"
        ),
        "\\bottomrule",
        "\\end{tabular}",
    ]
    (TABLE_DIR / "regime_setup_table.tex").write_text(
        "\n".join(regime_lines) + "\n", encoding="utf-8"
    )

    repro_lines = [
        "\\begin{tabular}{ll}",
        "\\toprule",
        "Item & Value \\\\",
        "\\midrule",
        "Repository & \\url{https://huggingface.co/cuber12/dropout-decay} \\\\",
        "Archived snapshot DOI & \\url{https://doi.org/10.5281/zenodo.20616633} \\\\",
        "Repository snapshot & Git tag \\texttt{v1.1-preprint} \\\\",
        "License & MIT-derived source with nanochat attribution retained \\\\",
        "Python & 3.11.15 \\\\",
        "PyTorch & 2.12.0 \\\\",
        "Operating system & macOS 26.4 arm64 \\\\",
        "Training backend & Apple Metal Performance Shaders (MPS); runner refuses CPU/CUDA \\\\",
        "\\bottomrule",
        "\\end{tabular}",
    ]
    (TABLE_DIR / "reproducibility_table.tex").write_text(
        "\n".join(repro_lines) + "\n", encoding="utf-8"
    )

    weight_lines = [
        "\\begin{tabular}{lll}",
        "\\toprule",
        "Cell property & Weight multiplier & Reason \\\\",
        "\\midrule",
        "Best grid point on dropout boundary & $0.30$ & Optimum may lie outside tested grid \\\\",
        "Quadratic minimum not bracketed & $0.50$ & Local parabola is less reliable \\\\",
        "Loss span below $0.02$ nats & $0.50$ & Very flat curve weakly identifies $p^\\star$ \\\\",
        (
            "Best-loss std. $\\sigma > 0$ & $1/(1+20\\sigma)$ & "
            "Noisier minima receive less influence \\\\"
        ),
        "\\bottomrule",
        "\\end{tabular}",
    ]
    (TABLE_DIR / "calibration_weight_table.tex").write_text(
        "\n".join(weight_lines) + "\n", encoding="utf-8"
    )

    protocol_lines = [
        "\\begin{tabular}{llll}",
        "\\toprule",
        "Experiment layer & What is run & Question answered & Statistical unit \\\\",
        "\\midrule",
        (
            "Static calibration & Per-regime dropout grid at fixed $(P,U,C)$ cells & "
            "Where does validation loss plateau over dropout? & 15--16 cells per regime \\\\"
        ),
        (
            "Coefficient fit & Per-regime regression from cell targets $p^\\star$ & "
            "Can pressure variables predict useful dropout? & MAE / leave-out MAE \\\\"
        ),
        (
            "4M final-prefix locked stream & Per-regime continuing model over growing prefixes & "
            "Does decay beat the best tested fixed dropout? & 5 paired seeds per regime \\\\"
        ),
        (
            "L20 8M extrapolation & Per-regime 51.84M-param model with frozen coefficients & "
            "Does the law extend to larger model/data scale? & 3 paired seeds per regime \\\\"
        ),
        (
            "Stage-wise audit & Per-regime decay vs best static at every prefix & "
            "Is the schedule final-loss or per-prefix optimal? & Prefix-level deltas \\\\"
        ),
        "\\bottomrule",
        "\\end{tabular}",
    ]
    (TABLE_DIR / "experiment_protocol_table.tex").write_text(
        "\n".join(protocol_lines) + "\n", encoding="utf-8"
    )

    budget_lines = [
        "\\begin{tabular}{llllrr}",
        "\\toprule",
        "Regime & Protocol & Prefix path & Sampled/stage & Batch & Block \\\\",
        "\\midrule",
    ]
    for key, paths in REGIMES.items():
        for protocol_key, protocol_label in [
            ("small", paths["small_protocol"]),
            ("l20", "8M final prefix / 3 seeds"),
        ]:
            budget = summaries[key][protocol_key]["budget"]
            prefix_path = " $\\rightarrow$ ".join(token_label(prefix) for prefix in budget["prefix_path"])
            sampled = ", ".join(token_label(value) for value in budget["sampled_per_stage"])
            batch = "/".join(str(value) for value in budget["batch_sizes"])
            block = str(budget["block_size"])
            budget_lines.append(
                f"{tex_escape(paths['label'])} & {tex_escape(protocol_label)} & "
                f"{prefix_path} & {sampled} & {batch} & {block} \\\\"
            )
    budget_lines.extend(["\\bottomrule", "\\end{tabular}"])
    (TABLE_DIR / "protocol_budget_table.tex").write_text(
        "\n".join(budget_lines) + "\n", encoding="utf-8"
    )

    coeff_lines = [
        "\\begin{tabular}{lrrrrrrrr}",
        "\\toprule",
        "Regime & Cells & A & B & D & C0 & Dropout MAE (abs. $p$) & L-model MAE & L-prefix MAE \\\\",
        "\\midrule",
    ]
    for key in REGIMES:
        row = coeffs[key]
        coeff_lines.append(
            f"{tex_escape(row['regime'])} & {row['cells']} & "
            f"{row['A']:.4f} & {row['B']:.4f} & {row['D']:.4f} & {row['C0']:.4f} & "
            f"{row['mae']:.4f} & {row['leave_model_mae']:.4f} & {row['leave_prefix_mae']:.4f} \\\\"
        )
    coeff_lines.extend(["\\bottomrule", "\\end{tabular}"])
    (TABLE_DIR / "coefficient_table.tex").write_text(
        "\n".join(coeff_lines) + "\n", encoding="utf-8"
    )

    evidence_lines = [
        "\\begin{tabular}{lrrrrr}",
        "\\toprule",
        "Regime & Calib. dropout MAE & 4M gain (nats) & 4M wins & 8M gain (nats) & 8M wins \\\\",
        "\\midrule",
    ]
    for key, paths in REGIMES.items():
        small = summaries[key]["small"]
        l20 = summaries[key]["l20"]
        evidence_lines.append(
            f"{tex_escape(paths['label'])} & {coeffs[key]['mae']:.4f} & "
            f"{small['improvement']:.4f} & {small['paired_wins']}/{small['n']} & "
            f"{l20['improvement']:.4f} & {l20['paired_wins']}/{l20['n']} \\\\"
        )
    evidence_lines.extend(["\\bottomrule", "\\end{tabular}"])
    (TABLE_DIR / "evidence_ladder_table.tex").write_text(
        "\n".join(evidence_lines) + "\n", encoding="utf-8"
    )

    val_lines = [
        "\\begin{tabular}{llrrrrr}",
        "\\toprule",
        "Regime & Protocol & Decay final (nats) & Best static final (nats) & Gain (nats) & Seeds & Paired wins \\\\",
        "\\midrule",
    ]
    for key, paths in REGIMES.items():
        for protocol_key, protocol_label in [
            ("small", paths["small_protocol"]),
            ("l20", "8M final prefix / 3 seeds"),
        ]:
            row = summaries[key][protocol_key]
            val_lines.append(
                f"{tex_escape(paths['label'])} & {tex_escape(protocol_label)} & "
                f"{row['condition_final_mean']:.4f} $\\pm$ {row['condition_final_std']:.4f} & "
                f"{row['best_static_mean']:.4f} $\\pm$ {row['best_static_std']:.4f} & "
                f"{row['improvement']:.4f} & {row['n']} & {row['paired_wins']}/{row['n']} \\\\"
            )
    val_lines.extend(["\\bottomrule", "\\end{tabular}"])
    (TABLE_DIR / "streaming_validation_table.tex").write_text(
        "\n".join(val_lines) + "\n", encoding="utf-8"
    )

    paired_lines = [
        "\\begin{tabular}{llrrrr}",
        "\\toprule",
        "Regime & Protocol & Paired gain mean & 95\\% bootstrap CI & Seeds & Wins \\\\",
        "\\midrule",
    ]
    for key, paths in REGIMES.items():
        for protocol_key, protocol_label in [
            ("small", paths["small_protocol"]),
            ("l20", "8M final prefix / 3 seeds"),
        ]:
            row = summaries[key][protocol_key]
            paired_lines.append(
                f"{tex_escape(paths['label'])} & {tex_escape(protocol_label)} & "
                f"{row['paired_gain_mean']:.4f} & "
                f"[{row['paired_gain_ci_low']:.4f}, {row['paired_gain_ci_high']:.4f}] & "
                f"{row['n']} & {row['paired_wins']}/{row['n']} \\\\"
            )
    paired_lines.extend(["\\bottomrule", "\\end{tabular}"])
    (TABLE_DIR / "paired_gain_ci_table.tex").write_text(
        "\n".join(paired_lines) + "\n", encoding="utf-8"
    )

    path_lines = [
        "\\begin{tabular}{llll}",
        "\\toprule",
        "Regime & Locked-stream prefix path & Frozen decay path & Final best static \\\\",
        "\\midrule",
    ]
    for key, paths in REGIMES.items():
        stage_rows = summaries[key]["l20"]["stage_deltas"]
        prefix_path = " $\\rightarrow$ ".join(token_label(row["prefix"]) for row in stage_rows)
        dropout_path = " $\\rightarrow$ ".join(
            f"{row['condition_dropout']:.3f}" for row in stage_rows
        )
        final_static = summaries[key]["l20"]["best_static_condition"].replace(
            "static_dropout_", "p="
        )
        path_lines.append(
            f"{tex_escape(paths['label'])} & {prefix_path} & {dropout_path} & "
            f"{tex_escape(final_static)} \\\\"
        )
    path_lines.extend(["\\bottomrule", "\\end{tabular}"])
    (TABLE_DIR / "l20_schedule_path_table.tex").write_text(
        "\n".join(path_lines) + "\n", encoding="utf-8"
    )

    l20_lines = [
        "\\begin{tabular}{lrrrrr}",
        "\\toprule",
        "Regime & Prefix & Decay dropout & Decay val (nats) & Best static val (nats) & Delta (nats) \\\\",
        "\\midrule",
    ]
    for key, paths in REGIMES.items():
        for row in summaries[key]["l20"]["stage_deltas"]:
            l20_lines.append(
                f"{tex_escape(paths['label'])} & {token_label(row['prefix'])} & "
                f"{row['condition_dropout']:.3f} & {row['condition_mean']:.4f} & "
                f"{row['best_static_mean']:.4f} & {row['delta']:+.4f} \\\\"
            )
    l20_lines.extend(["\\bottomrule", "\\end{tabular}"])
    (TABLE_DIR / "l20_stage_table.tex").write_text(
        "\n".join(l20_lines) + "\n", encoding="utf-8"
    )

    summary = {
        key: {
            "coefficients": coeffs[key],
            "small_stream": summaries[key]["small"],
            "l20_stream": summaries[key]["l20"],
        }
        for key in REGIMES
    }
    (PAPER_DIR / "paper_results_summary.json").write_text(
        json.dumps(summary, indent=2), encoding="utf-8"
    )


def main() -> None:
    FIG_DIR.mkdir(parents=True, exist_ok=True)
    TABLE_DIR.mkdir(parents=True, exist_ok=True)
    summaries = {}
    coeffs = {}
    for key, paths in REGIMES.items():
        coeffs[key] = coeff_summary(key, paths)
        summaries[key] = {
            "small": metric_stage_summary(paths["small_stream"], paths["small_condition"]),
            "l20": metric_stage_summary(paths["l20_stream"], paths["l20_condition"]),
        }
    write_calibration_curve_figure()
    write_static_optima_figure()
    write_protocol_schematic()
    write_l20_stage_delta_figure(summaries)
    write_final_improvement_figure(summaries)
    write_tables(summaries, coeffs)
    print(json.dumps({"figures": str(FIG_DIR), "tables": str(TABLE_DIR)}, indent=2))


if __name__ == "__main__":
    main()