scripts/claim5_covering.py

"""
CLAIM 5 -- Lemma 5.3 (covering number) + Lemma 5.4 (Rademacher complexity):
"The covering number of the dual multiplier class U is bounded as
 log N(delta,U,||.||_2,N) <= s log(1 + 2 B pi_max s / delta), used to derive the
 Rademacher complexity bound R_N(U) = O(s^1.5/sqrt(N))."

Independent re-derivation, executed:

(1) LIPSCHITZ TRANSFER: measure the true Lipschitz modulus of pi -> u(pi,.) in the empirical
    L2 metric and compare with the paper's L = 2 B sqrt(s) (Corollary 5.2).  We also compute
    the EXACT per-sample Lipschitz constant L_exact = max_{i,j} ||b_i - A_i x_j||_2 (u is a
    min of affine functions, so this is exact, not a bound).
(2) EXPLICIT delta-NET CONSTRUCTION of U for s = 1,2,3: build the grid net of Pi with
    ell_2 radius delta/L that Lemma 5.3's proof prescribes, MEASURE that it really is a
    delta-covering of U in the empirical L2 metric on N instances, and compare
    log(net size) against the paper's bound.
(3) greedy 2delta-PACKING of U -> a lower estimate of the true log covering number, i.e.
    how much slack Lemma 5.3 carries.
(4) DUDLEY CHAIN re-derivation: verify the footnote identity int_0^R sqrt(log(R/d)) dd
    = R sqrt(pi)/2, then numerically integrate the entropy integral with the paper's
    covering bound and compare to their closed form 3 sqrt(pi) B pi_max s^1.5/sqrt(N).
(5) MEASURED Rademacher complexity R_N(U): a 200k-point grid bracket for s=1, a
    rigorous two-sided bracket for s=2 (fine net + exact Lipschitz slack), and a multi-start
    lower estimate for s up to 16.  N- and s-exponents fitted; Lemma 5.4 checked for
    violation.
"""

import sys, os, time, itertools

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import numpy as np
from common import gen_instance, Pool, loglog_fit, dump_json

OUT = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "outputs"
)
os.makedirs(OUT, exist_ok=True)
SEED = 5052026
PIMAX = 1.0
t0 = time.time()
rng = np.random.default_rng(SEED)
res = {"seed": SEED, "pimax": PIMAX}


def make_pool(s, M):
    lst = []
    r = np.random.default_rng(SEED + 1000 * s)
    while len(lst) < M:
        it = gen_instance(r, s=s, m=2, p=5)
        if it is not None:
            lst.append(it)
    return Pool(lst, PIMAX)


def u_matrix(pool, pts, n_inst=None, chunk=400):
    """(n_pts, n_inst) matrix of u(pi, P_i) computed in chunks over pts."""
    n = pool.M if n_inst is None else n_inst
    a = pool.a[:n]  # (n,K)
    G = pool.G[:n]  # (n,K,s)
    out = np.empty((pts.shape[0], n))
    for st in range(0, pts.shape[0], chunk):
        blk = pts[st : st + chunk]  # (c,s)
        vals = a[:, :, None] + G @ blk.T  # (n,K,c)
        out[st : st + chunk] = vals.min(axis=1).T
    return out


def exact_lipschitz(pool, n_inst=None):
    n = pool.M if n_inst is None else n_inst
    finite = pool.a[:n] < 1e17
    nr = np.linalg.norm(pool.G[:n], axis=2)
    return float(nr[finite].max())


# =============================================================== (1) Lipschitz transfer
lip = {}
for s in [1, 2, 4, 8, 16]:
    pool = make_pool(s, 200)
    L = 2 * pool.B * np.sqrt(s)
    Lex = exact_lipschitz(pool)
    P1 = rng.uniform(0, PIMAX, (600, s))
    P2 = rng.uniform(0, PIMAX, (600, s))
    U1 = u_matrix(pool, P1)
    U2 = u_matrix(pool, P2)
    d = np.sqrt(((U1 - U2) ** 2).mean(axis=1))
    nrm = np.linalg.norm(P1 - P2, axis=1)
    worst = float(np.max(d / np.maximum(nrm, 1e-15)))
    lip[s] = {
        "B": pool.B,
        "L_paper_2Bsqrts": float(L),
        "L_exact_max_subgrad_norm": Lex,
        "max_measured_secant_slope": worst,
        "ratio_measured_over_L_paper": float(worst / L),
        "ratio_Lexact_over_Lpaper": float(Lex / L),
    }
    print(
        "(1) s=%2d  secant<=%.4f   L_exact=%.4f   L_paper=%.4f   (paper/exact = %.2fx)"
        % (s, worst, Lex, L, L / Lex)
    )
