code/test_smoke.py

"""Smoke tests for botlib: exact identities, solver cross-checks, tiny EntUCB run."""
import numpy as np
import sys
sys.path.insert(0, __file__.rsplit("/", 1)[0])
from botlib import (make_instance, kantorovich, sinkhorn_log, entropic_value,
                    run_entucb)


def main():
    rng = np.random.default_rng(0)
    inst = make_instance(4, 5, seed=1, cost_kind="smooth")
    N = inst.N

    # 1) basis is rho-orthonormal
    G = inst.B @ (inst.rho[:, None] * inst.B.T)
    err_orth = np.max(np.abs(G - np.eye(N)))
    print(f"[basis] max |B W B^T - I| = {err_orth:.2e}")
    assert err_orth < 1e-10

    # 2) Claim-1 mechanism: <c|pi> = theta* . a(pi) for random feasible plans
    import ot as pot
    errs = []
    for _ in range(50):
        # random feasible plan: mix of independent coupling and random sinkhorn plans
        M = rng.standard_normal((inst.K, inst.Kp))
        P = pot.sinkhorn(inst.mu, inst.nu, M - M.min(), reg=rng.uniform(0.05, 1.0))
        lhs = inst.pairing(inst.c_vec, P)
        rhs = float(inst.theta_star @ inst.embed(P))
        errs.append(abs(lhs - rhs))
    print(f"[claim1] max |<c|pi> - <theta*,a(pi)>| over 50 plans = {max(errs):.2e}")
    assert max(errs) < 1e-10

    # 3) Parseval: ||c||_L2(rho) = ||theta*||_2
    n1 = np.sqrt(inst.inner(inst.c_vec, inst.c_vec))
    n2 = np.linalg.norm(inst.theta_star)
    print(f"[parseval] ||c||_rho = {n1:.12f}  ||theta*||_2 = {n2:.12f}")
    assert abs(n1 - n2) < 1e-10

    # 4) our log-domain Sinkhorn vs POT (KL-vs-negentropy argmin equivalence)
    eps = 0.1
    P1, f, g = sinkhorn_log(inst.mu, inst.nu, inst.cost, eps)
    P2 = pot.bregman.sinkhorn_log(inst.mu, inst.nu, inst.cost, reg=eps, numItermax=20000, stopThr=1e-14)
    print(f"[sinkhorn] max |P_ours - P_pot| = {np.max(np.abs(P1 - P2)):.2e}")
    assert np.max(np.abs(P1 - P2)) < 1e-8
    # marginals
    print(f"[sinkhorn] marginal err = {np.max(np.abs(P1.sum(1) - inst.mu)):.2e}, {np.max(np.abs(P1.sum(0) - inst.nu)):.2e}")

    # 5) EOT value decreases to Kantorovich value as eps -> 0, gap >= 0
    kv, _ = kantorovich(inst)
    gaps = []
    for eps in [0.5, 0.1, 0.02, 0.004]:
        ev, _, _ = entropic_value(inst, eps)
        gaps.append(ev - kv)
    print(f"[eot->kant] kant={kv:.6f} gaps={['%.3e' % g for g in gaps]}")
    assert all(g > -1e-9 for g in gaps) and gaps[-1] < gaps[0]

    # 6) tiny EntUCB run end to end
    res = run_entucb(inst, T=150, sigma=0.05, delta=0.1, lam=1.0,
                     Cbar=1.1 * np.linalg.norm(inst.theta_star),
                     eps_schedule=lambda t: 0.75 * t ** -0.75, eta=0.75, seed=3)
    print(f"[entucb] T=150 kant_regret={res.kant_regret[-1]:.3f} ent_regret={res.ent_regret[-1]:.3f} "
          f"bound_lit={res.bound_literal[-1]:.1f} cert_viol={res.cert_viol_frac:.3f} covered_all={res.covered_all}")
    assert res.ent_regret[-1] <= res.bound_literal[-1]
    # entropic regret should be nonnegative-ish up to certificate tolerance and grow slowly
    print("smoke OK")


if __name__ == "__main__":
    main()