scripts/make_figures.py

"""Four poster figures, built from the raw result JSONs in outputs/."""

from __future__ import annotations

import json
import os

import matplotlib

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

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(ROOT, "outputs")
FIG = os.path.join(ROOT, "figs")
os.makedirs(FIG, exist_ok=True)

ACC = "#0F766E"
ACC2 = "#B45309"
GREY = "#555555"
plt.rcParams.update(
    {
        "font.size": 13,
        "axes.grid": True,
        "grid.alpha": 0.25,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "figure.dpi": 150,
        "savefig.bbox": "tight",
    }
)


def load(n):
    with open(os.path.join(OUT, n)) as f:
        return json.load(f)


c2, c3, c4, c5, c5s, c6 = (
    load("claim2_static.json"),
    load("claim3_dynamic.json"),
    load("claim4_highprob.json"),
    load("claim5_lowerbound.json"),
    load("claim5_supplement.json"),
    load("claim6_conjecture.json"),
)

# ------------------------------------------------------------------ figure 1: Theorem 3.1
B = c2["B_scalings"]
fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))
T = np.array(B["T_sweep"]["T"], float)
y = np.array(B["T_sweep"]["mean_regret"], float)
ax[0].loglog(T, y, "o-", color=ACC, label="measured $E[R_T(u)]$, $d=8$")
ax[0].loglog(
    T, y[0] * (T / T[0]) ** 0.5, "--", color=GREY, label=r"slope $1/2$ (Thm 3.1)"
)
d = np.array(B["d_sweep"]["d"], float)
yd = np.array(B["d_sweep"]["mean_regret"], float)
ax[0].loglog(d * 250, yd, "s-", color=ACC2, label="measured, sweeping $d$ ($T=2000$)")
ax[0].set_xlabel("$T$   (orange series: $250\\,d$)")
ax[0].set_ylabel("expected regret")
ax[0].set_title(
    "fitted exponents  $T^{%.3f\\pm%.3f}$,  $d^{%.3f\\pm%.3f}$"
    % (
        B["T_sweep"]["fitted_exponent"],
        B["T_sweep"]["stderr"],
        B["d_sweep"]["fitted_exponent"],
        B["d_sweep"]["stderr"],
    ),
    fontsize=12,
)
ax[0].legend(fontsize=9, frameon=False)

D = c2["D_kappa_mechanism"]["rows"]
dd = np.array([r["d"] for r in D], float)
real = np.array([r["mean_sum_ltilde_sq_over_V_T"] for r in D])
ax[1].loglog(
    dd,
    4 * dd**2,
    "--",
    color=ACC2,
    label=r"a.s. bound $4d^2$  $\Rightarrow\ \kappa=1$ (factor $d$)",
)
ax[1].loglog(
    dd,
    2 * dd,
    "--",
    color=ACC,
    label=r"$E$-bound $2d$  $\Rightarrow\ \kappa=\sqrt{d}$ (factor $\sqrt{d}$)",
)
ax[1].loglog(
    dd, real, "o-", color="black", label=r"realised $\sum_t\|\tilde\ell_t\|^2/V_T$"
)
ax[1].set_xlabel("$d$")
ax[1].set_ylabel(r"$\sum_t\|\tilde\ell_t\|^2 / V_T$")
ax[1].set_title(
    (
        r"realised dimension factor: $d^{%.3f\pm%.3f}$" % (0.5 * 1, 0.0)
        if False
        else r"realised exponent $%.3f\pm%.3f$ (0.5 $=\kappa=\sqrt{d}$)"
        % (
            c2["D_kappa_mechanism"]["fitted_d_exponent_of_realised_dimension_factor"],
            c2["D_kappa_mechanism"]["stderr"],
        )
    ),
    fontsize=12,
)
ax[1].legend(fontsize=9, frameon=False)
fig.savefig(os.path.join(FIG, "fig_claim2_static.png"))
plt.close(fig)

# ------------------------------------------------------------------ figure 2: Theorem 3.3
CD = c3["C_D_path_length_sweep"]
reg = CD["regression_of_squared_regret_on_P_T"]
P = np.array([r["P_T"] for r in CD["rows"]])
R = np.array([r["measured_alg6_no_prior_knowledge"] for r in CD["rows"]])
fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))
ax[0].plot(P, R**2, "o", color=ACC, ms=8, label="measured $R_T^2$")
xs = np.linspace(0, P.max() * 1.05, 100)
ax[0].plot(
    xs,
    reg["slope"] * xs + reg["intercept"],
    "-",
    color=ACC2,
    label="linear fit, $r^2=%.4f$" % reg["r_squared"],
)
ax[0].set_xlabel("path length $P_T$ (withheld from the algorithm)")
ax[0].set_ylabel("$R_T^2$")
ax[0].set_title(
    r"$R_T^2$ affine in $P_T\ \Leftrightarrow\ \sqrt{P_T}$ dependence", fontsize=12
)
ax[0].legend(fontsize=10, frameon=False)

