scripts/claim5_thm311_median3.py

"""
CLAIM 5 -- Theorem 3.11:  the pointwise median of three interpolators trained on
three independent samples of size n achieves E[L^gamma_D(M)] = O(d_gamma/n),
hence sample complexity O(d_gamma/eps), matching the Omega(d_gamma/eps) lower
bounds of Theorems 3.5 and 3.8.

Independent method
------------------
(a) Re-derive the key structural lemma executably: the median of three values is
    gamma-far from y only if at least two of them are (random + adversarial
    search over triples).
(b) Worst-case-over-distributions evaluation.  For the hard family of
    Theorem 3.5 with gamma-graph dimension d and a worst-case interpolator, the
    error probability of A(S) at a point of mass p is exactly (1-p)^n, so the
    median-of-three loss is  sum_i p_i * (3 q_i^2 - 2 q_i^3),  q_i = (1-p_i)^n.
    We MAXIMISE this over mass profiles subject to only d points being available
    -- i.e. we compute the worst case the theorem must dominate.
(c) Fit the empirical scaling exponents: log L vs log n (predicted -1) and
    log L vs log d_gamma (predicted +1), and check that n*L/d_gamma stays
    bounded (no hidden log n factor).
(d) Full Monte-Carlo with real sampling, real interpolators and a real median.
(e) Compare against a single interpolator, mean-of-three, and a proper learner
    on the Theorem 3.12 instance; and probe the boundary d_gamma = infinity.

Seeds: numpy default_rng(20260725 + offset).  CPU only.
"""

import math
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from core import (  # noqa: E402
    cutoff_loss,
    dump_json,
    fit_loglog_slope,
    is_gamma_graph_shattered,
    thm35_class,
)

OUT = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "outputs"
)
GAMMA = 0.1
SEED = 20260725
res = {"gamma": GAMMA, "seed": SEED, "theorem": "3.11"}

# ---------------------------------------------------------------------------
# (a) the structural lemma: median errs => at least 2 of 3 err
# ---------------------------------------------------------------------------
print("== structural lemma: median of 3 is gamma-far only if >= 2 of 3 are ==")
rng = np.random.default_rng(SEED)
N = 4_000_000
z = rng.random((N, 3))
y = rng.random(N)
far = np.abs(z - y[:, None]) > GAMMA
med = np.median(z, axis=1)
med_far = np.abs(med - y) > GAMMA
violations = int(np.sum(med_far & (far.sum(axis=1) < 2)))
res["lemma_random_trials"] = N
res["lemma_violations"] = violations
# adversarial grid search too
gz = np.linspace(0, 1, 41)
G = np.array(np.meshgrid(gz, gz, gz)).reshape(3, -1).T
adv_viol = 0
for yy in np.linspace(0, 1, 41):
    f = (np.abs(G - yy) > GAMMA).sum(axis=1)
    mf = np.abs(np.median(G, axis=1) - yy) > GAMMA
    adv_viol += int(np.sum(mf & (f < 2)))
res["lemma_grid_violations"] = adv_viol
res["lemma_holds"] = bool(violations == 0 and adv_viol == 0)
print(
    f"  {N:,} random triples + 41^4 grid triples: {violations} + {adv_viol} violations "
    f"-> lemma holds: {res['lemma_holds']}"
)


# ---------------------------------------------------------------------------
# (b) worst-case median-of-three loss on the d_gamma = d hard family
# ---------------------------------------------------------------------------
def median3_loss(p: np.ndarray, n: int) -> float:
    q = (1.0 - p) ** n
    return float(np.sum(p * (3 * q**2 - 2 * q**3)))


def single_interp_loss(p: np.ndarray, n: int) -> float:
    return float(np.sum(p * (1.0 - p) ** n))


def worst_case_profile(d: int, n: int):
    """Maximise the median-of-three loss over mass profiles supported on the
    gamma-graph-shattered set of size d (plus a heavy 'free' point that soaks up
    the leftover mass and is never mispredicted)."""
    best = (0.0, None)
    for u in np.geomspace(1e-6, 1.0 / d, 4000):
        p = np.full(d, u)
        val = median3_loss(p, n)
        if val > best[0]:
            best = (val, float(u))
    return best


print("\n== worst-case median-of-three loss vs n and d_gamma ==")
table = []
for d in (2, 4, 8, 16, 32, 64, 128):
    for n in (d, 2 * d, 5 * d, 10 * d, 50 * d, 100 * d, 500 * d, 1000 * d):
        L, u = worst_case_profile(d, n)
        table.append(
            {
                "d_gamma": d,
                "n": n,
                "n_over_d": n / d,
                "worst_E_loss": L,
                "n_L_over_d": n * L / d,
                "argmax_point_mass": u,
            }
        )