res["lipschitz_transfer"] = lip
res["lipschitz_verdict"] = (
    "Corollary 5.2's L = 2B sqrt(s) is a valid Lipschitz constant in every case; it "
    "over-estimates the exact constant max_{i,j}||b-Ax_j||_2 by roughly %.0fx"
    % np.mean(
        [v["L_paper_2Bsqrts"] / v["L_exact_max_subgrad_norm"] for v in lip.values()]
    )
)


# =============================================================== (2) explicit delta-net
def grid_net(s, h, pimax=PIMAX):
    ng = int(np.ceil(pimax / h))
    ax = np.clip((np.arange(ng) + 0.5) * h, 0, pimax)
    return np.array(list(itertools.product(ax, repeat=s))), ng


def min_dist_to_net(Uprobe, Unet):
    """min_j ||f_i - g_j||_{2,N} via the Gram trick.  U* are (n, Nsamp)."""
    Ns = Uprobe.shape[1]
    p2 = (Uprobe**2).sum(axis=1)
    n2 = (Unet**2).sum(axis=1)
    best = np.full(Uprobe.shape[0], np.inf)
    for st in range(0, Unet.shape[0], 20000):
        blk = Unet[st : st + 20000]
        d2 = p2[:, None] - 2 * (Uprobe @ blk.T) + n2[None, st : st + 20000]
        best = np.minimum(best, d2.min(axis=1))
    return np.sqrt(np.maximum(best, 0) / Ns)


net_res = []
NPOOL = 150
probes = {}
for s in [1, 2, 3]:
    pool = make_pool(s, NPOOL)
    L = 2 * pool.B * np.sqrt(s)
    pr = rng.uniform(0, PIMAX, size=(3000, s))
    Uprobe = u_matrix(pool, pr)
    for delta in [0.5, 0.25, 0.1]:
        rho = delta / L
        h = 2 * rho / np.sqrt(s)
        net, ng = grid_net(s, h)
        if net.shape[0] > 600000:
            print(
                "(2) s=%d delta=%g skipped (net of %d points)"
                % (s, delta, net.shape[0])
            )
            continue
        Unet = u_matrix(pool, net)
        mind = min_dist_to_net(Uprobe, Unet)
        bound = s * np.log(1 + 2 * pool.B * PIMAX * s / delta)
        row = {
            "s": s,
            "delta": delta,
            "B": pool.B,
            "n_net": int(net.shape[0]),
            "log_net_size": float(np.log(net.shape[0])),
            "lemma53_bound": float(bound),
            "max_probe_distance_to_net": float(mind.max()),
            "is_delta_covering": bool(mind.max() <= delta + 1e-9),
            "construction_within_bound": bool(np.log(net.shape[0]) <= bound + 1e-12),
        }
        net_res.append(row)
        print(
            "(2) s=%d delta=%.2f  |net|=%d  log|net|=%.3f <= bound %.3f : %s ; "
            "max dist %.4f <= delta : %s"
            % (
                s,
                delta,
                row["n_net"],
                row["log_net_size"],
                row["lemma53_bound"],
                row["construction_within_bound"],
                row["max_probe_distance_to_net"],
                row["is_delta_covering"],
            )
        )
res["explicit_nets"] = net_res
res["net_verdict"] = (
    "every explicitly constructed net is a genuine delta-covering of U in the empirical L2 "
    "metric AND its log-size obeys Lemma 5.3"
    if all(r["is_delta_covering"] and r["construction_within_bound"] for r in net_res)
    else "VIOLATED"
)

# =============================================================== (3) packing lower estimate
pack_res = []
for s in [1, 2, 3]:
    pool = make_pool(s, NPOOL)
    cand = rng.uniform(0, PIMAX, size=(4000, s))
    Uc = u_matrix(pool, cand)
    Ns = Uc.shape[1]
    for delta in [0.5, 0.25, 0.1]:
        kept = []
        Kmat = None
        for i in range(Uc.shape[0]):
            if kept:
                d = np.sqrt(((Uc[i][None, :] - Uc[kept]) ** 2).mean(axis=1))
                if d.min() <= 2 * delta:
                    continue
            kept.append(i)
        bound = s * np.log(1 + 2 * pool.B * PIMAX * s / delta)
        pack_res.append(
            {
                "s": s,
                "delta": delta,
                "greedy_2delta_packing_size": len(kept),
                "log_packing_lower_est": float(np.log(len(kept))),
                "lemma53_upper_bound": float(bound),
                "slack_nats": float(bound - np.log(len(kept))),
            }
        )
        print("(3)", pack_res[-1])
