scripts/make_figs.py

"""Build every figure in figs/ from the JSON outputs in outputs/."""

from __future__ import annotations

import json
import os

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

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

ACC = "#1f6f8b"
ACC2 = "#e07a5f"
ACC3 = "#3d5a80"
GREY = "#8d99ae"
plt.rcParams.update(
    {
        "figure.dpi": 150,
        "font.size": 9,
        "axes.spines.top": False,
        "axes.spines.right": False,
    }
)


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


# ------------------------------------------------------------------- claim 1
def fig1():
    d = load("claim1_pvalue_validity.json")
    fig, ax = plt.subplots(1, 3, figsize=(9.5, 3.0))

    A = d["A_oracle_uniformity"]
    ns = [30, 100, 300]
    for k, test in enumerate(("wilcoxon", "sign")):
        v = [A[f"n{n}|{test}|perm"]["type_I_at_0.05"] for n in ns]
        ax[0].plot(ns, v, "o-", color=[ACC, ACC2][k], label=f"SKO {test} (perm)")
        v = [A[f"n{n}|{test}|iid"]["type_I_at_0.05"] for n in ns]
        ax[0].plot(
            ns,
            v,
            "s--",
            color=[ACC, ACC2][k],
            alpha=0.55,
            label=f"SKO {test} (iid draws)",
        )
    ax[0].axhline(0.05, color="k", ls=":", lw=1, label="nominal 0.05")
    ax[0].set_xscale("log")
    ax[0].set_xlabel("n")
    ax[0].set_ylabel("type-I error at $\\alpha=0.05$")
    ax[0].set_title("Oracle $\\nu_j,\\rho_j$: null p-values", fontsize=9)
    ax[0].set_ylim(0, 0.09)
    ax[0].legend(fontsize=6, frameon=False)

    B = d["B_split_vs_nosplit"]
    names = ["SKO_Wcx_nosplit", "HRT_nosplit", "HRT_split", "HRT_split_oracle_sampler"]
    lab = [
        "SKO-Wcx\n(no split)",
        "HRT\n(no split)",
        "HRT\n(50/50 split)",
        "HRT split\n(oracle sampler)",
    ]
    v = [B[k]["type_I_at_0.05"] for k in names]
    c = [ACC, ACC2, ACC3, GREY]
    ax[1].bar(range(4), v, color=c)
    ax[1].axhline(0.05, color="k", ls=":", lw=1)
    ax[1].set_xticks(range(4))
    ax[1].set_xticklabels(lab, fontsize=6)
    ax[1].set_ylabel("type-I error at $\\alpha=0.05$")
    ax[1].set_title("Estimated $\\hat\\nu,\\hat\\rho$, GB black box", fontsize=9)
    for i, x in enumerate(v):
        ax[1].text(i, x + 0.02, f"{x:.3f}", ha="center", fontsize=6)

    C = d["C_heavy_tailed_elliptical"]
    v2 = [
        C["SKO_Wcx_oracle_nu"]["type_I_at_0.05"],
        C["HRT_split_gaussian_sampler"]["type_I_at_0.05"],
    ]
    ax[2].bar([0, 1], v2, color=[ACC, ACC2])
    ax[2].axhline(0.05, color="k", ls=":", lw=1)
    ax[2].set_xticks([0, 1])
    ax[2].set_xticklabels(
        ["SKO-Wcx\n(cond. mean only)", "HRT split\n(Gaussian knockoff\nsampler)"],
        fontsize=6,
    )
    ax[2].set_ylabel("type-I error at $\\alpha=0.05$")
    ax[2].set_title("Heavy-tailed $t_3$ elliptical design", fontsize=9)
    for i, x in enumerate(v2):
        ax[2].text(i, x + 0.004, f"{x:.3f}", ha="center", fontsize=6)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "claim1_pvalue_validity.png"), bbox_inches="tight")
    plt.close(fig)


