repro-bundle-v0/src/analysis.py

"""Shared curve-fitting / plotting helpers for the claim scripts."""
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt


def loglog_slope(xs, ys):
    """OLS fit of log(ys) ~ a + slope*log(xs). Returns (slope, intercept, r2)."""
    lx, ly = np.log(xs), np.log(ys)
    A = np.vstack([lx, np.ones_like(lx)]).T
    (slope, intercept), residuals, *_ = np.linalg.lstsq(A, ly, rcond=None)
    pred = A @ [slope, intercept]
    ss_res = np.sum((ly - pred) ** 2)
    ss_tot = np.sum((ly - ly.mean()) ** 2)
    r2 = 1 - ss_res / ss_tot if ss_tot > 0 else float('nan')
    return float(slope), float(intercept), float(r2)


def mean_sem(runs_by_x):
    """runs_by_x: dict x -> list of metric values (over seeds). Returns
    sorted (xs, means, sems)."""
    xs = sorted(runs_by_x.keys())
    means = np.array([np.mean(runs_by_x[x]) for x in xs])
    sems = np.array([np.std(runs_by_x[x], ddof=1) / np.sqrt(len(runs_by_x[x]))
                      if len(runs_by_x[x]) > 1 else 0.0 for x in xs])
    return np.array(xs, dtype=float), means, sems


def plot_scaling(xs, means, sems, slope, intercept, xlabel, title, out_path,
                  claimed_slope=None, extra_curves=None, ylabel='avg metric'):
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.2))

    ax = axes[0]
    ax.errorbar(xs, means, yerr=sems, fmt='o', ms=5, capsize=3, label='measured', color='#2563eb')
    xx = np.geomspace(xs.min(), xs.max(), 100)
    ax.plot(xx, np.exp(intercept) * xx ** slope, '--', color='#2563eb',
             label=f'fit slope={slope:.3f}')
    if claimed_slope is not None:
        c = np.exp(intercept) * xs[0] ** slope / xs[0] ** claimed_slope
        ax.plot(xx, c * xx ** claimed_slope, ':', color='#dc2626',
                 label=f'theory slope={claimed_slope:.3f}')
    if extra_curves:
        for label, (xc, yc) in extra_curves.items():
            ax.plot(xc, yc, '-.', label=label)
    ax.set_xscale('log'); ax.set_yscale('log')
    ax.set_xlabel(xlabel); ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.legend(fontsize=8)
    ax.grid(alpha=0.3)

    ax = axes[1]
    ax.errorbar(xs, means, yerr=sems, fmt='o-', ms=5, capsize=3, color='#16a34a')
    ax.set_xscale('log')
    ax.set_xlabel(xlabel); ax.set_ylabel(ylabel)
    ax.set_title(title + ' (linear y)')
    ax.grid(alpha=0.3)

    fig.tight_layout()
    fig.savefig(out_path, dpi=150)
    plt.close(fig)


def plot_collapse(theory_x, measured_y, out_path, xlabel, title):
    fig, ax = plt.subplots(figsize=(5, 4.2))
    ax.scatter(theory_x, measured_y, s=25, color='#7c3aed')
    slope, intercept, r2 = loglog_slope(theory_x, measured_y)
    xx = np.geomspace(theory_x.min(), theory_x.max(), 100)
    ax.plot(xx, np.exp(intercept) * xx ** slope, '--', color='#7c3aed',
            label=f'fit slope={slope:.3f}, R^2={r2:.3f}')
    ax.set_xscale('log'); ax.set_yscale('log')
    ax.set_xlabel(xlabel); ax.set_ylabel('measured avg metric')
    ax.set_title(title)
    ax.legend(fontsize=8)
    ax.grid(alpha=0.3)
    fig.tight_layout()
    fig.savefig(out_path, dpi=150)
    plt.close(fig)
    return slope, intercept, r2