scripts/claim1_sandwich.py

"""Claim 1 -- Def 2.1 / Eqs (3)-(4) safe KL divergence, and Prop 2.4:
the induced approximation satisfies (as printed in the paper) F_rho - O(rho) <= F <= F_rho.

Independent tests:
  T0  f_rho is a legal f-divergence generator (convex, finite exactly on [0,1/rho]),
      and f_rho -> f_0 = t log t + 1 - t as rho -> 0.
  T1  Lemma 2.3: (1/rho) log(1+rho e^s) - 1 really is the convex conjugate of f_rho
      (brute-force sup_t s t - f_rho(t) on a fine grid + Newton).
  T2  Duality (5) <-> (6): F_rho from the dual formula equals the *primal* value
      sup_nu { <phi,nu> - D_rho(nu,mu) } computed by constrained convex optimisation
      over the simplex with the hard density bound dnu/dmu <= 1/rho, and the
      maximiser matches the closed form e^{phi-a*}/(1+rho e^{phi-a*}).
  T3  Prop 2.4(i) monotonicity in rho, (ii) F_rho -> F as rho -> 0 with fitted rate.
  T4  DIRECTION AUDIT of the printed sandwich F_rho - O(rho) <= F <= F_rho.
  T5  Prop 2.4(iii) and (iv) explicit lower bounds, inside and outside their stated
      admissible rho-ranges (boundary audit).
  T6  Corollary 2.5 (LogSumExp form) over random inputs.
"""

import json
import os
import sys

import numpy as np
from scipy.optimize import minimize
from scipy.special import logsumexp

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from common import (
    F_rho,
    F_rho_bruteforce,
    F_true,
    alpha_star,
    d2_f_rho,
    d_f_rho,
    f_rho,
    f_rho_star,
)

OUT = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "outputs"
)
os.makedirs(OUT, exist_ok=True)
SEED = 20260725
res = {"seed": SEED}
rng = np.random.default_rng(SEED)

# ---------------------------------------------------------------- T0
rhos = [1e-4, 1e-3, 1e-2, 0.03, 0.1, 0.3, 0.5, 0.9]
t0 = {
    "convexity_min_second_deriv": {},
    "domain_ok": True,
    "limit_to_f0_max_abs_err": {},
}
for rho in rhos:
    tt = np.linspace(1e-9, 1.0 / rho - 1e-9, 200001)
    d2 = d2_f_rho(tt, rho)
    t0["convexity_min_second_deriv"][str(rho)] = float(d2.min())
    if not np.isinf(f_rho(np.array([1.0 / rho + 1e-6]), rho)[0]):
        t0["domain_ok"] = False
for rho in [1e-3, 1e-4, 1e-5, 1e-6]:
    tt = np.linspace(1e-6, 5.0, 20001)
    f0 = tt * np.log(tt) + 1 - tt
    t0["limit_to_f0_max_abs_err"][str(rho)] = float(np.max(np.abs(f_rho(tt, rho) - f0)))
res["T0_generator"] = t0

# ---------------------------------------------------------------- T1 conjugate
t1 = {"max_abs_err": 0.0, "n_checks": 0}
for rho in rhos:
    for s in np.linspace(-20, 20, 81):
        tt = np.linspace(1e-12, 1.0 / rho * (1 - 1e-12), 400001)
        sup = np.max(s * tt - f_rho(tt, rho))
        # refine with the analytic stationary point t = e^s/(1+rho e^s)
        tstar = np.exp(s - np.logaddexp(0.0, s + np.log(rho))) / 1.0
        tstar = min(max(tstar, 1e-14), 1.0 / rho - 1e-14)
        sup = max(sup, float(s * tstar - f_rho(np.array([tstar]), rho)[0]))
        err = abs(sup - float(f_rho_star(np.array([s]), rho)[0]))
        t1["max_abs_err"] = max(t1["max_abs_err"], err)
        t1["n_checks"] += 1
res["T1_lemma23_conjugate"] = t1