# ------------------------------------------------------------------- claim 2
def fig2():
    d = load("claim2_fdr.json")
    qs = d["q_grid"]
    fig, ax = plt.subplots(1, 2, figsize=(7.5, 3.0))
    keys = [
        ("oracle_gb", "oracle $\\nu,\\rho$ + GB", ACC, "o-"),
        ("oracle_lasso", "oracle $\\nu,\\rho$ + Lasso", ACC3, "s-"),
        ("estimated_gb", "estimated $\\hat\\nu,\\hat\\rho$ + GB", ACC2, "^-"),
        ("estimated_gb_nonlinear", "mis-specified imputer", GREY, "v--"),
    ]
    ax[0].plot([0, 0.55], [0, 0.55], "k:", lw=1, label="nominal $q$")
    for k, lab, col, st in keys:
        f = [d[k][str(q)]["FDR"] for q in qs]
        e = [1.96 * d[k][str(q)]["FDR_se"] for q in qs]
        ax[0].errorbar(qs, f, yerr=e, fmt=st, color=col, label=lab, ms=4, lw=1.2)
    ax[0].set_xlabel("nominal level $q$")
    ax[0].set_ylabel("realised FDR")
    ax[0].set_title("Theorem 3.4: FDR$(S_{SKO}) \\leq q$", fontsize=9)
    ax[0].legend(fontsize=6, frameon=False)

    for k, lab, col, st in keys:
        pw = [d[k][str(q)]["power"] for q in qs]
        ax[1].plot(qs, pw, st, color=col, label=lab, ms=4, lw=1.2)
    ax[1].set_xlabel("nominal level $q$")
    ax[1].set_ylabel("power")
    ax[1].set_title("Power at the knockoff threshold", fontsize=9)
    ax[1].set_ylim(-0.03, 1.05)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "claim2_fdr.png"), bbox_inches="tight")
    plt.close(fig)


# ------------------------------------------------------------------- claim 3
def fig3():
    d = load("claim3_stability.json")
    fig, ax = plt.subplots(1, 3, figsize=(9.5, 3.0))
    ex = d["figure1"]["example"]
    imp = np.array(ex["important"], bool)
    idx = np.arange(len(imp))
    ax[0].scatter(
        idx[~imp],
        np.array(ex["ridge"])[~imp],
        s=14,
        color=ACC,
        label="Ridge, unimportant",
    )
    ax[0].scatter(
        idx[imp], np.array(ex["ridge"])[imp], s=14, color=ACC2, label="Ridge, important"
    )
    ax[0].scatter(
        idx[~imp],
        np.array(ex["lasso"])[~imp],
        s=22,
        marker="x",
        color=ACC,
        alpha=0.6,
        label="Lasso, unimportant",
    )
    ax[0].scatter(
        idx[imp],
        np.array(ex["lasso"])[imp],
        s=22,
        marker="x",
        color=ACC2,
        alpha=0.6,
        label="Lasso, important",
    )
    ax[0].set_xlabel("Variable index $j$")
    ax[0].set_ylabel(r"$\|\tilde\theta^j-\hat\theta\|_2$")
    ax[0].set_title("Figure 1 reproduction ($n$=300, $p$=50)", fontsize=9)
    ax[0].legend(fontsize=5.5, frameon=False)

    for key, lab, col in (
        ("rate_ridge_lam1e-3", "Ridge $\\lambda=10^{-3}$", ACC),
        ("rate_ridge_lam1e-1", "Ridge $\\lambda=10^{-1}$", ACC3),
        ("rate_lasso_lam1e-2", "Lasso $\\lambda=10^{-2}$", ACC2),
        ("rate_ridge_heavy_t2", "$t_2$ design (E.2 broken)", GREY),
    ):
        r = d[key]
        ns = r["ns"]
        q = r["quantiles"]["0.05"]["q_values"]
        ax[1].loglog(
            ns,
            q,
            "o-",
            color=col,
            ms=3.5,
            lw=1.2,
            label=f"{lab}: {r['quantiles']['0.05']['fitted_exponent']:+.2f}",
        )
    ref = np.array(d["rate_ridge_lam1e-3"]["ns"], float)
    ax[1].loglog(ref, 3.5 * ref**-0.5, "k:", lw=1, label="slope $-1/2$")
    ax[1].set_xlabel("n")
    ax[1].set_ylabel("$Q_{0.95}$ of $\\|\\tilde\\theta^j-\\hat\\theta\\|_2$")
    ax[1].set_title("Rate in $n$ (null $j$)", fontsize=9)
    ax[1].legend(fontsize=5.5, frameon=False)

    nr = d["nu_rho_coefficients"]
    ax[2].loglog(
        nr["ns"],
        nr["null_q95"],
        "o-",
        color=ACC,
        ms=4,
        label=f"null $j$: {nr['null_fitted_exponent']:+.2f}",
    )
    ax[2].loglog(
        nr["ns"],
        nr["important_q95"],
        "s-",
        color=ACC2,
        ms=4,
        label=f"important $j$: {nr['important_fitted_exponent']:+.2f}",
    )
    r = np.array(nr["ns"], float)
    ax[2].loglog(r, 7.0 * r**-0.5, "k:", lw=1, label="slope $-1/2$")
    ax[2].set_xlabel("n")
    ax[2].set_ylabel(r"$Q_{0.95}\|\hat\theta_{\hat\rho}-\hat\theta_{\hat\nu}\|_2$")
    ax[2].set_title(r"$\hat\nu$ vs $\hat\rho$ coefficients", fontsize=9)
    ax[2].legend(fontsize=6, frameon=False)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "claim3_stability.png"), bbox_inches="tight")
    plt.close(fig)


