path: root/scripts/generate_paper_tables.py

#!/usr/bin/env python3
"""Generate LaTeX table fragments for the paper appendices from row-level CSVs.

Each fragment contains only a tabular environment; captions and labels live in
the paper source. Uncertainty follows the paper's nesting convention: feedback
draws are averaged within each forward initialization first, and standard
errors are computed across forward-initialization means.
"""

from __future__ import annotations

import argparse
import csv
import math
import statistics
from collections import defaultdict
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--depth-dir", type=Path, default=Path("outputs/aaai_depth_experiments"))
    parser.add_argument(
        "--cnn-dir", type=Path, default=Path("outputs/cnn_initialization_validation")
    )
    parser.add_argument(
        "--lr-dir", type=Path, default=Path("outputs/learning_rate_compensation")
    )
    parser.add_argument(
        "--teacher-dir", type=Path, default=Path("outputs/teacher_test_gap_sgd_T8000")
    )
    parser.add_argument("--outdir", type=Path, default=Path("outputs/tables"))
    return parser.parse_args()


def read_rows(path: Path) -> list[dict[str, str]]:
    if not path.exists():
        raise FileNotFoundError(path)
    with path.open() as handle:
        return list(csv.DictReader(handle))


def sem(values: list[float]) -> float:
    if len(values) < 2:
        return float("nan")
    return statistics.stdev(values) / math.sqrt(len(values))