res["worst_case_table"] = table
ratios = [t["n_L_over_d"] for t in table]
res["n_L_over_d_max"] = float(max(ratios))
res["n_L_over_d_min"] = float(min(ratios))
print(
    f"  n*L/d_gamma over the whole grid: min={min(ratios):.4f}  max={max(ratios):.4f}  "
    f"(bounded constant => L = O(d_gamma/n))"
)
for d in (8, 64):
    sub = [t for t in table if t["d_gamma"] == d]
    print(f"  d_gamma={d}:")
    for t in sub:
        print(
            f"    n={t['n']:7d} (={t['n_over_d']:6.0f} d)  L={t['worst_E_loss']:.6e}  n*L/d={t['n_L_over_d']:.4f}"
        )

# ---------------------------------------------------------------------------
# (c) scaling exponents
# ---------------------------------------------------------------------------
print("\n== fitted scaling exponents ==")
fits = {}
for d in (8, 32, 128):
    sub = [t for t in table if t["d_gamma"] == d]
    slope, r2 = fit_loglog_slope(
        [t["n"] for t in sub], [t["worst_E_loss"] for t in sub]
    )
    fits[f"log_L_vs_log_n_at_d{d}"] = {"slope": slope, "r2": r2}
    print(
        f"  d_gamma={d:4d}: d log L / d log n = {slope:+.4f}  (predicted -1)   R^2={r2:.5f}"
    )
# d_gamma-dependence must be measured at FIXED n (varying n with d would hold
# d/n constant and reveal nothing)
d_slopes = []
for n_fixed in (5_000, 20_000, 100_000):
    ds = [2, 4, 8, 16, 32, 64, 128]
    Ls = [worst_case_profile(dd, n_fixed)[0] for dd in ds]
    slope, r2 = fit_loglog_slope(ds, Ls)
    d_slopes.append(slope)
    fits[f"log_L_vs_log_d_at_n_{n_fixed}"] = {"slope": slope, "r2": r2}
    print(
        f"  n={n_fixed:7d} fixed:  d log L / d log d_gamma = {slope:+.4f}  (predicted +1)   R^2={r2:.5f}"
    )
res["fits"] = fits
res["d_exponent_mean"] = float(np.mean(d_slopes))

# sample complexity: n(eps) should be C * d/eps with C constant
print("\n== sample complexity n(eps) of median-of-three ==")
sc = []
for d in (8, 32, 128):
    for eps in (0.1, 0.05, 0.02, 0.01, 0.005):
        lo, hi = 1, 1 << 24
        while lo < hi:
            mid = (lo + hi) // 2
            if worst_case_profile(d, mid)[0] <= eps:
                hi = mid
            else:
                lo = mid + 1
        sc.append(
            {
                "d_gamma": d,
                "eps": eps,
                "n_eps": lo,
                "n_eps_over_d_over_eps": lo / (d / eps),
            }
        )
        print(
            f"  d_gamma={d:4d} eps={eps:<6} n(eps)={lo:8d}   n(eps)/(d/eps)={lo/(d/eps):.4f}"
        )
res["sample_complexity"] = sc
c = [s["n_eps_over_d_over_eps"] for s in sc]
res["sample_complexity_constant_mean"] = float(np.mean(c))
res["sample_complexity_constant_cv"] = float(np.std(c) / np.mean(c))
print(
    f"  constant C = n(eps)/(d_gamma/eps): mean={np.mean(c):.4f}  CV={np.std(c)/np.mean(c):.4f} "
    f"(constant => O(d_gamma/eps), matching the Omega(d_gamma/eps) lower bound)"
)


# ---------------------------------------------------------------------------
# (d) full Monte-Carlo with real samples, real interpolators, real median
# ---------------------------------------------------------------------------
def mc_median3(d, n, u, trials, seed, materialise=False):
    """Actual simulation: three independent samples, the worst-case interpolator
    on each, pointwise median, exact cutoff loss under D.

    The worst-case interpolator returns the member of H that is 0 on the
    observed points and gamma-far elsewhere, so its prediction vector is simply
    0 on observed / hi on unobserved.  For small d we ALSO materialise H and
    look the hypothesis up by its bit pattern, checking the shortcut agrees and
    that the point set really is gamma-graph shattered (Definition 3.1)."""
    hi = min(1.0, GAMMA + (1.0 - GAMMA) / 2.0)
    lookup = None
    if materialise:
        cls = thm35_class(d, GAMMA)
        witness = int(np.argmin(cls.values.sum(axis=1)))
        assert is_gamma_graph_shattered(cls, range(d), GAMMA, witness)
        codes = (cls.values > 0).astype(int) @ (1 << np.arange(d))
        lookup = {int(c): i for i, c in enumerate(codes)}
    p = np.zeros(d + 1)
    p[:d] = u
    p[d] = 1.0 - d * u  # heavy point outside the shattered set, never mispredicted
    rng = np.random.default_rng(seed)
    labels = np.zeros(d)
    masses = np.full(d, u)
    out = np.empty(trials)
    single = np.empty(trials)
    for t in range(trials):
        preds = []
        for _ in range(3):
            draws = rng.choice(d + 1, size=n, p=p)
            mask = np.zeros(d, bool)
            mask[draws[draws < d]] = True
            pr = np.where(mask, 0.0, hi)
            if lookup is not None:
                b = (~mask).astype(int)
                assert np.allclose(pr, cls.values[lookup[int(b @ (1 << np.arange(d)))]])
            preds.append(pr)
        P = np.stack(preds)
        out[t] = cutoff_loss(np.median(P, axis=0), labels, masses, GAMMA)
        single[t] = cutoff_loss(P[0], labels, masses, GAMMA)
    return (
        float(out.mean()),
        float(out.std(ddof=1) / math.sqrt(trials)),
        float(single.mean()),
    )


