code/exp_claim6_gpu.py

# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "scipy", "torch"]
# ///
"""Claim 6 GPU-scaled audit — uniform coverage of the OFUL-style confidence sets
(paper Eqs. (4)-(5) + Lemma B.2) with thousands of Monte-Carlo replications.

Batched over M replications in torch (float64). Designs are sequences of
transport plans produced by batched log-domain Sinkhorn on random / adaptively
chosen (estimate-dependent) costs. Verifies
    P( theta* in C_t(delta) for ALL t <= T ) >= 1 - delta/2
and the controls (width/4, noise 4x) that break the theorem's conditions.

Writes results JSON to --out (mounted bucket path inside a HF Job).
"""
import argparse
import json
import math
import os
import sys
import time

import numpy as np
import torch


# ---- minimal instance construction (self-contained; mirrors botlib) --------

def trig_features_1d(u, n_feats):
    feats = np.empty((n_feats, len(u)))
    freqs = np.empty(n_feats)
    for m in range(n_feats):
        if m == 0:
            feats[m] = 1.0; freqs[m] = 0
        else:
            k = (m + 1) // 2
            ph = 2 * np.pi * k * u
            feats[m] = np.cos(ph) if (m % 2 == 1) else np.sin(ph)
            freqs[m] = k
    return feats, freqs


def build_instance(K, Kp, seed):
    rng = np.random.default_rng(seed)
    x = np.sort(rng.uniform(0.02, 0.98, K)) + rng.normal(0, 1e-3, K)
    y = np.sort(rng.uniform(0.02, 0.98, Kp)) + rng.normal(0, 1e-3, Kp)
    mu = rng.dirichlet(np.full(K, 5.0)); mu = np.maximum(mu, 1e-3); mu /= mu.sum()
    nu = rng.dirichlet(np.full(Kp, 5.0)); nu = np.maximum(nu, 1e-3); nu /= nu.sum()
    rho = np.outer(mu, nu).ravel()
    N = K * Kp
    fx, frx = trig_features_1d(x, 2 * K + 3)
    fy, fry = trig_features_1d(y, 2 * Kp + 3)
    pairs = [(m1, m2) for m1 in range(len(fx)) for m2 in range(len(fy))]
    pairs.sort(key=lambda p: (max(frx[p[0]], fry[p[1]]), frx[p[0]] + fry[p[1]], p[0] + p[1], p))
    cands = np.asarray([np.outer(fx[m1], fy[m2]).ravel() for (m1, m2) in pairs])
    B = []
    for f in cands:
        g = f.astype(float).copy()
        for _ in range(2):
            if B:
                Bm = np.asarray(B)
                g = g - Bm.T @ (Bm @ (rho * g))
        nrm = np.sqrt(g @ (rho * g))
        if nrm > 1e-8:
            B.append(g / nrm)
        if len(B) == N:
            break
    B = np.asarray(B)
    X, Y = np.meshgrid(x, y, indexing="ij")
    cost = np.cos(2 * np.pi * (X - Y)) + 0.5 * np.sin(2 * np.pi * X) * np.cos(4 * np.pi * Y) + 0.3 * X * Y
    theta_star = B @ (rho * cost.ravel())
    return dict(mu=mu, nu=nu, rho=rho, B=B, cost=cost, theta_star=theta_star, K=K, Kp=Kp, N=N)


def batched_sinkhorn(mu, nu, M, eps, n_iter=60):
    """Batched log-domain Sinkhorn. M: (B,K,K'), eps: scalar or (B,1,1)."""
    B = M.shape[0]
    logmu = torch.log(mu)[None, :, None]      # (1,K,1)
    lognu = torch.log(nu)[None, None, :]      # (1,1,K')
    f = torch.zeros(M.shape[0], M.shape[1], 1, dtype=M.dtype, device=M.device)
    g = torch.zeros(M.shape[0], 1, M.shape[2], dtype=M.dtype, device=M.device)
    for _ in range(n_iter):
        f = -eps * torch.logsumexp((g - M) / eps + lognu, dim=2, keepdim=True)
        g = -eps * torch.logsumexp((f - M) / eps + logmu, dim=1, keepdim=True)
    P = torch.exp((f + g - M) / eps + logmu + lognu)
    P = torch.clamp(P, min=0)
    P = P / P.sum(dim=(1, 2), keepdim=True)
    return P