res["packing_lower_estimates"] = pack_res

# =============================================================== (4) Dudley chain
R = 3.0
xs = np.linspace(1e-12, R, 4000001)
num = float(np.trapezoid(np.sqrt(np.log(R / xs)), xs))
res["dudley_identity"] = {
    "R": R,
    "numeric_integral": num,
    "closed_form_R_sqrt_pi_over_2": float(R * np.sqrt(np.pi) / 2),
    "rel_err": float(abs(num - R * np.sqrt(np.pi) / 2) / (R * np.sqrt(np.pi) / 2)),
}
print("(4a)", res["dudley_identity"])

dud = []
for s in [1, 2, 4, 8, 16, 32]:
    Bv, L, D = 1.0, 2.0 * np.sqrt(s), PIMAX * np.sqrt(s)
    for N in [64, 256, 1024]:
        dd = np.linspace(1e-10, L * D, 800001)
        integ = float(
            np.trapezoid(np.sqrt(s * np.log(1 + 2 * Bv * PIMAX * s / dd) / N), dd)
        )
        closed = 3 * np.sqrt(np.pi) * Bv * PIMAX * s**1.5 / np.sqrt(N)
        dud.append(
            {
                "s": s,
                "N": N,
                "numeric_entropy_integral": integ,
                "paper_closed_form": float(closed),
                "numeric_over_closed": integ / closed,
            }
        )
    print(
        "(4b) s=%2d" % s,
        {k: (round(v, 5) if isinstance(v, float) else v) for k, v in dud[-1].items()},
    )
res["dudley_chain"] = dud
res["dudley_verdict"] = (
    "the paper's closed form 3 sqrt(pi) B pi_max s^1.5/sqrt(N) dominates the exact entropy "
    "integral for every (s,N) tested"
    if all(d["numeric_over_closed"] <= 1.0 for d in dud)
    else "closed form is NOT an upper bound"
)
res["dudley_exponents"] = {
    "s_exponent_at_N1024": loglog_fit(
        [d["s"] for d in dud if d["N"] == 1024],
        [d["numeric_entropy_integral"] for d in dud if d["N"] == 1024],
    ),
    "N_exponent_at_s16": loglog_fit(
        [d["N"] for d in dud if d["s"] == 16],
        [d["numeric_entropy_integral"] for d in dud if d["s"] == 16],
    ),
}
print(
    "(4c) exact entropy-integral exponents (paper: s^1.5, N^-0.5):",
    res["dudley_exponents"],
)

# =============================================================== (5) measured Rademacher
NSIG = 600
NLIST = [25, 50, 100, 200, 400]
rad = []

# ---- s = 1 : essentially exact -- 200k-point grid, exact Lipschitz slack ~ 1e-5
pool1 = make_pool(1, max(NLIST))
brk = np.linspace(0.0, PIMAX, 200001)[:, None]
U1 = u_matrix(pool1, brk)  # (ngrid, M)
L1 = exact_lipschitz(pool1)
slack1 = L1 * (PIMAX / 200000) / 2
for N in NLIST:
    sig = rng.choice([-1.0, 1.0], size=(NSIG, N))
    vals = (sig @ U1[:, :N].T) / N
    r = float(vals.max(axis=1).mean())
    bound = 3 * np.sqrt(np.pi) * pool1.B * PIMAX * 1**1.5 / np.sqrt(N)
    rad.append(
        {
            "s": 1,
            "N": N,
            "R_hat": r,
            "R_hat_lo": r,
            "R_hat_hi": r + slack1,
            "method": "fine-grid-bracket",
            "n_eval_points": int(brk.shape[0]),
            "lipschitz_slack": float(slack1),
            "lemma54_bound": float(bound),
            "bound_respected": bool(r + slack1 <= bound),
        }
    )
    print(
        "(5) s=1 N=%3d  R_hat in [%.6f, %.6f] (%d grid points)  Lemma5.4 bound=%.5f  ok=%s"
        % (N, r, r + slack1, brk.shape[0], bound, rad[-1]["bound_respected"])
    )