# ------------------------------------------------------------------- claim 4
def fig4():
    d = load("claim4_double_robustness.json")
    fig, ax = plt.subplots(1, 4, figsize=(12, 2.9))
    for k, (kind, lab) in enumerate(
        (("rf", "Random Forest"), ("nn", "Neural Network"), ("gb", "Gradient Boosting"))
    ):
        b = np.array(d["figure3"][kind]["blue_sample"])
        o = np.array(d["figure3"][kind]["orange_sample"])
        lim = np.quantile(np.abs(b), 0.995)
        bins = np.linspace(-lim, lim, 61)
        ax[k].hist(
            b,
            bins=bins,
            color=ACC,
            alpha=0.75,
            label=r"$l(\hat m(\tilde X'_1),y)-l(\hat m(\tilde X'_2),y)$",
        )
        ax[k].hist(
            o,
            bins=bins,
            color=ACC2,
            alpha=0.75,
            label=r"$l(\hat m(\tilde X'_1),y)-l(\hat m(\tilde X_1),y)$",
        )
        ax[k].set_title(
            f"$\\hat m$ = {lab}\nsd ratio "
            f"{d['figure3'][kind]['concentration_ratio_std']:.1f}$\\times$",
            fontsize=8,
        )
        ax[k].set_xlabel("Loss difference value")
        if k == 0:
            ax[k].set_ylabel("Frequency")
            ax[k].legend(fontsize=5.5, frameon=False)

    for key, lab, col, st in (
        ("rate_linear_null", "linear, null $j$", ACC, "o-"),
        ("rate_linear_important", "linear, important $j$", ACC2, "s-"),
        ("rate_nn_null", "NN, null $j$", ACC3, "^-"),
    ):
        if key not in d:
            continue
        r = d[key]
        ax[3].loglog(
            r["ns"],
            r["D_n_q90_abs"],
            st,
            color=col,
            ms=3.5,
            lw=1.2,
            label=f"{lab}: {r['exponent_D_n_q90_abs']:+.2f}",
        )
    ref = np.array(d["rate_linear_null"]["ns"], float)
    ax[3].loglog(
        ref,
        ref**-1.0 * d["rate_linear_null"]["D_n_q90_abs"][0] * ref[0],
        "k:",
        lw=1,
        label="slope $-1$",
    )
    ax[3].loglog(
        ref,
        ref**-0.5 * d["rate_linear_important"]["D_n_q90_abs"][0] * ref[0] ** 0.5,
        "k--",
        lw=1,
        label="slope $-1/2$",
    )
    ax[3].set_xlabel("n")
    ax[3].set_ylabel("$Q_{0.9}|l(\\hat m(\\tilde X'))-l(\\hat m(\\tilde X))|$")
    ax[3].set_title("Compound rate $a_n b_n$", fontsize=8)
    ax[3].legend(fontsize=5.5, frameon=False)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "claim4_double_robustness.png"), bbox_inches="tight")
    plt.close(fig)