def write_fragment(path: Path, lines: list[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(lines) + "\n")
    print(f"wrote {path}")


def initialization_table(depth_dir: Path, outdir: Path) -> None:
    rows = read_rows(depth_dir / "initialization_rows.csv")
    cells: dict[tuple[int, int], dict[str, dict[str, float]]] = defaultdict(dict)
    for row in rows:
        key = (int(row["depth"]), int(row["init_seed"]))
        cells[key][row["rule"]] = {
            "pred": float(row["predicted_deficit"]),
            "mean": float(row["empirical_deficit"]),
            "se": float(row["empirical_stderr"]),
        }
    lines = [
        r"\begin{tabular}{@{}rrccc@{}}",
        r"\toprule",
        r"Depth & Init & Prediction & FA measured (SE) & DFA measured (SE) \\",
        r"\midrule",
    ]
    for (depth, init_seed), by_rule in sorted(cells.items()):
        fa, dfa = by_rule["FA"], by_rule["DFA"]
        assert abs(fa["pred"] - dfa["pred"]) < 1e-12, (depth, init_seed)
        lines.append(
            f"{depth} & {init_seed - 10_000} & {fa['pred']:.6f} & "
            f"{fa['mean']:.6f} ({fa['se']:.2g}) & "
            f"{dfa['mean']:.6f} ({dfa['se']:.2g}) \\\\"
        )
    lines += [r"\bottomrule", r"\end{tabular}"]
    write_fragment(outdir / "initialization_rows.tex", lines)


def finite_time_table(depth_dir: Path, outdir: Path) -> None:
    rows = read_rows(depth_dir / "finite_time_rows.csv")
    cells: dict[tuple[int, int], dict[int, dict[str, list[float]]]] = defaultdict(
        lambda: defaultdict(lambda: defaultdict(list))
    )
    fields = ["empirical_gap", "fixed_gap", "velocity_gap", "retangent_gap"]
    for row in rows:
        cell = cells[(int(row["depth"]), int(row["width"]))][int(row["init_seed"])]
        for field in fields:
            cell[field].append(float(row[field]))
    lines = [
        r"\begin{tabular}{@{}rrccccc@{}}",
        r"\toprule",
        r"Depth & Width & Measured & SE & Frozen & Linear & Relinearized \\",
        r"\midrule",
    ]
    for (depth, width), by_init in sorted(cells.items()):
        init_means = {
            field: [statistics.mean(values[field]) for _, values in sorted(by_init.items())]
            for field in fields
        }
        measured = statistics.mean(init_means["empirical_gap"])
        lines.append(
            f"{depth} & {width} & {measured:.6f} & {sem(init_means['empirical_gap']):.2g} & "
            f"{statistics.mean(init_means['fixed_gap']):.6f} & "
            f"{statistics.mean(init_means['velocity_gap']):.6f} & "
            f"{statistics.mean(init_means['retangent_gap']):.6f} \\\\"
        )
    lines += [r"\bottomrule", r"\end{tabular}"]
    write_fragment(outdir / "finite_time_cells.tex", lines)


def cnn_table(cnn_dir: Path, outdir: Path) -> None:
    rows = read_rows(cnn_dir / "cnn_initialization_rows.csv")
    lines = [
        r"\begin{tabular}{@{}lrccc@{}}",
        r"\toprule",
        r"Rule & Init & Prediction & Measured & SE \\",
        r"\midrule",
    ]
    for row in sorted(rows, key=lambda r: (r["rule"], int(r["init_seed"]))):
        lines.append(
            f"{row['rule']} & {int(row['init_seed']) - 10_000} & "
            f"{float(row['predicted_deficit']):.6f} & "
            f"{float(row['empirical_deficit']):.6f} & "
            f"{float(row['empirical_stderr']):.2g} \\\\"
        )
    lines += [r"\bottomrule", r"\end{tabular}"]
    write_fragment(outdir / "cnn_initialization_rows.tex", lines)


def learning_rate_table(lr_dir: Path, outdir: Path) -> None:
    rows = read_rows(lr_dir / "evaluation_rows.csv")
    regime_order = {"same_lr": 0, "initial_compensation": 1, "independently_tuned": 2}
    regime_label = {
        "same_lr": "Same rate",
        "initial_compensation": "Initial compensation",
        "independently_tuned": "Independently selected",
    }
    cells: dict[tuple[int, str], dict[int, list[float]]] = defaultdict(lambda: defaultdict(list))
    fa_rates: dict[tuple[int, str], list[float]] = defaultdict(list)
    bp_rates: dict[tuple[int, str], set[float]] = defaultdict(set)
    for row in rows:
        assert row["stable"] == "True", row
        key = (int(row["depth"]), row["regime"])
        cells[key][int(row["init_seed"])].append(float(row["gap_to_bp"]))
        fa_rates[key].append(float(row["fa_lr"]))
        bp_rates[key].add(float(row["bp_lr"]))
    lines = [
        r"\begin{tabular}{@{}rlcccc@{}}",
        r"\toprule",
        r"Depth & Regime & BP rate & Mean FA rate & Gap & SE \\",
        r"\midrule",
    ]
    for (depth, regime), by_init in sorted(
        cells.items(), key=lambda item: (item[0][0], regime_order[item[0][1]])
    ):
        init_means = [statistics.mean(values) for _, values in sorted(by_init.items())]
        (bp_rate,) = bp_rates[(depth, regime)]
        lines.append(
            f"{depth} & {regime_label[regime]} & {bp_rate:g} & "
            f"{statistics.mean(fa_rates[(depth, regime)]):.4g} & "
            f"{statistics.mean(init_means):.5f} & {sem(init_means):.2g} \\\\"
        )
    lines += [r"\bottomrule", r"\end{tabular}"]
    write_fragment(outdir / "learning_rate_regimes.tex", lines)


def teacher_table(teacher_dir: Path, outdir: Path) -> None:
    rows = [r for r in read_rows(teacher_dir / "runs.csv") if r["run_type"] == "fa"]
    by: dict[int, dict[int, dict[str, list[float]]]] = defaultdict(
        lambda: defaultdict(lambda: defaultdict(list))
    )
    for row in rows:
        cell = by[int(row["width"])][int(row["init_seed"])]
        cell["train"].append(float(row["train_gap_to_bp"]))
        cell["test"].append(float(row["test_gap_to_bp"]))
    lines = [
        r"\begin{tabular}{@{}rcc@{}}",
        r"\toprule",
        r"Width & Train difference & Test difference \\",
        r"\midrule",
    ]
    for width, by_init in sorted(by.items()):
        train_means = [statistics.mean(v["train"]) for _, v in sorted(by_init.items())]
        test_means = [statistics.mean(v["test"]) for _, v in sorted(by_init.items())]
        lines.append(
            f"{width} & \\({statistics.mean(train_means):.4f}\\pm{sem(train_means):.4f}\\) & "
            f"\\({statistics.mean(test_means):.4f}\\pm{sem(test_means):.4f}\\) \\\\"
        )
    lines += [r"\bottomrule", r"\end{tabular}"]
    write_fragment(outdir / "teacher_nested_se.tex", lines)


TRAIN_SIZES = [64, 96, 128, 160, 192, 224, 256, 320]


def _corr(xs: list[float], ys: list[float]) -> float:
    mx, my = statistics.mean(xs), statistics.mean(ys)
    cov = sum((x - mx) * (y - my) for x, y in zip(xs, ys))
    return cov / math.sqrt(
        sum((x - mx) ** 2 for x in xs) * sum((y - my) ** 2 for y in ys)
    )


def train_size_table(outputs: Path, outdir: Path) -> None:
    lines = [
        r"\begin{tabular}{@{}rcccc@{}}",
        r"\toprule",
        r"\(N\) & Measured & Frozen & Linear & Relin.\ \\",
        r"\midrule",
    ]
    for n in TRAIN_SIZES:
        rows = read_rows(
            outputs
            / f"compressed_operator_s20_256traj_T50_width64_N{n}"
            / "compressed_operator_rows.csv"
        )
        measured = [float(r["empirical_gap"]) for r in rows]
        maes = {
            field: statistics.mean(
                abs(float(r[field]) - m) for r, m in zip(rows, measured)
            )
            for field in ("fixed_gap", "linear_gap", "retangent_gap")
        }
        lines.append(
            f"{n} & {statistics.mean(measured):.5f} & {maes['fixed_gap']:.5f} & "
            f"{maes['linear_gap']:.5f} & {maes['retangent_gap']:.5f} \\\\"
        )
    lines += [r"\bottomrule", r"\end{tabular}"]
    write_fragment(outdir / "train_size_predictors.tex", lines)


def snapshot_location_table(outputs: Path, outdir: Path) -> None:
    rows = [
        r
        for n in TRAIN_SIZES
        for r in read_rows(
            outputs
            / f"compressed_operator_retangent_T50_width64_N{n}"
            / "compressed_operator_rows.csv"
        )
    ]
    lines = [
        r"\begin{tabular}{@{}rccccc@{}}",
        r"\toprule",
        r"\(s\) & \(s/T\) & Frozen & Linear & Relin.\ & Corr.\ \\",
        r"\midrule",
    ]
    for s in ("5", "10", "20"):
        sub = [r for r in rows if r["early_steps"] == s]
        measured = [float(r["empirical_gap"]) for r in sub]
        maes = {
            field: statistics.mean(
                abs(float(r[field]) - m) for r, m in zip(sub, measured)
            )
            for field in ("fixed_gap", "linear_gap", "retangent_gap")
        }
        corr = _corr([float(r["linear_gap"]) for r in sub], measured)
        lines.append(
            f"{s} & {int(s)/50:.1f} & {maes['fixed_gap']:.5f} & "
            f"{maes['linear_gap']:.5f} & {maes['retangent_gap']:.5f} & {corr:.5f} \\\\"
        )
    lines += [r"\bottomrule", r"\end{tabular}"]
    write_fragment(outdir / "snapshot_location.tex", lines)


def stress_grid_table(outputs: Path, outdir: Path) -> None:
    rows = read_rows(outputs / "operator_stress_grid_summary" / "stress_grid_metrics.csv")
    lines = [
        r"\begin{tabular}{@{}rrrrrcccc@{}}",
        r"\toprule",
        r"\(d\) & \(w\) & \(T\) & \(s\) & Rows & Frozen MAE & Linear MAE & Relin.\ MAE & Linear corr.\ \\",
        r"\midrule",
    ]
    for row in rows:
        lines.append(
            f"{int(row['hidden_layers'])} & {int(row['width'])} & "
            f"{int(row['target_steps'])} & {int(row['early_steps'])} & {int(row['rows'])} & "
            f"{float(row['fixed_mae']):.5f} & {float(row['linear_mae']):.5f} & "
            f"{float(row['retangent_mae']):.5f} & {float(row['linear_corr']):.5f} \\\\"
        )
    lines += [r"\bottomrule", r"\end{tabular}"]
    write_fragment(outdir / "stress_grid.tex", lines)


def main() -> None:
    args = parse_args()
    initialization_table(args.depth_dir, args.outdir)
    finite_time_table(args.depth_dir, args.outdir)
    cnn_table(args.cnn_dir, args.outdir)
    learning_rate_table(args.lr_dir, args.outdir)
    teacher_table(args.teacher_dir, args.outdir)
    outputs = args.depth_dir.parent
    train_size_table(outputs, args.outdir)
    snapshot_location_table(outputs, args.outdir)
    stress_grid_table(outputs, args.outdir)


if __name__ == "__main__":
    main()