batch_numeric_audit.py

# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy>=1.26", "scipy>=1.11"]
# ///
"""Small independent numerical audits for six ICML 2026 reproductions.

These are mechanism checks, not substitutes for the papers' full experiments.
The single script is used locally and in one shared Hugging Face CPU Job.
"""

from __future__ import annotations

import argparse
import itertools
import json
from pathlib import Path

import numpy as np
from scipy.linalg import solve_discrete_lyapunov


def mnl_design(rng: np.random.Generator) -> dict:
    n, d, k = 8, 3, 3
    x = rng.normal(size=(n, d))
    theta = rng.normal(size=d) / 3
    sets = [s for r in range(1, k + 1) for s in itertools.combinations(range(n), r)]
    infos = []
    for s in sets:
        z = x[list(s)]
        e = np.exp(z @ theta)
        p = e / (1 + e.sum())
        mean = (p[:, None] * z).sum(0)
        h = (p[:, None, None] * np.einsum("ni,nj->nij", z, z)).sum(0) - np.outer(mean, mean)
        infos.append(h + 1e-3 * np.eye(d))
    infos = np.asarray(infos)
    w = np.ones(len(sets)) / len(sets)
    initial = float(np.linalg.slogdet(np.einsum("i,ijk->jk", w, infos))[1])
    lmo_match = True
    for t in range(300):
        m = np.einsum("i,ijk->jk", w, infos)
        scores = np.einsum("ij,kji->k", np.linalg.inv(m), infos)
        idx_vec = int(scores.argmax())
        idx_loop = max(range(len(infos)), key=lambda j: float(np.trace(np.linalg.solve(m, infos[j]))))
        lmo_match &= idx_vec == idx_loop
        gamma = 2 / (t + 2)
        w *= 1 - gamma
        w[idx_vec] += gamma
    final = float(np.linalg.slogdet(np.einsum("i,ijk->jk", w, infos))[1])
    return {
        "candidate_assortments": len(sets),
        "vectorized_lmo_matches_bruteforce": bool(lmo_match),
        "logdet_initial": initial,
        "logdet_final": final,
        "logdet_gain": final - initial,
        "active_support": int((w > 1e-5).sum()),
    }


def wdist(a: np.ndarray, b: np.ndarray, p: int) -> float:
    a, b = a.copy(), b.copy()
    i = j = 0
    cost = 0.0
    while i < len(a) and j < len(b):
        flow = min(a[i], b[j])
        cost += flow * abs(i - j) ** p
        a[i] -= flow
        b[j] -= flow
        if a[i] <= 1e-14:
            i += 1
        if b[j] <= 1e-14:
            j += 1
    return cost ** (1 / p)


def wasserstein_noise(rng: np.random.Generator) -> dict:
    n = 32
    base = np.ones(n) / n
    sigmas = np.geomspace(2e-5, 8e-4, 7)
    means = {1: [], 2: []}
    for sigma in sigmas:
        vals = {1: [], 2: []}
        for _ in range(300):
            eps = rng.normal(scale=sigma, size=n)
            eps -= eps.mean()
            noisy = np.maximum(base + eps, 1e-12)
            noisy /= noisy.sum()
            for p in (1, 2):
                vals[p].append(wdist(base, noisy, p))
        for p in (1, 2):
            means[p].append(float(np.mean(vals[p])))
    slopes = {
        f"W{p}_loglog_slope": float(np.polyfit(np.log(sigmas), np.log(means[p]), 1)[0])
        for p in (1, 2)
    }
    return {
        **slopes,
        "expected_mechanism_slopes": {"W1": 1.0, "W2": 0.5},
        "sigma_grid": sigmas.tolist(),
        "mean_W1": means[1],
        "mean_W2": means[2],
    }


def riemannian_dueling(rng: np.random.Generator) -> dict:
    d = 20
    x = rng.normal(size=d)
    x /= np.linalg.norm(x)
    a = rng.normal(size=d)
    g = a - x * (x @ a)
    g /= np.linalg.norm(g)
    estimates = {}
    for m in (200, 2_000, 20_000):
        u = rng.normal(size=(m, d))
        u -= (u @ x)[:, None] * x
        u /= np.linalg.norm(u, axis=1)[:, None]
        est = (np.sign(u @ g)[:, None] * u).mean(0)
        estimates[str(m)] = float(est @ g / np.linalg.norm(est))
    return {
        "dimension": d,
        "cosine_to_normalized_gradient": estimates,
        "tangent_violation": float(abs(x @ est)),
    }