def run_variant(inst, M_reps, T, sigma, delta, lam, Cbar, seed, device,
                beta_scale=1.0, sigma_run=None, adaptive=False, dtype=torch.float64):
    torch.manual_seed(seed)
    K, Kp, N = inst["K"], inst["Kp"], inst["N"]
    sig_run = sigma if sigma_run is None else sigma_run
    mu = torch.tensor(inst["mu"], dtype=dtype, device=device)
    nu = torch.tensor(inst["nu"], dtype=dtype, device=device)
    Bmat = torch.tensor(inst["B"], dtype=dtype, device=device)          # (N,N)
    th_star = torch.tensor(inst["theta_star"], dtype=dtype, device=device)  # (N,)

    A = torch.zeros(M_reps, N, N, dtype=dtype, device=device)
    Vinv = torch.eye(N, dtype=dtype, device=device).expand(M_reps, N, N).clone() / lam
    bvec = torch.zeros(M_reps, N, dtype=dtype, device=device)
    ld = torch.zeros(M_reps, dtype=dtype, device=device)   # logdet(I + lam^-1 A_t)
    covered = torch.ones(M_reps, dtype=torch.bool, device=device)
    min_slack = torch.full((M_reps,), float("inf"), dtype=dtype, device=device)
    theta_hat = torch.zeros(M_reps, N, dtype=dtype, device=device)
    log_const = math.log(4.0 / delta ** 2)
    cov_curve = []

    for t in range(1, T + 1):
        if adaptive and t > 5:
            base = torch.einsum("kn,mk->mn", Bmat, theta_hat)  # (M,N) func values
            cmat = base.reshape(M_reps, K, Kp) + 0.3 * torch.randn(M_reps, K, Kp, dtype=dtype, device=device)
        else:
            cmat = torch.randn(M_reps, K, Kp, dtype=dtype, device=device)
        eps = torch.rand(M_reps, 1, 1, dtype=dtype, device=device) * 0.45 + 0.05
        P = batched_sinkhorn(mu, nu, cmat, eps)
        a = torch.einsum("nz,mz->mn", Bmat, P.reshape(M_reps, N))       # (M,N)
        Cfb = a @ th_star + sig_run * torch.randn(M_reps, dtype=dtype, device=device)

        w = torch.einsum("mij,mj->mi", Vinv, a)
        s = 1.0 + (a * w).sum(1)
        ld = ld + torch.log(s)
        Vinv = Vinv - torch.einsum("mi,mj->mij", w, w) / s[:, None, None]
        A = A + torch.einsum("mi,mj->mij", a, a)
        bvec = bvec + a * Cfb[:, None]
        theta_hat = torch.einsum("mij,mj->mi", Vinv, bvec)

        d = th_star[None, :] - theta_hat
        Vd = lam * d + torch.einsum("mij,mj->mi", A, d)
        dist = torch.sqrt(torch.clamp((d * Vd).sum(1), min=0))
        beta = beta_scale * (sigma * torch.sqrt(torch.clamp(ld + log_const, min=0)) + math.sqrt(lam) * Cbar)
        covered &= dist <= beta
        min_slack = torch.minimum(min_slack, beta - dist)
        if t in (1, 2, 5, 10, 20, 50, 100, 200, 400, 600, 800, T):
            cov_curve.append((t, float(covered.float().mean())))

    return {
        "uniform_coverage": float(covered.float().mean()),
        "mean_min_slack": float(min_slack.mean()),
        "coverage_curve": cov_curve,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="results")
    ap.add_argument("--reps", type=int, default=4096)
    ap.add_argument("--T", type=int, default=800)
    ap.add_argument("--K", type=int, default=8)
    args = ap.parse_args()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"device={device} torch={torch.__version__} reps={args.reps} T={args.T}")
    t0 = time.time()

    sigma, delta, lam = 0.1, 0.2, 1.0
    inst = build_instance(args.K, args.K, seed=1000)
    Cbar = 1.1 * float(np.linalg.norm(inst["theta_star"]))

    res = {"config": dict(K=args.K, N=inst["N"], T=args.T, reps=args.reps,
                          sigma=sigma, delta=delta, lam=lam, device=device),
           "target_uniform_coverage": 1 - delta / 2}
    variants = {
        "random_correct": dict(),
        "adaptive_perturbed": dict(adaptive=True),
        "control_beta_quarter": dict(beta_scale=0.25),
        "control_noise_4x": dict(sigma_run=4 * sigma),
    }
    for name, kw in variants.items():
        r = run_variant(inst, args.reps, args.T, sigma, delta, lam, Cbar,
                        seed=hash(name) % 2 ** 31, device=device, **kw)
        res[name] = r
        print(f"[{name}] uniform_coverage={r['uniform_coverage']:.4f} "
              f"mean_min_slack={r['mean_min_slack']:.3f}", flush=True)

    res["wall_seconds"] = time.time() - t0
    os.makedirs(args.out, exist_ok=True)
    path = os.path.join(args.out, "claim6_gpu.json")
    with open(path, "w") as f:
        json.dump(res, f, indent=2)
    print(f"saved {path} ({res['wall_seconds']:.0f}s)")


if __name__ == "__main__":
    main()