print("\n== Monte-Carlo cross-check (real sampling / real median) ==")
mc = []
for d, n in ((8, 40), (8, 200), (12, 120), (32, 160), (32, 800), (128, 640)):
    L_pred, u = worst_case_profile(d, n)
    m_med, se, m_single = mc_median3(d, n, u, 4000, SEED + n, materialise=(d <= 12))
    mc.append(
        {
            "d_gamma": d,
            "n": n,
            "point_mass_u": u,
            "analytic_median3": L_pred,
            "mc_median3": m_med,
            "mc_stderr": se,
            "mc_single_interpolator": m_single,
            "rel_err": abs(m_med - L_pred) / L_pred,
        }
    )
    print(
        f"  d={d:4d} n={n:5d}: analytic={L_pred:.6f}  MC={m_med:.6f} +/- {2*se:.6f}  "
        f"(rel.err {abs(m_med-L_pred)/L_pred*100:.2f}%)   single interpolator MC={m_single:.6f}"
    )
res["monte_carlo"] = mc
res["mc_max_rel_err"] = float(max(r["rel_err"] for r in mc))

# ---------------------------------------------------------------------------
# (e) comparison + boundary probe
# ---------------------------------------------------------------------------
print("\n== comparison and boundary probes ==")
cmp_rows = []
for d in (8, 32, 128):
    n = 100 * d
    Lm, u = worst_case_profile(d, n)
    p = np.full(d, u)
    cmp_rows.append(
        {
            "d_gamma": d,
            "n": n,
            "median_of_3": Lm,
            "single_interpolator_same_profile": single_interp_loss(p, n),
            "single_interpolator_worst_profile": max(
                single_interp_loss(np.full(d, uu), n)
                for uu in np.geomspace(1e-7, 1.0 / d, 3000)
            ),
        }
    )
    r = cmp_rows[-1]
    print(
        f"  d={d:4d} n={n:6d}: median-of-3 = {r['median_of_3']:.3e}   "
        f"single interpolator (its own worst profile) = {r['single_interpolator_worst_profile']:.3e}   "
        f"ratio = {r['single_interpolator_worst_profile']/r['median_of_3']:.2f}x"
    )
res["comparison"] = cmp_rows

# boundary: d_gamma = infinity (Theorem 3.10 class) -- the bound must go vacuous
print("\n  boundary probe: d_gamma = infinity (Theorem 3.10 class)")
rng2 = np.random.default_rng(SEED + 31)
s = 2000
ku = s * s
n = 500
losses = []
for _ in range(200):
    A = rng2.choice(ku, size=s, replace=False)
    Zs = []
    for _ in range(3):
        obs = np.unique(A[rng2.integers(0, s, size=n)])
        fill = rng2.integers(0, ku, size=s - len(obs))
        Zs.append(np.union1d(obs, fill))
    # median of three: a point is mispredicted unless >= 2 of the 3 zero it out
    hits = sum(np.isin(A, Z).astype(int) for Z in Zs)
    losses.append(float(np.mean(hits < 2)))
res["boundary_infinite_dgamma_median3_loss"] = float(np.mean(losses))
print(
    f"    median-of-three on the d_gamma = infinity class (n={n}, ku={ku}): "
    f"E[L] = {np.mean(losses):.5f}  -> the O(d_gamma/n) bound is vacuous there, as it must be"
)

ok = (
    res["lemma_holds"]
    and res["n_L_over_d_max"] < 1.0
    and abs(fits["log_L_vs_log_n_at_d32"]["slope"] + 1) < 0.05
    and res["sample_complexity_constant_cv"] < 0.1
    and res["mc_max_rel_err"] < 0.08
    and abs(res["d_exponent_mean"] - 1.0) < 0.05
)
res["verdict"] = "verified" if ok else "partial"
print(f"\nverdict = {res['verdict']}")
dump_json(os.path.join(OUT, "claim5_thm311.json"), res)