code/exp_claim6.py

"""Claim 6 audit — RLS confidence sets in L^2(rho) with log-det width.

Paper Eqs. (4)-(5) [OpenReview Eqs. (11)-(12)] + Lemma B.2:
    C_t(delta) = { f : ||f - f_hat_t||_{V_t^lambda} <= beta_t(delta) },
    beta_t(delta) = sigma sqrt(log(4 det(I + lam^{-1} A_t)/delta^2)) + sqrt(lam) Cbar,
    P( f* in  intersection_t C_t(delta) ) >= 1 - delta/2.

Monte-Carlo audit of UNIFORM (over t<=T) coverage under
  (i)   EntUCB-collected adaptive designs           -> expect >= 1 - delta/2
  (ii)  random-plan designs                          -> expect >= 1 - delta/2
  (iii) CONTROL: width shrunk by 4                   -> coverage collapses
  (iv)  CONTROL: true noise 4x the assumed sigma     -> coverage collapses
Also records beta_t and logdet(I + lam^{-1} A_t) trajectories (the claim that
the width is *controlled by the log-determinant of the design operator*).
"""
import argparse
import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from botlib import make_instance, run_entucb, sinkhorn_log, beta_width

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "results")
os.makedirs(OUT, exist_ok=True)


def random_design_run(inst, T, sigma, delta, lam, Cbar, seed, beta_scale=1.0,
                      sigma_run=None, adaptive=False):
    """Coverage of theta* under random (or estimate-perturbed adaptive) plans."""
    rng = np.random.default_rng(seed)
    sig_run = sigma if sigma_run is None else sigma_run
    N = inst.N
    A = np.zeros((N, N))
    bvec = np.zeros(N)
    Vinv = np.eye(N) / lam
    logdet = 0.0
    covered = True
    min_slack = np.inf
    betas, logdets, ts = [], [], []
    rec = np.unique(np.clip(np.round(np.geomspace(1, T, 40)).astype(int), 1, T))
    theta_hat = np.zeros(N)
    for t in range(1, T + 1):
        if adaptive and t > 5:
            # action depends on past noise through theta_hat -> adaptive design
            cmat = (inst.B.T @ (theta_hat + 0.3 * rng.standard_normal(N))).reshape(inst.K, inst.Kp)
        else:
            cmat = rng.standard_normal((inst.K, inst.Kp))
        P, _, _ = sinkhorn_log(inst.mu, inst.nu, cmat, float(rng.uniform(0.05, 0.5)), n_iter=500, tol=1e-9)
        a = inst.embed(P)
        C = float(inst.theta_star @ a) + sig_run * rng.standard_normal()
        w = Vinv @ a
        Vinv -= np.outer(w, w) / (1.0 + a @ w)
        A += np.outer(a, a)
        bvec += a * C
        theta_hat = Vinv @ bvec
        d = inst.theta_star - theta_hat
        Vt = lam * np.eye(N) + A
        dist = np.sqrt(d @ (Vt @ d))
        sgn, ld = np.linalg.slogdet(np.eye(N) + A / lam)
        beta = beta_scale * beta_width(sigma, delta, lam, Cbar, ld)
        covered = covered and (dist <= beta)
        min_slack = min(min_slack, beta - dist)
        if t in rec:
            ts.append(t); betas.append(beta); logdets.append(ld)
    return covered, min_slack, np.array(ts), np.array(betas), np.array(logdets)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--full", action="store_true")
    ap.add_argument("--tag", default="")
    args = ap.parse_args()

    K = Kp = 5
    T = 400 if args.full else 120
    M_adapt = 200 if args.full else 12
    M_rand = 400 if args.full else 20
    sigma, delta, lam = 0.1, 0.2, 1.0

    res = {"config": dict(K=K, Kp=Kp, T=T, M_adapt=M_adapt, M_rand=M_rand,
                          sigma=sigma, delta=delta, lam=lam)}

    # ---- (i) adaptive EntUCB designs -----------------------------------
    cov = []
    for m in range(M_adapt):
        inst = make_instance(K, Kp, seed=1000 + m % 8, cost_kind="smooth")
        Cbar = 1.1 * np.linalg.norm(inst.theta_star)
        r = run_entucb(inst, T=T, sigma=sigma, delta=delta, lam=lam, Cbar=Cbar,
                       eps_schedule=lambda t: 0.75 * t ** -0.75, seed=20000 + m,
                       opt_iters=4, n_record=8)
        cov.append(r.covered_all)
    res["entucb_uniform_coverage"] = float(np.mean(cov))
    print(f"[claim6] EntUCB adaptive designs: uniform coverage {np.mean(cov):.3f} "
          f"(target >= {1 - delta / 2:.2f}, M={M_adapt})")

    # ---- (ii)-(iv) random / control designs ----------------------------
    variants = {
        "random_correct": dict(beta_scale=1.0, sigma_run=None, adaptive=False),
        "adaptive_perturbed": dict(beta_scale=1.0, sigma_run=None, adaptive=True),
        "control_beta_quarter": dict(beta_scale=0.25, sigma_run=None, adaptive=False),
        "control_noise_4x": dict(beta_scale=1.0, sigma_run=4 * sigma, adaptive=False),
    }
    inst0 = make_instance(K, Kp, seed=1000, cost_kind="smooth")
    Cbar0 = 1.1 * np.linalg.norm(inst0.theta_star)
    for name, kw in variants.items():
        covs, slacks = [], []
        for m in range(M_rand):
            c, s, ts, betas, lds = random_design_run(inst0, T, sigma, delta, lam, Cbar0,
                                                     seed=31000 + m, **kw)
            covs.append(c); slacks.append(s)
        res[name] = {"uniform_coverage": float(np.mean(covs)),
                     "mean_min_slack": float(np.mean(slacks))}
        print(f"[claim6] {name}: uniform coverage {np.mean(covs):.3f}, mean min slack {np.mean(slacks):.3f}")

    # ---- width vs logdet relation (deterministic identity check) --------
    c, s, ts, betas, lds = random_design_run(inst0, T, sigma, delta, lam, Cbar0, seed=99)
    pred = sigma * np.sqrt(lds + np.log(4 / delta ** 2)) + np.sqrt(lam) * Cbar0
    res["width_logdet_identity_maxerr"] = float(np.max(np.abs(betas - pred)))
    res["width_curve"] = {"t": ts.tolist(), "beta": betas.tolist(), "logdet": lds.tolist()}
    print(f"[claim6] beta_t == sigma sqrt(logdet + log(4/d^2)) + sqrt(lam) Cbar: "
          f"max err {res['width_logdet_identity_maxerr']:.2e}")

    tag = args.tag or ("full" if args.full else "smoke")
    with open(os.path.join(OUT, f"claim6_{tag}.json"), "w") as f:
        json.dump(res, f, indent=2)
    print(f"saved results/claim6_{tag}.json")


if __name__ == "__main__":
    main()