def sinkhorn_kernel(rng: np.random.Generator) -> dict:
    x = np.sort(rng.uniform(size=120))
    k = np.exp(-((x[:, None] - x[None, :]) ** 2) / (2 * 0.08**2)) + 1e-8
    d = np.ones(len(x))
    errors = []
    for _ in range(40):
        row = d * (k @ d)
        d /= np.sqrt(row)
        s = d[:, None] * k * d[None, :]
        errors.append(float(np.max(abs(s.sum(1) - 1))))
    eig = np.linalg.eigvalsh(s)
    return {
        "iterations_to_row_error_below_0.1pct": next((i + 1 for i, e in enumerate(errors) if e < 1e-3), None),
        "final_row_error": errors[-1],
        "symmetry_error": float(np.max(abs(s - s.T))),
        "minimum_entry": float(s.min()),
        "eigenvalue_min": float(eig.min()),
        "eigenvalue_max": float(eig.max()),
    }


def knockoff_threshold(w: np.ndarray, q: float) -> float:
    for t in np.sort(np.unique(np.abs(w[w != 0]))):
        if (1 + np.sum(w <= -t)) / max(1, np.sum(w >= t)) <= q:
            return float(t)
    return float("inf")


def semi_knockoffs(rng: np.random.Generator) -> dict:
    p, nonnull, reps, q = 400, 80, 1_000, 0.1
    fdps, powers = [], []
    for _ in range(reps):
        w = rng.normal(size=p)
        w[:nonnull] += 4.0
        t = knockoff_threshold(w, q)
        sel = np.flatnonzero(w >= t)
        false = np.sum(sel >= nonnull)
        fdps.append(false / max(1, len(sel)))
        powers.append(np.sum(sel < nonnull) / nonnull)
    return {
        "target_fdr": q,
        "empirical_fdr": float(np.mean(fdps)),
        "mean_power": float(np.mean(powers)),
        "replicates": reps,
        "scope": "sign-flip/knockoff+ calibration proxy; not the released Semi-knockoffs implementation",
    }


def sgmcmc_covariance(rng: np.random.Generator) -> dict:
    h = np.diag([0.7, 1.4, 2.2])
    target = np.linalg.inv(h)
    rows = []
    for lam in (0.005, 0.01, 0.02, 0.04, 0.08):
        a = np.eye(3) - lam * h
        exact = solve_discrete_lyapunov(a, 2 * lam * np.eye(3))
        rel = np.linalg.norm(exact - target, 2) / np.linalg.norm(target, 2)
        state = np.zeros(3)
        samples = []
        for t in range(180_000):
            state = a @ state + np.sqrt(2 * lam) * rng.normal(size=3)
            if t >= 20_000 and t % 10 == 0:
                samples.append(state.copy())
        empirical = np.cov(np.asarray(samples), rowvar=False)
        mc = np.linalg.norm(empirical - exact, 2) / np.linalg.norm(exact, 2)
        rows.append({"lambda": lam, "proxy_relative_error": float(rel), "mc_to_exact_error": float(mc)})
    slope = float(np.polyfit(np.log([r["lambda"] for r in rows]), np.log([r["proxy_relative_error"] for r in rows]), 1)[0])
    return {"step_size_error_loglog_slope": slope, "rows": rows}


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", default="batch_results.json")
    args = parser.parse_args()
    rng = np.random.default_rng(20260725)
    results = {
        "scope": "Independent lightweight numerical audits; not full benchmark replications or proof replacements.",
        "optimal_design_mnl": mnl_design(rng),
        "wasserstein_noise": wasserstein_noise(rng),
        "riemannian_dueling": riemannian_dueling(rng),
        "sinkhorn_diffusion": sinkhorn_kernel(rng),
        "semi_knockoffs": semi_knockoffs(rng),
        "sgmcmc_covariance": sgmcmc_covariance(rng),
    }
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()