scripts/make_figs.py

"""Figures for the reproduction poster / logbook."""

import json
import os
import sys

import matplotlib

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

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "..", "outputs")
FIG = os.path.join(HERE, "..", "figs")
os.makedirs(FIG, exist_ok=True)
ACC = "#B5453C"
ACC2 = "#1F4E79"
GREY = "#6B7280"
plt.rcParams.update(
    {
        "font.size": 11,
        "axes.grid": True,
        "grid.alpha": 0.25,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "figure.dpi": 190,
        "savefig.bbox": "tight",
    }
)


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


# ---------------------------------------------------------------- fig 1
def fig_mechanism():
    fig, ax = plt.subplots(figsize=(7.4, 3.5))
    ax.axis("off")
    ax.set_xlim(0, 10)
    ax.set_ylim(0, 6)
    ax.text(
        5,
        5.6,
        "one +C / -C unit pair, two accumulation orders",
        ha="center",
        fontsize=13,
        fontweight="bold",
        color="#111827",
    )

    def chain(y, items, colour, label):
        ax.text(0.15, y + 0.62, label, fontsize=11, fontweight="bold", color=colour)
        x = 0.4
        for i, (txt, res) in enumerate(items):
            ax.add_patch(
                plt.Rectangle(
                    (x, y - 0.32), 1.75, 0.72, fc="white", ec=colour, lw=1.6, zorder=2
                )
            )
            ax.text(x + 0.875, y + 0.19, txt, ha="center", fontsize=10.5, zorder=3)
            ax.text(
                x + 0.875,
                y - 0.16,
                res,
                ha="center",
                fontsize=10.5,
                color=colour,
                fontweight="bold",
                zorder=3,
            )
            if i < len(items) - 1:
                ax.annotate(
                    "",
                    xy=(x + 2.05, y + 0.04),
                    xytext=(x + 1.78, y + 0.04),
                    arrowprops=dict(arrowstyle="->", color=GREY, lw=1.4),
                )
            x += 2.3

    chain(
        3.4,
        [
            ("start", "t"),
            ("(+) C", "C  (t lost)"),
            ("(-) C", "0"),
            ("(+) u", "u"),
        ],
        ACC,
        "AD accumulation of grad(x)",
    )
    chain(
        1.1,
        [("start", "0"), ("(+) C", "C"), ("(-) C", "0"), ("(+) v", "v")],
        ACC2,
        "forward accumulation",
    )
    ax.text(
        5,
        0.05,
        "identical units, opposite downstream weights: the pair is invisible to f "
        "and lethal to D",
        ha="center",
        fontsize=10.5,
        color="#374151",
        style="italic",
    )
    fig.savefig(os.path.join(FIG, "fig1_mechanism.png"))
    plt.close(fig)


# ---------------------------------------------------------------- fig 2
def fig_exactness():
    rows = load("claim12_theorem31.json")
    acts = ["relu", "elu", "gelu", "swish", "sigmoid", "tanh"]
    dts = ["float32", "float64"]
    M = np.zeros((2, 6))
    for r in rows:
        M[dts.index(r["dtype"]), acts.index(r["activation"])] += (
            r["value_mismatches"] + r["grad_mismatches"]
        )
    fig, ax = plt.subplots(figsize=(7.4, 2.5))
    ax.imshow(M, cmap="Greens_r", vmin=0, vmax=1, aspect="auto")
    for i in range(2):
        for j in range(6):
            ax.text(
                j,
                i,
                "0",
                ha="center",
                va="center",
                fontsize=15,
                fontweight="bold",
                color="#14532d",
            )
    ax.set_xticks(range(6))
    ax.set_xticklabels(["ReLU", "ELU", "GELU", "Swish", "Sigmoid", "tanh"])
    ax.set_yticks(range(2))
    ax.set_yticklabels(["binary32", "binary64"])
    ax.set_title(
        "Theorem 3.1, L = 9: mismatches out of 33 domain points x 3 seeds\n"
        "(value AND gradient, bit-exact)",
        fontsize=11.5,
    )
    ax.grid(False)
    fig.savefig(os.path.join(FIG, "fig2_exactness.png"))
    plt.close(fig)


# ---------------------------------------------------------------- fig 3
def fig_lemma34_boundary():
    rows = [r for r in load("claim4_lemma34.json") if r["dtype"] == "float32"]
    acts = [r["activation"] for r in rows]
    ok = [np.log2(r["largest_abs_y_still_suppressed"]) for r in rows]
    leak = [np.log2(r["smallest_abs_y_that_leaks"]) for r in rows]
    x = np.arange(len(acts))
    fig, ax = plt.subplots(figsize=(7.4, 3.1))
    ax.bar(x - 0.19, ok, 0.38, color=ACC2, label="largest |y| still fully suppressed")
    ax.bar(x + 0.19, leak, 0.38, color=ACC, label="smallest |y| that leaks")
    ax.axhline(127, ls="--", color=GREY, lw=1.2)
    ax.text(
        len(acts) - 0.5, 122, "$\\Omega = 2^{128}$", ha="right", fontsize=10, color=GREY
    )
    ax.set_xticks(x)
    ax.set_xticklabels(
        [a.upper() if a in ("relu", "elu") else a.capitalize() for a in acts]
    )
    ax.set_ylabel("$\\log_2 |y|$")
    ax.set_title(
        "Lemma 3.4 (binary32): input-gradient range over which the AD\n"
        "gradient is exactly 0 while f = f* stays bit-exact",
        fontsize=11.5,
    )
    ax.legend(fontsize=9.5, loc="upper left")
    fig.savefig(os.path.join(FIG, "fig3_lemma34_range.png"))
    plt.close(fig)


# ---------------------------------------------------------------- fig 4
def fig_thm32():
    try:
        d = load("claim3_theorem32.json")["y_dependent_gradient"]
    except Exception:
        return
    ch = [c for c in d["checks"] if c["y"] > 0]
    ys = np.array([c["y"] for c in ch])
    gs = np.array([abs(c["got"]) for c in ch])
    order = np.argsort(ys)
    ys, gs = ys[order], gs[order]
    fig, ax = plt.subplots(figsize=(7.4, 3.3))
    ax.loglog(
        ys,
        gs,
        "o-",
        color=ACC,
        lw=1.8,
        ms=5,
        label="IEEE-754 AD gradient $|D_{f,x}(y)|$",
    )
    ref = gs[0] * ys / ys[0]
    ax.loglog(
        ys,
        ref,
        "--",
        color=ACC2,
        lw=1.8,
        label="exact real arithmetic (chain rule): $\\propto y$",
    )
    ax.set_xlabel("input gradient $y = \\varphi'_x(f(x))$")
    ax.set_ylabel("$|D_{f,x}(y)|$")
    ax.set_title(
        "Theorem 3.2: one network, gradient prescribed on 24 bands\n"
        "f(x) = 0 everywhere; $|D/y|$ spans "
        f"{d['ratio_spread_orders']:.1f} orders of magnitude",
        fontsize=11.5,
    )
    ax.legend(fontsize=9.5)
    fig.savefig(os.path.join(FIG, "fig4_thm32.png"))
    plt.close(fig)


if __name__ == "__main__":
    fig_mechanism()
    fig_exactness()
    fig_lemma34_boundary()
    fig_thm32()
    print("figures:", sorted(os.listdir(FIG)))