batch20_audit.py

#!/usr/bin/env python3
"""Small deterministic numerical audits for seven ICML 2026 theory papers.

These checks test mechanisms or necessary consequences of the stated results;
they are deliberately not presented as substitutes for the papers' proofs or
full experiments.
"""

from __future__ import annotations

import json
import math
import time

import numpy as np


def bandit() -> dict:
    # Adversarial linear-bandit lower-bound sanity check: for independent
    # Rademacher coordinate losses, the best fixed unit-ball comparator has
    # hindsight loss -||sum_t g_t||_2 and the zero learner has loss zero.
    rng = np.random.default_rng(1)
    out = {}
    for d, t in [(4, 400), (16, 400), (16, 1600)]:
        vals = []
        for _ in range(3000):
            gsum = rng.choice((-1.0, 1.0), size=(t, d)).sum(0) / math.sqrt(d)
            vals.append(np.linalg.norm(gsum))
        mean = float(np.mean(vals))
        out[f"d{d}_T{t}"] = {
            "mean_regret": mean,
            "sqrt_dT": math.sqrt(d * t),
            "normalized": mean / math.sqrt(t),
        }
    return {
        "check": "Rademacher lower-bound scaling",
        "results": out,
        "note": "With ||g_t||=1, regret scales as sqrt(T); dimension enters through the hard partial-feedback construction, not this full-information sanity check.",
    }


def renyi() -> dict:
    # Data processing for discrete Rényi divergence and additive Gaussian RDP.
    alpha = 2.0
    p = np.array([0.40, 0.35, 0.25])
    q = np.array([0.20, 0.45, 0.35])
    d_before = math.log(float(np.sum(p**alpha * q ** (1 - alpha)))) / (alpha - 1)
    pp = np.array([p[0] + p[1], p[2]])
    qq = np.array([q[0] + q[1], q[2]])
    d_after = math.log(float(np.sum(pp**alpha * qq ** (1 - alpha)))) / (alpha - 1)
    sensitivity = 1.0
    sigmas = np.array([1.2, 2.0, 3.0])
    eps = alpha * sensitivity**2 / (2 * sigmas**2)
    return {
        "alpha": alpha,
        "renyi_before_postprocess": d_before,
        "renyi_after_postprocess": d_after,
        "data_processing_holds": d_after <= d_before + 1e-12,
        "adaptive_composition_sum": float(eps.sum()),
        "per_round_eps": eps.tolist(),
        "precision_sum": float(np.sum(1 / sigmas**2)),
    }


def universality() -> dict:
    # A slow-decaying inner map makes ordinary sums diverge, whereas a
    # p-summable weighted version converges.
    ns = np.array([10, 100, 1000, 10000])
    harmonic = np.array([np.sum(1 / np.arange(1, n + 1)) for n in ns])
    weighted = np.array([np.sum(1 / np.arange(1, n + 1) ** 2) for n in ns])
    # Duplication sensitivity of an unnormalised sum vs a mean.
    x = np.array([1.0, 2.0, 4.0])
    duplicated = np.repeat(x, 3)
    return {
        "n": ns.tolist(),
        "ordinary_partial_sums": harmonic.tolist(),
        "weighted_partial_sums": weighted.tolist(),
        "weighted_limit_pi2_over6": math.pi**2 / 6,
        "ordinary_growth_last_minus_first": float(harmonic[-1] - harmonic[0]),
        "weighted_tail_last_minus_prev": float(weighted[-1] - weighted[-2]),
        "sum_duplication_factor": float(duplicated.sum() / x.sum()),
        "mean_duplication_error": float(abs(duplicated.mean() - x.mean())),
    }


def causal_selection() -> dict:
    # Collider/selection bias check: X and Y are independent marginally, but
    # selection on X+Y+noise induces conditional dependence.
    rng = np.random.default_rng(2)
    n = 500_000
    x = rng.normal(size=n)
    y = rng.normal(size=n)
    s = x + y + 0.35 * rng.normal(size=n) > 1.0
    corr_all = float(np.corrcoef(x, y)[0, 1])
    corr_sel = float(np.corrcoef(x[s], y[s])[0, 1])
    # A second selection generation amplifies the dependency.
    s2 = s & (0.8 * x + 0.8 * y + rng.normal(size=n) > 0.5)
    corr_sel2 = float(np.corrcoef(x[s2], y[s2])[0, 1])
    return {
        "n": n,
        "selected_fraction_generation1": float(s.mean()),
        "selected_fraction_generation2": float(s2.mean()),
        "corr_unselected_population": corr_all,
        "corr_after_one_selection": corr_sel,
        "corr_after_two_selections": corr_sel2,
    }