# ---------------------------------------------------------------- T2 primal-dual
t2 = {"cases": [], "max_abs_gap": 0.0, "max_density_err": 0.0}
for trial in range(12):
    n = int(rng.integers(4, 8))
    phi = rng.normal(0, 2.0, size=n)
    w = rng.dirichlet(np.ones(n))
    rho = float(rng.choice([0.03, 0.1, 0.3, 0.5]))
    dual, a = F_rho(phi, rho, w)

    def neg(r):
        return -(np.sum(w * phi * r) - np.sum(w * f_rho(r, rho)))

    cons = [{"type": "eq", "fun": lambda r: np.sum(w * r) - 1.0}]
    bnds = [(1e-10, 1.0 / rho - 1e-10)] * n
    best, bestr = -np.inf, None
    for _ in range(6):
        r0 = rng.uniform(1e-3, 1.0 / rho, size=n)
        r0 = r0 / np.sum(w * r0)
        r0 = np.clip(r0, 1e-9, 1.0 / rho - 1e-9)
        try:
            sol = minimize(
                neg,
                r0,
                bounds=bnds,
                constraints=cons,
                method="SLSQP",
                options={"maxiter": 800, "ftol": 1e-14},
            )
            if sol.success and -sol.fun > best:
                best, bestr = -sol.fun, sol.x
        except Exception:
            pass
    closed = np.exp(phi - a) / (1.0 + rho * np.exp(phi - a))
    t2["cases"].append(
        {
            "n": n,
            "rho": rho,
            "dual_Frho": dual,
            "primal_sup": float(best),
            "abs_gap": abs(dual - best),
            "density_max_abs_err": float(np.max(np.abs(closed - bestr))),
            "density_max": float(closed.max()),
            "one_over_rho": 1.0 / rho,
        }
    )
    t2["max_abs_gap"] = max(t2["max_abs_gap"], abs(dual - best))
    t2["max_density_err"] = max(
        t2["max_density_err"], float(np.max(np.abs(closed - bestr)))
    )
res["T2_primal_dual"] = t2


