scripts/make_figs_123.py

"""Figures for the analytic claims 1-3."""

import json
import os
import sys

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from common import F_rho, F_true, cvar, d2_f_rho, d2_f_rho_star

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
FIGS = os.path.join(ROOT, "figs")
os.makedirs(FIGS, exist_ok=True)
plt.rcParams.update({"font.size": 11, "figure.dpi": 130})

rng = np.random.default_rng(20260725)

# ---------------------------------------------------------------- claim 1
fig, ax = plt.subplots(1, 2, figsize=(9.5, 3.6))
rhos = np.logspace(-6, np.log10(0.9), 40)
for lbl, phi in [
    ("Gaussian n=200", rng.normal(0, 1.5, 200)),
    (
        "bimodal n=200",
        np.concatenate([rng.normal(-3, 0.5, 100), rng.normal(3, 0.5, 100)]),
    ),
    ("spiky n=200", np.concatenate([rng.normal(0, 0.3, 199), [8.0]])),
]:
    w = np.full(len(phi), 1 / len(phi))
    Fv = F_true(phi, w)
    gap = np.array([Fv - F_rho(phi, r, w)[0] for r in rhos])
    ax[0].loglog(rhos, gap, marker="o", ms=2.5, lw=1.2, label=lbl)
ax[0].loglog(rhos, rhos / 2, "k--", lw=1, label=r"$\rho/2$ (slope 1)")
ax[0].set_xlabel(r"$\rho$")
ax[0].set_ylabel(r"$F - F_\rho$  (always $>0$)")
ax[0].set_title(r"Approximation error is $O(\rho)$")
ax[0].legend(fontsize=8)
ax[0].grid(alpha=0.3, which="both")

d = json.load(open(os.path.join(ROOT, "outputs", "claim1_sandwich.json")))
a = d["T4_direction_audit"]
ax[1].bar(
    ["$F \\leq F_\\rho$\n(as printed)", "$F_\\rho \\leq F$\n(actual)"],
    [a["F_le_Frho_holds"], a["Frho_le_F_holds"]],
    color=["#c0392b", "#27ae60"],
)
ax[1].set_ylabel("instances holding (of 3000)")
ax[1].set_title("Direction audit of the printed sandwich")
for i, v in enumerate([a["F_le_Frho_holds"], a["Frho_le_F_holds"]]):
    ax[1].text(i, v + 60, str(v), ha="center", fontweight="bold")
ax[1].set_ylim(0, 3400)
fig.tight_layout()
fig.savefig(os.path.join(FIGS, "claim1.png"), bbox_inches="tight")
plt.close(fig)

# ---------------------------------------------------------------- claim 2
fig, ax = plt.subplots(1, 2, figsize=(9.5, 3.6))
for rho in [0.03, 0.1, 0.3]:
    t = np.linspace(1e-6, 1 / rho * (1 - 1e-6), 4000)
    ax[0].plot(t * rho, d2_f_rho(t, rho) / rho, lw=1.4, label=rf"$\rho={rho}$")
ax[0].axhline(4, color="k", ls="--", lw=1, label="tight modulus $4$")
ax[0].axhline(1, color="r", ls=":", lw=1.2, label=r"claimed modulus $1$")
ax[0].set_ylim(0, 20)
ax[0].set_xlabel(r"$\rho t$  (domain $[0,1]$)")
ax[0].set_ylabel(r"$f_\rho''(t)\,/\,\rho$")
ax[0].legend(fontsize=8)
ax[0].set_title(r"$f_\rho$ is $4\rho$-strongly convex $\Rightarrow$ $\rho$-SC")
ax[0].grid(alpha=0.3)
for rho in [0.03, 0.1, 0.3]:
    s = np.linspace(-8 - np.log(rho), 8 - np.log(rho), 4000)
    ax[1].plot(
        s + np.log(rho), d2_f_rho_star(s, rho) * rho, lw=1.4, label=rf"$\rho={rho}$"
    )
ax[1].axhline(0.25, color="k", ls="--", lw=1, label="tight $1/4$")
ax[1].axhline(1.0, color="r", ls=":", lw=1.2, label="claimed $1$")
ax[1].set_xlabel(r"$s + \log\rho$")
ax[1].set_ylabel(r"$\rho\,(f_\rho^*)''(s)$")
ax[1].set_title(
    r"$f_\rho^*$ is $\frac{1}{4\rho}$-smooth $\Rightarrow \frac{1}{\rho}$-smooth"
)
ax[1].legend(fontsize=8)
ax[1].grid(alpha=0.3)
fig.tight_layout()
fig.savefig(os.path.join(FIGS, "claim2.png"), bbox_inches="tight")
plt.close(fig)

# ---------------------------------------------------------------- claim 3
fig, ax = plt.subplots(1, 2, figsize=(9.5, 3.6))
phi = np.random.default_rng(2026).normal(0, 1.5, 300)
w = np.full(300, 1 / 300)
lams = np.logspace(-4, 0.5, 30)
for rho in [0.05, 0.1, 0.25]:
    C = cvar(phi, rho, w)
    vals = np.array(
        [lam * F_rho(phi / lam, rho, w)[0] - lam * (np.log(rho) - 1) for lam in lams]
    )
    ax[0].semilogx(
        lams,
        vals,
        lw=1.4,
        label=rf"$\lambda F_\rho(\phi/\lambda)-\lambda(\log\rho-1)$, $\rho={rho}$",
    )
    ax[0].axhline(C, ls="--", lw=1, color=ax[0].lines[-1].get_color())
ax[0].set_xlabel(r"$\lambda$")
ax[0].set_ylabel("value")
ax[0].set_title(r"CVaR limit as $\lambda\to0$ (dashed = CVaR$_\rho$)")
ax[0].legend(fontsize=7)
ax[0].grid(alpha=0.3)

rhos = np.logspace(-6, -0.3, 30)
for lam in [0.5, 1.0, 2.0]:
    tgt = lam * F_true(phi / lam, w)
    err = np.array([abs(lam * F_rho(phi / lam, r, w)[0] - tgt) for r in rhos])
    ax[1].loglog(rhos, err, marker="o", ms=2.5, lw=1.2, label=rf"$\lambda={lam}$")
ax[1].loglog(rhos, rhos, "k--", lw=1, label="slope 1")
ax[1].set_xlabel(r"$\rho$")
ax[1].set_ylabel(r"$|\lambda F_\rho(\phi/\lambda) - $LogSumExp$|$")
ax[1].set_title(r"LogSumExp limit as $\rho\to0$")
ax[1].legend(fontsize=8)
ax[1].grid(alpha=0.3, which="both")
fig.tight_layout()
fig.savefig(os.path.join(FIGS, "claim3.png"), bbox_inches="tight")
plt.close(fig)
print("figures written")