scripts/summarize_cross_regime_backtest.py

#!/usr/bin/env python3
"""
Summarize cross-regime dropout coefficient backtests.

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 statistics


ROOT = Path("runs/coefficient_calibration/cross_regime_backtest")
REPORT_PATH = Path("docs/cross_regime_backtest_report.md")


FIT_DIRS = {
    "openwebtext10k_main_base": ROOT / "openwebtext10k_main_base",
    "openwebtext10k_main_interaction": ROOT / "openwebtext10k_main_interaction",
    "openwebtext10k_plus_5m_base": ROOT / "openwebtext10k_plus_5m_base",
    "openwebtext10k_plus_5m_interaction": ROOT / "openwebtext10k_plus_5m_interaction",
    "pooled_previous_plus_corpus_probes_base": ROOT
    / "pooled_previous_plus_corpus_probes_base",
    "pooled_previous_plus_corpus_probes_interaction": ROOT
    / "pooled_previous_plus_corpus_probes_interaction",
    "tinystories_all_base": ROOT / "tinystories_all_base",
    "tinystories_all_interaction": ROOT / "tinystories_all_interaction",
    "tinystories_all_quadratic": ROOT / "tinystories_all_quadratic",
    "corpus_pressure_probe_pooled_base": ROOT / "corpus_pressure_probe_pooled_base",
}


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


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


def coefficients(fit_dir: Path) -> dict[str, float]:
    payload = read_json(fit_dir / "coefficients.json")
    return {key: float(value) for key, value in payload["coefficients"].items()}


def fit_payload(name: str) -> dict:
    payload = read_json(FIT_DIRS[name] / "coefficients.json")
    payload["name"] = name
    return payload


def predict(cell: dict, coef: dict[str, float], feature_set: str) -> float:
    x = float(cell["x_model_pressure"])
    y = float(cell["x_sample_pressure"])
    if feature_set == "base":
        return coef["A"] * x + coef["B"] * y + coef["C0"]
    if feature_set == "interaction":
        return coef["A"] * x + coef["B"] * y + coef["D"] * x * y + coef["C0"]
    if feature_set == "quadratic":
        return (
            coef["A"] * x
            + coef["B"] * y
            + coef["Qp"] * x * x
            + coef["Qc"] * y * y
            + coef["C0"]
        )
    raise ValueError(f"unsupported feature set: {feature_set}")


def error_metrics(
    fit_name: str,
    test_cells_name: str,
) -> dict[str, float | str]:
    fit = fit_payload(fit_name)
    coef = coefficients(FIT_DIRS[fit_name])
    cells = read_cells(FIT_DIRS[test_cells_name])
    errors = []
    for cell in cells:
        errors.append(predict(cell, coef, fit["feature_set"]) - float(cell["target_dropout"]))
    return {
        "fit": fit_name,
        "test": test_cells_name,
        "feature_set": fit["feature_set"],
        "n": len(errors),
        "rmse": math.sqrt(statistics.fmean(error * error for error in errors)),
        "mae": statistics.fmean(abs(error) for error in errors),
        "bias": statistics.fmean(errors),
    }


def fmt(value: float) -> str:
    return f"{value:.4f}"


def coef_summary(name: str) -> str:
    payload = fit_payload(name)
    coef = payload["coefficients"]
    ordered = []
    for key in ["A", "B", "D", "Qp", "Qc", "C0"]:
        if key in coef:
            ordered.append(f"{key}={coef[key]:+.4f}")
    return ", ".join(ordered)


def metrics_row(name: str, label: str) -> str:
    payload = fit_payload(name)
    metrics = payload["metrics"]
    cv = payload["cv"]
    leave_model = cv.get("leave_model", {})
    leave_prefix = cv.get("leave_prefix", {})
    leave_source = cv.get("leave_source", {})
    return (
        f"| {label} | `{payload['feature_set']}` | {int(metrics['n'])} | "
        f"{fmt(metrics['mae'])} | {fmt(metrics['rmse'])} | "
        f"{fmt(leave_model.get('mae', float('nan')))} | "
        f"{fmt(leave_prefix.get('mae', float('nan')))} | "
        f"{fmt(leave_source.get('mae', float('nan')))} | "
        f"`{coef_summary(name)}` |"
    )


def transfer_row(item: dict[str, float | str]) -> str:
    return (
        f"| `{item['fit']}` | `{item['test']}` | `{item['feature_set']}` | "
        f"{item['n']} | {fmt(float(item['mae']))} | "
        f"{fmt(float(item['rmse']))} | {fmt(float(item['bias']))} |"
    )


def corpus_probe_rows() -> list[str]:
    rows = []
    for cell in read_cells(FIT_DIRS["corpus_pressure_probe_pooled_base"]):
        rows.append(
            "| "
            f"`{cell['source']}` | `{cell['model_name']}` | "
            f"{float(cell['x_model_pressure']):.4f} | "
            f"{float(cell['x_sample_pressure']):.4f} | "
            f"{float(cell['target_dropout']):.4f} | "
            f"{cell['boundary_optimum']} | {cell['bracketed_optimum']} |"
        )
    return rows


def write_transfer_csv(items: list[dict[str, float | str]]) -> None:
    path = ROOT / "cross_regime_transfer.csv"
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(
            handle,
            fieldnames=["fit", "test", "feature_set", "n", "rmse", "mae", "bias"],
        )
        writer.writeheader()
        for item in items:
            writer.writerow(item)


def main() -> None:
    transfer_items = [
        error_metrics("openwebtext10k_plus_5m_base", "tinystories_all_base"),
        error_metrics("tinystories_all_base", "openwebtext10k_plus_5m_base"),
        error_metrics("openwebtext10k_plus_5m_interaction", "tinystories_all_interaction"),
        error_metrics("tinystories_all_interaction", "openwebtext10k_plus_5m_interaction"),
        error_metrics(
            "pooled_previous_plus_corpus_probes_interaction",
            "tinystories_all_interaction",
        ),
        error_metrics(
            "tinystories_all_interaction",
            "pooled_previous_plus_corpus_probes_interaction",
        ),
    ]
    write_transfer_csv(transfer_items)

    lines = [
        "# Cross-Regime Backtest Report",
        "",
        "Date: 2026-05-30",
        "",
        "This report validates saved previous-regime results before launching any new",
        "training. No MPS training was run for this report; all numbers are offline",
        "fits or backtests from saved `model_selection.csv` and `metrics.jsonl` files.",
        "",
        "## Result Sources",
        "",
        "| Regime/source | Run directories | Role |",
        "|---|---|---|",
        "| OpenWebText10K main | `runs/screen_static/20260525-133008` | OpenWebText10K static screen, 15 cells |",
        "| OpenWebText10K + 5M | OpenWebText10K main + `runs/screen_static/20260525-122824` | OpenWebText10K regime plus low-pressure 5M extension, 18 cells |",
        "| corpus probes | `runs/corpus_difficulty_pressure_{local,tinystories,wikitext103}` | diagnostic rows for corpus sensitivity |",
        "| current TinyStories | `runs/regime_calibration_tinystories_*` | current coefficient evidence, 16 cells |",
        "",
        "## Within-Regime Fits",
        "",
        "| Fit | Feature set | Cells | MAE | RMSE | Leave-model MAE | Leave-prefix MAE | Leave-source MAE | Coefficients |",
        "|---|---|---:|---:|---:|---:|---:|---:|---|",
        metrics_row("openwebtext10k_main_base", "OpenWebText10K main"),
        metrics_row("openwebtext10k_main_interaction", "OpenWebText10K main"),
        metrics_row("openwebtext10k_plus_5m_base", "OpenWebText10K + 5M"),
        metrics_row("openwebtext10k_plus_5m_interaction", "OpenWebText10K + 5M"),
        metrics_row("tinystories_all_base", "current TinyStories"),
        metrics_row("tinystories_all_interaction", "current TinyStories"),
        metrics_row("tinystories_all_quadratic", "current TinyStories"),
        "",
        "### Reading",
        "",
        "- The OpenWebText10K regime supports the interaction form: adding `D*x*y`",
        "  reduces MAE from `0.0389` to `0.0148` on the local+5M cells.",
        "- The current TinyStories regime also supports the interaction form: MAE",
        "  drops from `0.0435` to `0.0180`.",
        "- The quadratic alternative does not beat interaction on TinyStories, so",
        "  the interaction form remains the preferred complexity level.",
        "",
        "## Cross-Regime Transfer Diagnostics",
        "",
        "These rows deliberately apply coefficients fitted on one regime to cells from",
        "another regime. Good within-regime fit with weaker raw transfer supports the",
        "current framing: formula structure transfers, coefficients are regime-specific.",
        "",
        "| Fit coefficients from | Test cells from | Feature set | n | MAE | RMSE | Bias |",
        "|---|---|---|---:|---:|---:|---:|",
    ]
    lines.extend(transfer_row(item) for item in transfer_items)
    lines.extend(
        [
            "",
            "### Reading",
            "",
            "- Raw coefficient transfer is poor compared with within-regime error.",
            "- Previous-local interaction coefficients overpredict/underpredict current",
            "  TinyStories cells enough that we should not claim universal numeric",
            "  coefficients.",
            "- This validates the plan's regime-specific coefficient rule.",
            "",
            "## Corpus-Probe Limitation",
            "",
            "The corpus probe rows are useful as a warning, not as a complete new regime",
            "fit. They use almost identical pressure variables but have different dropout",
            "targets, so a formula using only `P/U` and `C/U` cannot distinguish them.",
            "",
            "| Source | Model | x=log10(P/U) | y=log10(C/U) | Target dropout | Boundary | Bracketed |",
            "|---|---|---:|---:|---:|---|---|",
        ]
    )
    lines.extend(corpus_probe_rows())
    lines.extend(
        [
            "",
            "### Reading",
            "",
            "Because corpus identity changed while pressure variables were effectively",
            "held fixed, these probe rows demonstrate corpus sensitivity. They should not",
            "be pooled into a single universal coefficient fit unless the formula gains a",
            "corpus/regime feature or each corpus receives its own calibrated coefficients.",
            "",
            "## Decision",
            "",
            "The previous regime validation does not block the current plan. It strengthens",
            "the formula-family claim because both the OpenWebText10K regime and the current",
            "TinyStories regime prefer the interaction pressure law over first-order ABC.",
            "",
            "The validated claim remains:",
            "",
            "```text",
            "pressure-law structure transfers across regimes;",
            "numeric coefficients are regime-calibrated.",
            "```",
            "",
            "The next step is not a new broad sweep. The next step is to use this report",
            "as the offline backtest gate, then proceed to narrowed TinyStories multi-seed",
            "streaming only if we accept the regime-specific coefficient framing.",
            "",
            "## Artifacts",
            "",
            "- Fit outputs: `runs/coefficient_calibration/cross_regime_backtest/`",
            "- Transfer table: `runs/coefficient_calibration/cross_regime_backtest/cross_regime_transfer.csv`",
        ]
    )
    REPORT_PATH.write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(json.dumps({"report": str(REPORT_PATH), "transfer_rows": len(transfer_items)}, indent=2))


if __name__ == "__main__":
    main()