ax[1].plot(P, R, "o", color=ACC, ms=8, label="PABLO + Alg. 6 (no $P_T$)")
ax[1].plot(
    xs,
    np.sqrt(reg["slope"] * xs + reg["intercept"]),
    "-",
    color=ACC2,
    label=r"$\sqrt{a+bP_T}$",
)
ax[1].set_xlabel("path length $P_T$")
ax[1].set_ylabel("dynamic regret")
ax[1].set_title(
    "$T=%d$, $d=%d$, %d seeds/point" % (CD["T"], CD["d"], CD["seeds"]), fontsize=12
)
ax[1].legend(fontsize=10, frameon=False)
fig.savefig(os.path.join(FIG, "fig_claim3_dynamic.png"))
plt.close(fig)

# --------------------------------------------------- figure 3: Theorem 5.2 and Theorem 4.3
fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))
rows = c5s["rows"]
dv = np.array([r["d"] for r in rows], float)
ev = np.array([r["empirical_minimax_envelope"] for r in rows])
sq = np.array([r["sqrt_dT"] for r in rows])
fl = np.array([r["theorem_floor"] for r in rows])
ax[0].loglog(dv, ev, "o-", color="black", label="best of 3 algorithms (envelope)")
ax[0].loglog(dv, sq, "--", color=ACC, label=r"$\sqrt{dT}$")
ax[0].loglog(dv, fl, ":", color=ACC2, label=r"Thm 5.2 floor $\sqrt{dT}/64\wedge T/12d$")
ax[0].set_xlabel("$d$   (with $T = 500\\,d$)")
ax[0].set_ylabel("direction regret $R^Z_T$")
ax[0].set_title(
    r"Thm 5.2 construction: $d^{%.3f\pm%.3f}$ (pred. 1)"
    % (c5s["fitted_d_exponent"], c5s["stderr"]),
    fontsize=12,
)
ax[0].legend(fontsize=9, frameon=False)

Q = c4["T2_T3_T4_quantiles"]["quantiles"]
rr = [r for r in Q["rows"] if r["u_norm"] == 4.0]
x = np.array([np.sqrt(np.log(Q["T"] / r["delta"])) for r in rr])
q = np.array([r["quantile_1_minus_3delta"] for r in rr])
b = np.array([r["theorem_bound"] for r in rr])
ax[1].plot(x, q, "o-", color=ACC, label=r"measured $(1-3\delta)$ quantile")
ax[1].plot(x, b, "s--", color=ACC2, label="Theorem 4.3 bound")
ax[1].set_yscale("log")
ax[1].set_xlabel(r"$\sqrt{\log(T/\delta)}$")
ax[1].set_ylabel("regret quantile")
ax[1].set_title(
    r"high-probability bound, %d seeds, $\delta$ down to 0.002" % Q["seeds"],
    fontsize=12,
)
ax[1].legend(fontsize=10, frameon=False)
fig.savefig(os.path.join(FIG, "fig_claim45.png"))
plt.close(fig)

# ------------------------------------------------------------- figure 4: Conjecture 5.3
fig, ax = plt.subplots(figsize=(6.4, 4.2))
for dd_, col in zip([2, 4, 8], [ACC, ACC2, "black"]):
    rr = [r for r in c6["T1_achievability_shape"]["rows"] if r["d"] == dd_]
    ax.semilogx(
        [r["u_norm"] for r in rr],
        [r["ratio"] for r in rr],
        "o-",
        color=col,
        label="$d=%d$" % dd_,
    )
    ax.axvline(np.exp(dd_), color=col, ls=":", alpha=0.5)
ax.axhline(1.0, color=GREY, ls="--", lw=1)
ax.set_xlabel(r"$\|u\|$   (dotted: the transition $\log\|u\|=d$)")
ax.set_ylabel(r"measured $/\ \|u\|\sqrt{T(d\vee\log\|u\|)}$")
ax.set_title("Conjecture 5.3 (open): the ratio is bounded but decreasing", fontsize=12)
ax.legend(fontsize=10, frameon=False)
fig.savefig(os.path.join(FIG, "fig_claim6.png"))
plt.close(fig)

print("figures written to", FIG)