code/exp_claim1.py

"""Claim 1 audit — the Fourier embedding that reduces BOT to a linear bandit.

Paper (arXiv:2502.07397v2, Sec. 3.1 / App. A-B): the operator
    F: phi in L^2(R^d; rho) |-> int phi(x) e^{-2 pi i <x|.>} d rho(x)
"is an isometry on L^2(R^d; rho)" (Thm A.2 as stated for general rho), and
Eq. (2) [OpenReview Eq. (7)]:
    <c|pi> = <F c(-.) | F pi>_{L^2(rho)} .

We audit three regimes:

(A) VALID MECHANISM (what the theorems actually need): for ANY discrete BOT
    instance, an rho-orthonormal basis of L^2(rho) reduces <c|pi> to a
    Hilbert-space inner product exactly.  Also: the properly *paired* Fourier
    transforms are unitary — (i) unitary DFT (uniform grid <-> integer
    frequencies), (ii) self-dual Plancherel on R^d (fine-grid FFT proxy).

(B) LITERAL CLAIM FALSE FOR ATOMIC rho: the operator F_rho evaluated on
    supp(rho), as an operator on L^2(rho) (same rho on both sides, as printed),
    is NOT an isometry: singular values far from 1 for generic atoms.

(C) CONSTRUCTIVE DEGENERACY ("Fourier operator = 1"): marginals supported on
    the INTEGER LATTICE make e^{-2 pi i <x|xi>} = 1 for every x, xi in
    supp(rho): F_rho collapses to the rank-one operator phi -> (int phi drho) 1.
    Every transport plan embeds to the SAME element; Eq. (2)'s right-hand side
    equals int c d rho regardless of pi; the frequency-domain linear model is
    misspecified and EntUCB-with-F_rho suffers linear regret.

Outputs: results/claim1.json + figures.
"""
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, kantorovich, sinkhorn_log

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


def l2rho_operator_singvals(points: np.ndarray, weights: np.ndarray) -> np.ndarray:
    """Singular values of F_rho: L^2(rho) -> L^2(rho) for atomic rho.

    (F phi)(xi) = sum_x phi(x) e^{-2 pi i <x|xi>} w(x), evaluated at xi in supp(rho).
    As a matrix in the *weighted* geometry: A = W^{1/2} K W^{1/2},
    K_{xi,x} = e^{-2 pi i <x|xi>}.  Isometry <=> all singular values = 1.
    """
    if points.ndim == 1:
        points = points[:, None]
    G = points @ points.T  # <x|xi> for all pairs
    K = np.exp(-2j * np.pi * G)
    Ws = np.sqrt(weights)
    A = (Ws[:, None] * K) * Ws[None, :]
    return np.linalg.svd(A, compute_uv=False)


