code/experiments/make_fig3_plot.py

"""Render the Figure 3 reproduction (mean FP32 reconstruction rel. error by
exponent, bf16 & fp16 panels) as a Plotly HTML figure for the logbook."""

import csv

import plotly.graph_objects as go
from plotly.subplots import make_subplots

# validated categorical palette (dataviz reference instance), fixed order
COLORS = {"none": "#2a78d6", "float": "#008300", "ulp_int8": "#e87ba4",
          "ulp_int16": "#eda100"}
NAMES = {"none": "No correction", "float": "BF16/FP16 + same-format float err",
         "ulp_int8": "ULP + INT8 (ours, 24-bit)", "ulp_int16": "ULP + INT16 (ours, 32-bit)"}
INK, MUTED, GRID = "#0b0b0b", "#898781", "#e1e0d9"

rows = list(csv.DictReader(open("repro_flashoptim/outputs/fig3_reconstruction.csv")))
fig = make_subplots(rows=2, cols=1, shared_xaxes=False, vertical_spacing=0.14,
                    subplot_titles=("Target: BF16 (θ′ bf16)", "Target: FP16 (θ′ fp16)"))

for r_i, target in enumerate(["bf16", "fp16"], start=1):
    for m in ["none", "float", "ulp_int8", "ulp_int16"]:
        pts = [(int(r["exp_unbiased"]), float(r["mean_rel_err"]), int(r["n"]))
               for r in rows if r["target"] == target and r["method"] == m
               and int(r["exp_biased"]) > 0 and float(r["mean_rel_err"]) > 0]
        pts.sort()
        fig.add_trace(go.Scatter(
            x=[p[0] for p in pts], y=[p[1] for p in pts], mode="lines",
            name=NAMES[m], legendgroup=m, showlegend=(r_i == 1),
            line=dict(color=COLORS[m], width=2),
            hovertemplate="exp 2^%{x}: mean rel err %{y:.2e}<extra>" + NAMES[m] + "</extra>"),
            row=r_i, col=1)
    # denormal boundaries (dotted verticals, as in the paper)
    for xb in ([-126] if target == "bf16" else [-14, 15]):
        fig.add_vline(x=xb, line_dash="dot", line_color=MUTED, line_width=1, row=r_i, col=1)

fig.update_yaxes(type="log", title_text="mean relative error", exponentformat="power",
                 gridcolor=GRID, zeroline=False, color=MUTED, title_font_color=INK)
fig.update_xaxes(title_text="fp32 value exponent (unbiased)", gridcolor=GRID,
                 zeroline=False, color=MUTED, title_font_color=INK)
fig.update_layout(
    title=dict(text="FP32 reconstruction error by exponent — exhaustive over all 4.28e9 finite FP32 values<br>"
                    "<sup>Reproduction of FlashOptim Fig. 3 protocol (independent reimplementation of Alg. 1)</sup>",
               font=dict(color=INK, size=15)),
    font=dict(family="system-ui, -apple-system, 'Segoe UI', sans-serif", color=INK),
    paper_bgcolor="#fcfcfb", plot_bgcolor="#fcfcfb",
    legend=dict(orientation="h", yanchor="bottom", y=1.03, x=0, font=dict(size=11)),
    height=640, margin=dict(l=70, r=20, t=110, b=50), hovermode="x unified")
fig.write_html("repro_flashoptim/outputs/fig3_reconstruction.html", include_plotlyjs="cdn")
print("wrote repro_flashoptim/outputs/fig3_reconstruction.html")