repro-bundle/scripts/spectrum_fit_diagnostics.py

#!/usr/bin/env python3
"""Compute log-log power-law fit diagnostics from persisted spectrum CSVs."""

import argparse
import csv
import json
from collections import defaultdict
from pathlib import Path

import numpy as np


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("csv", type=Path)
    ap.add_argument("--output", type=Path, required=True)
    args = ap.parse_args()
    grouped: dict[str, list[tuple[int, float]]] = defaultdict(list)
    with args.csv.open(encoding="utf-8") as handle:
        for row in csv.DictReader(handle):
            grouped[row["model"]].append((int(row["rank"]), float(row["eigenvalue"])))
    diagnostics = {}
    for model, values in grouped.items():
        arr = np.asarray(values)
        selected = (arr[:, 0] >= 11) & (arr[:, 0] <= 500)
        x, y = np.log(arr[selected, 0]), np.log(arr[selected, 1])
        coef = np.polyfit(x, y, 1)
        pred = np.polyval(coef, x)
        r2 = 1.0 - float(np.sum((y - pred) ** 2) / np.sum((y - y.mean()) ** 2))
        diagnostics[model] = {"alpha": float(-coef[0]), "log_log_r_squared": r2, "fit_ranks": [11, 500]}
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(diagnostics, indent=2), encoding="utf-8")
    print(json.dumps(diagnostics, indent=2))


if __name__ == "__main__":
    main()