def cgp() -> dict:
    # Certified 1-D Lipschitz pruning with exact observations. The upper
    # envelope U(x)=min_i y_i+L|x-x_i| can only eliminate x if U(x)<best.
    grid = np.linspace(0, 1, 5001)
    f = 1 - np.abs(grid - 0.37)
    L = 1.0
    sample_counts = [2, 4, 8, 16, 32, 64]
    rows = []
    for m in sample_counts:
        idx = np.linspace(0, len(grid) - 1, m, dtype=int)
        ys = f[idx]
        upper = np.min(ys[:, None] + L * np.abs(grid[None, :] - grid[idx, None]), axis=0)
        best = float(ys.max())
        active = upper >= best - 1e-12
        xstar_idx = int(np.argmax(f))
        rows.append(
            {
                "samples": m,
                "active_fraction": float(active.mean()),
                "optimizer_retained": bool(active[xstar_idx]),
                "best_gap": float(f.max() - best),
            }
        )
    return {"rows": rows, "all_optimizer_retained": all(r["optimizer_retained"] for r in rows)}


def conformal() -> dict:
    # Split-conformal threshold on nested node sets K_tau={v: score(v)>=tau}.
    rng = np.random.default_rng(3)
    nodes = 36
    n_cal = 50
    n_test = 20_000
    phi = 0.75
    # Each route is a random contiguous interval on a path graph; node score is
    # peaked at the centre and the nonconformity is the minimum score on route.
    score = 1 - np.abs(np.arange(nodes) - (nodes - 1) / 2) / ((nodes - 1) / 2)

    def sample_route():
        a = int(rng.integers(0, nodes - 1))
        b = int(rng.integers(a + 1, nodes))
        return np.arange(a, b + 1)

    cal = np.array([float(score[sample_route()].min()) for _ in range(n_cal)])
    # Conservative split-conformal order statistic for P(route subset K)>=phi.
    rank = max(1, math.floor((n_cal + 1) * (1 - phi)))
    tau = float(np.sort(cal)[rank - 1])
    covered = np.mean([score[sample_route()].min() >= tau for _ in range(n_test)])
    sets = [set(np.where(score >= t)[0]) for t in np.linspace(0, 1, 21)]
    nested = all(sets[i + 1].issubset(sets[i]) for i in range(len(sets) - 1))
    return {
        "calibration_routes": n_cal,
        "test_routes": n_test,
        "target_phi": phi,
        "threshold": tau,
        "empirical_coverage": float(covered),
        "nested_threshold_sets": nested,
        "selected_nodes": int(np.sum(score >= tau)),
    }


def forest() -> dict:
    # Forest matrix diagonal and stochastic diagonal estimation. A simple
    # Hutchinson audit checks O(l n) accumulation and 1/sqrt(l) error decay;
    # this is not the paper's GSCF sampler.
    rng = np.random.default_rng(4)
    n = 180
    a = np.zeros((n, n))
    for i in range(n - 1):
        w = 0.5 + rng.random()
        sign = rng.choice((-1.0, 1.0))
        a[i, i + 1] = a[i + 1, i] = sign * w
    degree = np.diag(np.sum(np.abs(a), axis=1))
    lap = degree - a
    q = np.linalg.inv(np.eye(n) + lap)
    exact = np.diag(q)
    rows = []
    for l in [20, 100, 500, 2000]:
        tic = time.perf_counter()
        est = np.zeros(n)
        for _ in range(l):
            z = rng.choice((-1.0, 1.0), size=n)
            est += z * (q @ z)
        est /= l
        rows.append(
            {
                "samples": l,
                "mean_relative_error": float(np.mean(np.abs(est - exact) / np.maximum(exact, 1e-12))),
                "seconds": time.perf_counter() - tic,
            }
        )
    return {"n": n, "rows": rows, "error_reduced": rows[-1]["mean_relative_error"] < rows[0]["mean_relative_error"]}


def main() -> None:
    payload = {
        "scope": "independent lightweight numerical audits; not proof replacements or full benchmark replications",
        "bandit": bandit(),
        "renyi": renyi(),
        "universality": universality(),
        "causal_selection": causal_selection(),
        "cgp": cgp(),
        "conformal": conformal(),
        "forest": forest(),
    }
    print(json.dumps(payload, indent=2))


if __name__ == "__main__":
    main()