# ------------------------------------------------------------------- claim 5
def fig5():
    d = load("claim5_power.json")
    meth = ["LOCO_Wcx", "HRT", "SKO_Wcx", "SKO_Wcx_p5"]
    cols = [GREY, ACC3, ACC, ACC2]
    panels = [
        ("figure4_adjacent_gb", "Fig. 4 - adjacent support (GB)"),
        ("figure5_masked_nn", "Fig. 5 - masked correlation (NN)"),
    ]
    fig, ax = plt.subplots(2, 3, figsize=(9.5, 5.2))
    for r, (key, title) in enumerate(panels):
        dd = d[key]
        for c, met in enumerate(("power", "type_I", "auc")):
            v = [dd[f"{m}|{met}"] for m in meth]
            e = [1.96 * dd[f"{m}|{met}_se"] for m in meth]
            ax[r, c].barh(
                range(4), v, xerr=e, color=cols, height=0.65, error_kw={"lw": 0.8}
            )
            ax[r, c].set_yticks(range(4))
            ax[r, c].set_yticklabels(meth, fontsize=7)
            ax[r, c].set_xlabel(
                {"power": "Power", "type_I": "Type-I Error", "auc": "AUC"}[met],
                fontsize=8,
            )
            if met == "type_I":
                ax[r, c].axvline(0.05, color="k", ls=":", lw=1)
            for i, x in enumerate(v):
                ax[r, c].text(x, i, f" {x:.3f}", va="center", fontsize=6)
            if c == 0:
                ax[r, c].set_title(title, fontsize=9, loc="left")
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "claim5_power.png"), bbox_inches="tight")
    plt.close(fig)


# ------------------------------------------------------------------- claim 6
def fig6():
    d = load("claim6_wdbc.json")
    kinds = [
        ("rf", "Random Forest"),
        ("nn", "Neural Network"),
        ("gb", "Gradient Boosting"),
    ]
    meth = ["HRT_split", "SKO_Wcx", "SKO_Wcx_p5"]
    cols = [ACC3, ACC, ACC2]
    fig, ax = plt.subplots(1, 4, figsize=(11.5, 3.0))
    for k, (kind, lab) in enumerate(kinds):
        v = [d[kind][m]["mean_discoveries"] for m in meth]
        e = [d[kind][m]["sd_discoveries"] for m in meth]
        t1 = [d[kind][m]["artificial_null_rejection_rate"] for m in meth]
        ax[k].barh(range(3), v, xerr=e, color=cols, height=0.6, error_kw={"lw": 0.8})
        ax[k].set_yticks(range(3))
        ax[k].set_yticklabels(meth, fontsize=7)
        ax[k].set_xlabel("Discoveries (of 30 real features)", fontsize=8)
        ax[k].set_title(f"$\\hat m$ = {lab}", fontsize=9)
        for i, (x, tt) in enumerate(zip(v, t1)):
            ax[k].text(x, i, f"  {x:.1f}   (type-I {tt:.2f})", va="center", fontsize=6)
        ax[k].set_xlim(0, max(v) * 1.9)
    cm = d["cross_model"]["per_feature_rate_by_model"]
    names = [n for n in cm if n != "ARTIFICIAL_NULL"]
    M = np.array([cm[n] for n in names])
    order = np.argsort(-M.mean(1))
    im = ax[3].imshow(M[order], aspect="auto", cmap="Blues", vmin=0, vmax=1)
    ax[3].set_xticks(range(3))
    ax[3].set_xticklabels(["RF", "NN", "GB"], fontsize=7)
    ax[3].set_yticks(range(len(names)))
    ax[3].set_yticklabels([names[i][:22] for i in order], fontsize=4.5)
    ax[3].set_title("SKO-Wcx-p5 rejection rate\nacross seeds", fontsize=8)
    fig.colorbar(im, ax=ax[3], fraction=0.04)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "claim6_wdbc.png"), bbox_inches="tight")
    plt.close(fig)