# ---------------------------------------------------------------- random instances
def random_instance(rng, kind=None):
    kind = kind or rng.choice(["gauss", "heavy", "spiky", "uniform", "bimodal"])
    n = int(rng.integers(5, 400))
    if kind == "gauss":
        phi = rng.normal(0, rng.uniform(0.2, 4.0), size=n)
    elif kind == "heavy":
        phi = rng.standard_t(2.0, size=n) * rng.uniform(0.5, 3.0)
    elif kind == "spiky":
        phi = rng.normal(0, 0.3, size=n)
        phi[0] += rng.uniform(3, 12)
    elif kind == "uniform":
        phi = rng.uniform(-6, 6, size=n)
    else:
        phi = np.concatenate(
            [rng.normal(-3, 0.5, n // 2), rng.normal(3, 0.5, n - n // 2)]
        )
    w = rng.dirichlet(np.ones(n) * rng.uniform(0.3, 3.0))
    return phi, w, str(kind)


# ---------------------------------------------------------------- T3 monotone + limit
rho_grid = np.array([0.5, 0.3, 0.1, 0.03, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6])
mono_viol, n_mono = 0, 0
rate_fits = []
for trial in range(400):
    phi, w, _ = random_instance(rng)
    vals = [F_rho(phi, r, w)[0] for r in rho_grid]
    for i in range(len(rho_grid) - 1):  # rho_grid decreasing
        n_mono += 1
        if not (vals[i] <= vals[i + 1] + 1e-10):  # F_{rho'} <= F_rho for rho <= rho'
            mono_viol += 1
    Ftrue = F_true(phi, w)
    gaps = np.array([Ftrue - v for v in vals])
    small = rho_grid <= 1e-2
    p = np.polyfit(np.log(rho_grid[small]), np.log(np.maximum(gaps[small], 1e-300)), 1)
    rate_fits.append(float(p[0]))
    if trial < 200:
        rate_fits[-1] = float(p[0])
res["T3_monotone_and_limit"] = {
    "n_pairs": n_mono,
    "monotonicity_violations": mono_viol,
    "gap_rate_exponent_mean": float(np.mean(rate_fits)),
    "gap_rate_exponent_std": float(np.std(rate_fits)),
    "gap_rate_exponent_min": float(np.min(rate_fits)),
    "gap_rate_exponent_max": float(np.max(rate_fits)),
    "note": "exponent of (F - F_rho) ~ rho^p fitted on rho <= 1e-2",
}

# ---------------------------------------------------------------- T4 DIRECTION AUDIT
audit = {
    "n": 0,
    "F_le_Frho_holds": 0,
    "Frho_le_F_holds": 0,
    "max_(F - Frho)": -np.inf,
    "min_(F - Frho)": np.inf,
    "examples": [],
}
for trial in range(3000):
    phi, w, kind = random_instance(rng)
    rho = float(rng.choice([1e-4, 1e-3, 1e-2, 0.03, 0.1, 0.3, 0.5, 0.7, 0.9]))
    Fv = F_true(phi, w)
    Fr, _ = F_rho(phi, rho, w)
    d = Fv - Fr
    audit["n"] += 1
    if Fv <= Fr + 1e-12:
        audit["F_le_Frho_holds"] += 1
    if Fr <= Fv + 1e-12:
        audit["Frho_le_F_holds"] += 1
    audit["max_(F - Frho)"] = max(audit["max_(F - Frho)"], float(d))
    audit["min_(F - Frho)"] = min(audit["min_(F - Frho)"], float(d))
    if len(audit["examples"]) < 5 and rho in (0.1, 0.3):
        audit["examples"].append(
            {
                "kind": kind,
                "n": int(phi.size),
                "rho": rho,
                "F": float(Fv),
                "F_rho": float(Fr),
                "F_minus_Frho": float(d),
            }
        )
# a fully explicit hand-checkable example: mu = delta_0, phi = 0
rho = 0.1
Fr_pt, _ = F_rho(np.array([0.0]), rho, np.array([1.0]))
closed = np.log(1 - rho) - 1 - np.log(1 - rho) / rho
audit["point_mass_example"] = {
    "rho": rho,
    "F": 0.0,
    "F_rho_numeric": float(Fr_pt),
    "F_rho_closed_form": float(closed),
    "closed_form": "log(1-rho) - 1 - log(1-rho)/rho",
}
audit["bruteforce_crosscheck"] = float(
    abs(F_rho_bruteforce(np.array([0.0]), rho, np.array([1.0])) - Fr_pt)
)
res["T4_direction_audit"] = audit

# ---------------------------------------------------------------- T5 Prop 2.4(iii),(iv)
t5 = {
    "iii_inside": {"n": 0, "viol": 0, "min_slack": np.inf},
    "iii_outside": {"n": 0, "viol": 0, "min_slack": np.inf},
    "iv_inside": {"n": 0, "viol": 0, "min_slack": np.inf},
    "iv_outside": {"n": 0, "viol": 0, "min_slack": np.inf},
}
for trial in range(1500):
    phi, w, _ = random_instance(rng)
    F1 = F_true(phi, w)
    F2 = F_true(2 * phi, w)
    a_cap = 0.25 * np.exp(2 * F1 - F2)  # admissible rho for (iii)
    M = float(phi.max())
    b_cap = np.exp(F1 - M)  # admissible rho for (iv)
    for rho in [1e-4, 1e-3, 1e-2, 0.03, 0.1, 0.3, 0.5, 0.9]:
        if rho >= 1:
            continue
        Fr, _ = F_rho(phi, rho, w)
        lb3 = F1 + rho / 2 - 4 * rho * np.exp(F2 - 2 * F1)
        key = "iii_inside" if rho <= a_cap else "iii_outside"
        t5[key]["n"] += 1
        slack = Fr - lb3
        t5[key]["min_slack"] = min(t5[key]["min_slack"], float(slack))
        if slack < -1e-9:
            t5[key]["viol"] += 1
        lb4 = F1 - rho * np.exp(M - F1)
        key = "iv_inside" if rho <= b_cap else "iv_outside"
        t5[key]["n"] += 1
        slack = Fr - lb4
        t5[key]["min_slack"] = min(t5[key]["min_slack"], float(slack))
        if slack < -1e-9:
            t5[key]["viol"] += 1
for k in t5:
    t5[k]["min_slack"] = (
        float(t5[k]["min_slack"]) if np.isfinite(t5[k]["min_slack"]) else None
    )
res["T5_explicit_bounds"] = t5

# ---------------------------------------------------------------- T6 Corollary 2.5
t6 = {"n": 0, "lower_viol": 0, "upper_viol": 0, "max_gap_over_rho": 0.0, "worst": None}
for trial in range(3000):
    n = int(rng.integers(2, 300))
    a = rng.normal(0, rng.uniform(0.2, 5.0), size=n)
    rho = float(rng.choice([1e-3, 1e-2, 0.03, 0.1, 0.3, 0.6, 0.9]))
    lse = float(logsumexp(a))
    # inf_alpha alpha - 1 + (1/rho) sum_i log(1 + rho e^{a_i - alpha})
    # == log n + F_{rho/n}(id; mu_n) with mu_n uniform  -> reuse F_rho with weights 1/n
    val = np.log(n) + F_rho(a, rho / n, np.full(n, 1.0 / n))[0]
    t6["n"] += 1
    if val < lse - rho - 1e-9:
        t6["lower_viol"] += 1
    if val > lse + 1e-9:
        t6["upper_viol"] += 1
    g = (lse - val) / rho
    if g > t6["max_gap_over_rho"]:
        t6["max_gap_over_rho"] = float(g)
        t6["worst"] = {
            "n": n,
            "rho": rho,
            "LSE": lse,
            "approx": float(val),
            "gap": float(lse - val),
        }
res["T6_corollary25"] = t6

with open(os.path.join(OUT, "claim1_sandwich.json"), "w") as fh:
    json.dump(res, fh, indent=2, default=float)
print(json.dumps(res, indent=2, default=float)[:6000])