def main():
    rng = np.random.default_rng(0)
    res = {}

    # ---------------- (A) valid mechanism --------------------------------
    # A1: exact reduction via rho-orthonormal basis, many random instances
    max_err, max_parseval = 0.0, 0.0
    import ot as pot
    for s in range(20):
        inst = make_instance(rng.integers(3, 7), rng.integers(3, 7), seed=100 + s)
        for _ in range(20):
            M = rng.standard_normal((inst.K, inst.Kp))
            P = pot.sinkhorn(inst.mu, inst.nu, M - M.min(), reg=float(rng.uniform(0.05, 1.0)))
            lhs = inst.pairing(inst.c_vec, P)
            rhs = float(inst.theta_star @ inst.embed(P))
            max_err = max(max_err, abs(lhs - rhs))
        n1 = np.sqrt(inst.inner(inst.c_vec, inst.c_vec))
        max_parseval = max(max_parseval, abs(n1 - np.linalg.norm(inst.theta_star)))
    res["A1_pairing_identity_max_err"] = max_err
    res["A1_parseval_max_err"] = max_parseval
    print(f"[A1] exact reduction <c|pi>=<theta*,a(pi)>: max err {max_err:.3e} "
          f"(400 plans, 20 instances); Parseval max err {max_parseval:.3e}")

    # A2: unitary DFT pairing (uniform grid <-> integer frequencies)
    n = 64
    Fdft = np.exp(-2j * np.pi * np.outer(np.arange(n), np.arange(n)) / n) / np.sqrt(n)
    res["A2_dft_unitarity_err"] = float(np.max(np.abs(Fdft.conj().T @ Fdft - np.eye(n))))
    print(f"[A2] unitary DFT (paired measures): ||F*F - I|| = {res['A2_dft_unitarity_err']:.3e}")

    # A3: self-dual Plancherel on R (Lebesgue), fine-grid FFT proxy
    L, m = 40.0, 2 ** 14
    dx = L / m
    x = (np.arange(m) - m // 2) * dx
    phi = np.exp(-np.pi * x ** 2) * (1 + 0.3 * np.cos(3 * x))
    ph_hat = np.fft.fftshift(np.fft.fft(np.fft.ifftshift(phi))) * dx
    nrm_x = np.sum(np.abs(phi) ** 2) * dx
    dxi = 1.0 / L
    nrm_f = np.sum(np.abs(ph_hat) ** 2) * dxi
    res["A3_plancherel_lebesgue_relerr"] = float(abs(nrm_x - nrm_f) / nrm_x)
    print(f"[A3] Plancherel wrt Lebesgue on R (FFT proxy): rel err {res['A3_plancherel_lebesgue_relerr']:.3e}")

    # ---------------- (B) literal claim for atomic rho -------------------
    sv_generic = l2rho_operator_singvals(rng.uniform(0, 1, 12), np.full(12, 1 / 12))
    res["B_generic_atoms_singvals_minmax"] = [float(sv_generic.min()), float(sv_generic.max())]
    # uniform grid on [0,1) with SAME rho on both sides (not the DFT pairing!)
    g = np.arange(12) / 12
    sv_grid = l2rho_operator_singvals(g, np.full(12, 1 / 12))
    res["B_grid_same_measure_singvals_minmax"] = [float(sv_grid.min()), float(sv_grid.max())]
    # two-atom counterexample {0,1}: kernel e^{-2 pi i x xi} = 1 on supp x supp
    sv2 = l2rho_operator_singvals(np.array([0.0, 1.0]), np.array([0.5, 0.5]))
    res["B_two_atom_lattice_singvals"] = [float(v) for v in sv2]
    print(f"[B] singular values of F_rho on L^2(rho): generic atoms [{sv_generic.min():.3f}, {sv_generic.max():.3f}] "
          f"(isometry requires all =1); grid/12 [{sv_grid.min():.3f}, {sv_grid.max():.3f}]; "
          f"lattice {{0,1}}: {np.round(sv2, 6).tolist()}")

    # ---------------- (C) lattice degeneracy -----------------------------
    # marginals on integer points; rho = mu x nu on Z^2
    K = Kp = 5
    xi_ = np.arange(K, dtype=float)          # 0..4 integers
    yj = np.arange(Kp, dtype=float)
    mu = np.array([0.1, 0.3, 0.2, 0.25, 0.15])
    nu = np.array([0.2, 0.15, 0.3, 0.1, 0.25])
    rho = np.outer(mu, nu)
    pairs = np.array([[xi, yj_] for xi in xi_ for yj_ in yj])   # supp(rho) in Z^2
    # kernel on supp(rho) x supp(rho)
    Kmat = np.exp(-2j * np.pi * (pairs @ pairs.T))
    res["C_kernel_equals_one"] = float(np.max(np.abs(Kmat - 1.0)))
    sv = l2rho_operator_singvals(pairs, rho.ravel())
    res["C_lattice_singvals_top3"] = [float(v) for v in sv[:3]]
    res["C_lattice_rank"] = int(np.sum(sv > 1e-10))
    print(f"[C] integer-lattice instance: max|kernel-1| = {res['C_kernel_equals_one']:.2e}; "
          f"F_rho singular values top3 = {np.round(sv[:3], 6).tolist()}, numerical rank = {res['C_lattice_rank']} "
          f"(rank-one: every phi maps to the constant int phi drho)")

    # every transport plan embeds to F pi == 1 on supp(rho)
    cost = np.abs(xi_[:, None] - yj[None, :])          # c(x,y)=|x-y|, Lipschitz
    kant_val, plan_opt = kantorovich_discrete(mu, nu, cost)
    plans = {
        "independent": np.outer(mu, nu),
        "optimal": plan_opt,
        "sinkhorn_eps0.1": sinkhorn_log(mu, nu, cost, 0.1)[0],
    }
    emb = {}
    for name, P in plans.items():
        # F pi evaluated on supp(rho): sum_z e^{-2 pi i <z|xi>} P(z) = sum P = 1
        Fpi = (Kmat @ P.ravel())
        emb[name] = [float(np.max(np.abs(Fpi - 1.0))), float(np.sum(P * cost))]
    res["C_all_plans_embed_to_one"] = {k: v[0] for k, v in emb.items()}
    res["C_plan_costs"] = {k: v[1] for k, v in emb.items()}
    # Eq (2)/(7) RHS: <F c(-.) | F pi>_rho = (int c drho) * 1  for every pi
    rhs = float(np.sum(rho * cost))
    res["C_eq7_rhs_constant"] = rhs
    res["C_eq7_lhs_range"] = [min(v[1] for v in emb.values()), max(v[1] for v in emb.values())]
    print(f"[C] all plans embed to F pi == 1 (max dev {max(v[0] for v in emb.values()):.2e}); "
          f"Eq.(7) RHS = int c drho = {rhs:.4f} for EVERY plan, while <c|pi> ranges "
          f"{res['C_eq7_lhs_range'][0]:.4f} .. {res['C_eq7_lhs_range'][1]:.4f} -> identity FAILS")

    # C2: bandit consequence — frequency-domain model misspecified: any
    # estimator sees identical features a_t == a for all actions.  The best
    # frequency-domain prediction is constant, so any algorithm relying on the
    # F_rho embedding cannot distinguish plans: regret >= gap * T for one of
    # the candidate plans.  Demonstrate with RLS-UCB on the rank-one feature.
    T = 2000
    sigma = 0.05
    rng2 = np.random.default_rng(7)
    # features: a(pi) = F_rho embedding coefficients == same vector for all pi
    # -> UCB over {independent, optimal} plans is a coin flip forever; regret of
    #    always playing "independent" (what constant-model tie-breaks to):
    gap = res["C_plan_costs"]["independent"] - kant_val
    regret_curve = gap * np.arange(1, T + 1)
    res["C_linear_regret_slope"] = float(gap)
    print(f"[C2] misspecified frequency model => constant features; playing the model-indistinguishable "
          f"independent coupling gives regret slope {gap:.4f} per step (linear regret), "
          f"vs a valid embedding (exp_bandit claims 2-5) which achieves sublinear regret on the same instance.")

    with open(os.path.join(OUT, "claim1.json"), "w") as f:
        json.dump(res, f, indent=2)
    print("saved results/claim1.json")


def kantorovich_discrete(mu, nu, cost):
    import ot as pot
    plan = pot.emd(mu, nu, np.ascontiguousarray(cost))
    return float(np.sum(plan * cost)), plan


if __name__ == "__main__":
    main()