def fig5_poster():
    """Compact single-row version of the Figure-4 panel for the poster."""
    d = load("claim5_power.json")
    meth = ["LOCO_Wcx", "HRT", "SKO_Wcx", "SKO_Wcx_p5"]
    cols = [GREY, ACC3, ACC, ACC2]
    dd = d["figure4_adjacent_gb"]
    fig, ax = plt.subplots(1, 3, figsize=(9.5, 2.5))
    for c, met in enumerate(("power", "type_I", "auc")):
        v = [dd[f"{m}|{met}"] for m in meth]
        e = [1.96 * dd[f"{m}|{met}_se"] for m in meth]
        ax[c].barh(range(4), v, xerr=e, color=cols, height=.65,
                   error_kw={"lw": .8})
        ax[c].set_yticks(range(4))
        ax[c].set_yticklabels(meth, fontsize=7)
        ax[c].set_xlabel({"power": "Power", "type_I": "Type-I Error",
                          "auc": "AUC"}[met], fontsize=8)
        if met == "type_I":
            ax[c].axvline(0.05, color="k", ls=":", lw=1)
        for i, x in enumerate(v):
            ax[c].text(x, i, f" {x:.3f}", va="center", fontsize=6)
    ax[0].set_title("Fig. 4 - adjacent support, gradient boosting, n=300, p=50, 50 reps",
                    fontsize=9, loc="left")
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "claim5_power_poster.png"), bbox_inches="tight")
    plt.close(fig)


# ------------------------------------------------------------------ claim 5b
def fig5b():
    d = load("claim5b_derandomization.json")
    c = load("claim5c_derand_validity.json")
    perms = d["permutation_grid"]
    ns = [50, 75, 100, 150, 300]
    fig, ax = plt.subplots(1, 3, figsize=(9.5, 3.0))
    cmap = plt.cm.viridis(np.linspace(0.1, 0.85, len(ns)))
    for i, n in enumerate(ns):
        v = [d[f"n{n}"][f"perm{k}"]["power"] for k in perms]
        ax[0].plot(perms, v, "o-", color=cmap[i], ms=4, label=f"n={n}")
        t = [d[f"n{n}"][f"perm{k}"]["type_I"] for k in perms]
        ax[1].plot(perms, t, "o-", color=cmap[i], ms=4, label=f"n={n}")
    ax[0].set_xscale("log")
    ax[0].set_xticks(perms)
    ax[0].set_xticklabels(perms)
    ax[0].set_xlabel("permutations (derandomisation)")
    ax[0].set_ylabel("power")
    ax[0].set_title("Masked correlation: power gain", fontsize=9)
    ax[0].legend(fontsize=6, frameon=False)
    ax[1].axhline(0.05, color="k", ls=":", lw=1)
    ax[1].set_xscale("log")
    ax[1].set_xticks(perms)
    ax[1].set_xticklabels(perms)
    ax[1].set_xlabel("permutations (derandomisation)")
    ax[1].set_ylabel("type-I error at $\\alpha=0.05$")
    ax[1].set_title("...bought with type-I inflation", fontsize=9)

    for key, lab, col in (
        ("oracle_n300", "oracle $\\nu=\\rho$, n=300", ACC),
        ("estimated_n100", "estimated, n=100", ACC2),
        ("estimated_n300", "estimated, n=300", ACC3),
        ("estimated_n1000", "estimated, n=1000", GREY),
    ):
        if key not in c:
            continue
        v = [c[key][f"perm{k}"]["type_I_at_0.05"] for k in perms]
        ax[2].plot(perms, v, "o-", color=col, ms=4, label=lab)
    ax[2].axhline(0.05, color="k", ls=":", lw=1)
    ax[2].set_xscale("log")
    ax[2].set_xticks(perms)
    ax[2].set_xticklabels(perms)
    ax[2].set_xlabel("permutations")
    ax[2].set_ylabel("type-I error at $\\alpha=0.05$")
    ax[2].set_title("Global null: the inflation is\nan ESTIMATION effect", fontsize=9)
    ax[2].legend(fontsize=6, frameon=False)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "claim5b_derandomization.png"), bbox_inches="tight")
    plt.close(fig)


if __name__ == "__main__":
    for fn in (fig1, fig2, fig3, fig4, fig5, fig5_poster, fig5b, fig6):
        try:
            fn()
            print("ok", fn.__name__)
        except Exception as exc:  # noqa: BLE001
            print("FAILED", fn.__name__, type(exc).__name__, exc)