1import json
  2import math
  3import time
  4import numpy as np
  5
  6
  7def kl_bits(q, p):
  8    return float(np.sum(q * (np.log(q) - np.log(p)) / np.log(2.0)))
  9
 10
 11def l1(q, p):
 12    return float(np.sum(np.abs(q - p)))
 13
 14
 15def poisson_race(p, q, rng, max_candidates=100000):
 16    """Exact marked-Poisson race with theorem stopping rule."""
 17    logp, logq = np.log(p), np.log(q)
 18    logw_max = float(np.max(logq - logp))
 19    # Draw the first arrival; subsequent arrivals are carried into the next iteration.
 20    arrival = -math.log(max(rng.random(), np.finfo(float).tiny))
 21    best_score, best_x = math.inf, -1
 22    for i in range(1, max_candidates + 1):
 23        x = int(rng.choice(len(p), p=p))
 24        score = math.log(arrival) + logp[x] - logq[x]
 25        if score < best_score:
 26            best_score, best_x = score, x
 27        next_arrival = arrival - math.log(max(rng.random(), np.finfo(float).tiny))
 28        if math.log(next_arrival) - logw_max >= best_score:
 29            return best_x, i
 30        arrival = next_arrival
 31    raise RuntimeError("candidate cap exceeded")
 32
 33
 34def rejection_sample(p, q, rng, max_candidates=100000):
 35    """Standard rejection sampler with envelope M=max_x Q(x)/P(x)."""
 36    m = float(np.max(q / p))
 37    for i in range(1, max_candidates + 1):
 38        x = int(rng.choice(len(p), p=p))
 39        if rng.random() <= q[x] / (m * p[x]):
 40            return x, i
 41    raise RuntimeError("rejection candidate cap exceeded")
 42
 43
 44def run_case(name, p, q, n=30000, seed=1234):
 45    rng = np.random.default_rng(seed)
 46    counts = np.zeros(len(p), dtype=np.int64)
 47    race_idx = np.empty(n, dtype=np.int64)
 48    t0 = time.perf_counter()
 49    for k in range(n):
 50        x, i = poisson_race(p, q, rng)
 51        counts[x] += 1
 52        race_idx[k] = i
 53    race_seconds = time.perf_counter() - t0
 54
 55    rng = np.random.default_rng(seed + 1)
 56    rej_idx = np.empty(n, dtype=np.int64)
 57    rej_counts = np.zeros(len(p), dtype=np.int64)
 58    t0 = time.perf_counter()
 59    for k in range(n):
 60        x, i = rejection_sample(p, q, rng)
 61        rej_counts[x] += 1
 62        rej_idx[k] = i
 63    rej_seconds = time.perf_counter() - t0
 64
 65    empirical_race = counts / n
 66    empirical_rej = rej_counts / n
 67    return {
 68        "name": name,
 69        "kl_bits": kl_bits(q, p),
 70        "l1": l1(q, p),
 71        "upper_bound_bits": kl_bits(q, p) + 1.45 * l1(q, p),
 72        "theorem_lower_bound_bits": 0.5 * max(kl_bits(q, p), l1(q, p)),
 73        "race_tv": 0.5 * float(np.sum(np.abs(empirical_race - q))),
 74        "rejection_tv": 0.5 * float(np.sum(np.abs(empirical_rej - q))),
 75        "race_mean_log2_index": float(np.mean(np.log2(race_idx))),
 76        "race_mean_index": float(np.mean(race_idx)),
 77        "rejection_mean_index": float(np.mean(rej_idx)),
 78        "race_p95_index": float(np.quantile(race_idx, 0.95)),
 79        "rejection_p95_index": float(np.quantile(rej_idx, 0.95)),
 80        "race_seconds": race_seconds,
 81        "rejection_seconds": rej_seconds,
 82        "race_target_evals_per_sample": 1.0,
 83        "rejection_target_evals_per_sample": 1.0,
 84    }
 85
 86
 87def main():
 88    # All probabilities are positive, making the finite-support KL well defined.
 89    v = 16
 90    p = np.full(v, 1.0 / v)
 91    cases = []
 92    for name, alpha in [("identical", 0.0), ("close", 0.10), ("moderate", 0.30), ("farther", 0.60)]:
 93        # Q is a smooth tilt of P toward the first symbol; alpha controls distance.
 94        q = np.full(v, (1.0 - alpha) / (v - 1))
 95        q[0] = alpha if alpha > 0 else 1.0 / v
 96        if alpha == 0:
 97            q = p.copy()
 98        cases.append(run_case(name, p, q))
 99    result = {
100        "description": "Finite categorical verification of KL-budgeted Poisson-race sampling",
101        "n_samples_per_case": 30000,
102        "vocab": v,
103        "cases": cases,
104        "notes": [
105            "Race samples use iid P marks and Exp(1) arrivals; stopping uses known max Q/P.",
106            "TV is halved L1, while the theorem's displayed L1 is reported separately.",
107            "Target evaluation count is not a speed advantage here: q is explicitly available; this is an exactness/candidate-index test."
108        ]
109    }
110    with open("results.json", "w") as f:
111        json.dump(result, f, indent=2)
112    print(json.dumps(result, indent=2))
113
114
115if __name__ == "__main__":
116    main()