# ---- s = 2 : rigorous bracket with a fine grid + exact Lipschitz slack
pool2 = make_pool(2, max(NLIST))
ng = 400
net2, _ = grid_net(2, PIMAX / ng)
L2 = exact_lipschitz(pool2)
slack2 = L2 * (PIMAX / ng) * np.sqrt(2) / 2
U2 = u_matrix(pool2, net2)
for N in NLIST:
    sig = rng.choice([-1.0, 1.0], size=(NSIG, N))
    vals = (sig @ U2[:, :N].T) / N
    r = float(vals.max(axis=1).mean())
    bound = 3 * np.sqrt(np.pi) * pool2.B * PIMAX * 2**1.5 / np.sqrt(N)
    rad.append(
        {
            "s": 2,
            "N": N,
            "R_hat": r,
            "R_hat_lo": r,
            "R_hat_hi": r + slack2,
            "method": "grid-bracket",
            "n_eval_points": int(net2.shape[0]),
            "lipschitz_slack": float(slack2),
            "lemma54_bound": float(bound),
            "bound_respected": bool(r + slack2 <= bound),
        }
    )
    print(
        "(5) s=2 N=%3d  R_hat in [%.5f, %.5f]  Lemma5.4 bound=%.5f  ok=%s"
        % (N, r, r + slack2, bound, rad[-1]["bound_respected"])
    )
res["rademacher"] = rad
res["rademacher_N_exponents"] = {
    str(s): loglog_fit(
        [d["N"] for d in rad if d["s"] == s], [d["R_hat"] for d in rad if d["s"] == s]
    )
    for s in [1, 2]
}
print(
    "(5b) N exponents of measured R_hat (predicted -0.5):",
    res["rademacher_N_exponents"],
)


# ---- s-scaling: multi-start lower estimate
def rademacher_multistart(pool, N, rng, nsig=60, nrestart=200, niter=50):
    s = pool.s
    idx = np.arange(N)
    tot = 0.0
    for _ in range(nsig):
        sig = rng.choice([-1.0, 1.0], size=N)
        starts = rng.uniform(0, pool.pimax, size=(nrestart, s))
        Ustart = u_matrix(pool, starts, n_inst=N)
        best = float((Ustart @ sig / N).max())
        pi = starts[int(np.argmax(Ustart @ sig / N))].copy()
        step = pool.pimax / 4
        cur = best
        for _ in range(niter):
            _, g = pool.grad_all(pi, idx)
            d = (sig[:, None] * g).sum(axis=0) / N
            nd = np.linalg.norm(d)
            if nd < 1e-12:
                break
            cand = np.clip(pi + step * d / nd, 0, pool.pimax)
            v = float(sig @ pool.u_all(cand, idx)[0] / N)
            if v > cur:
                pi, cur = cand, v
            else:
                step *= 0.6
        tot += max(cur, best)
    return tot / nsig


rad_s = []
for s in [1, 2, 4, 8, 16]:
    pool = make_pool(s, 200)
    r = rademacher_multistart(pool, 200, rng)
    bound = 3 * np.sqrt(np.pi) * pool.B * PIMAX * s**1.5 / np.sqrt(200)
    rad_s.append(
        {
            "s": s,
            "N": 200,
            "R_hat_multistart_lower_est": r,
            "lemma54_bound": float(bound),
            "bound_respected": bool(r <= bound),
        }
    )
    print(
        "(5c)",
        {k: (round(v, 5) if isinstance(v, float) else v) for k, v in rad_s[-1].items()},
    )
res["rademacher_s_scaling_lower_est"] = rad_s
res["rademacher_s_exponent"] = loglog_fit(
    [d["s"] for d in rad_s], [d["R_hat_multistart_lower_est"] for d in rad_s]
)
print(
    "(5d) s-exponent of measured R_hat (Lemma 5.4 bound exponent = 1.5):",
    res["rademacher_s_exponent"],
)
res["rademacher_verdict"] = (
    "Lemma 5.4's bound is never violated"
    if all(d["bound_respected"] for d in rad)
    and all(d["bound_respected"] for d in rad_s)
    else "VIOLATED"
)

res["wall_time_s"] = time.time() - t0
dump_json(os.path.join(OUT, "claim5_covering.json"), res)
print("done in", round(time.time() - t0, 1), "s")