make_plots.py

# /// script
# requires-python = ">=3.11"
# dependencies = ["pandas", "plotly", "matplotlib", "kaleido"]
# ///
"""Build logbook figures (interactive Plotly HTML) and poster figures (PNG) from
the experiment's results.csv / lp_quality_gap.csv / scaling.csv."""
import argparse
from pathlib import Path

import pandas as pd
import plotly.graph_objects as go
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

PIPELINE_LABEL = {
    "oracle__lagrangian": "Oracle (true effects)",
    "buoplr_net__lagrangian": "BUOPLR (proposed)",
    "independent_mlp__lagrangian": "Indep. per-outcome net + Lagrangian",
    "linear_additive__lagrangian": "Linear-additive + Lagrangian",
    "buoplr_net__greedy": "BUOPLR net + greedy",
    "independent_mlp__greedy": "Indep. net + greedy",
    "linear_additive__greedy": "Linear-additive + greedy",
    "control_no_treatment": "No-treatment control",
}
COLORS = {
    "oracle__lagrangian": "#9AA5B1",
    "buoplr_net__lagrangian": "#B8860B",
    "independent_mlp__lagrangian": "#6E7B8B",
    "linear_additive__lagrangian": "#B0B8C1",
    "buoplr_net__greedy": "#D9B96A",
    "independent_mlp__greedy": "#AEB7C0",
    "linear_additive__greedy": "#C7CDD3",
    "control_no_treatment": "#E3E6E9",
}


def fig_pipeline_comparison(results_csv: Path, out_html: Path, out_png: Path):
    df = pd.read_csv(results_csv)
    df = df[df["pipeline"] != "control_no_treatment"].copy()
    df["label"] = df["pipeline"].map(lambda p: PIPELINE_LABEL.get(p, p))
    df = df.sort_values("total_main_uplift", ascending=True)
    colors = [COLORS.get(p, "#888") for p in df["pipeline"]]

    fig = go.Figure(go.Bar(
        x=df["total_main_uplift"], y=df["label"], orientation="h",
        marker_color=colors,
        text=[f"{v:.1f}" for v in df["total_main_uplift"]], textposition="outside",
    ))
    fig.update_layout(
        title="Realized incremental main-outcome uplift by pipeline (synthetic proxy)",
        xaxis_title="Total realized main-outcome uplift on held-out users",
        template="plotly_white", height=480, margin=dict(l=260, r=40, t=60, b=50),
    )
    fig.write_html(str(out_html), include_plotlyjs="cdn")

    fig_mpl, ax = plt.subplots(figsize=(13.0, 8.7), dpi=380)  # AR ~1.49 to match poster hero-stage
    ax.barh(df["label"], df["total_main_uplift"], color=colors, edgecolor="#33363b", linewidth=0.6)
    for i, v in enumerate(df["total_main_uplift"]):
        ax.text(v, i, f" {v:.0f}", va="center", fontsize=9)
    ax.set_xlabel("Total realized main-outcome uplift (held-out users)", fontsize=9)
    ax.set_title("BUOPLR vs baselines — offline pipeline comparison", fontsize=11)
    ax.tick_params(labelsize=8)
    ax.set_xlim(0, df["total_main_uplift"].max() * 1.15)
    plt.tight_layout()
    fig_mpl.savefig(out_png)
    plt.close(fig_mpl)


def fig_lp_quality_gap(lp_csv: Path, out_html: Path, out_png: Path):
    df = pd.read_csv(lp_csv)
    labels = {"linear_additive": "Linear-additive", "independent_mlp": "Independent per-outcome net",
              "buoplr_net": "BUOPLR net"}
    df["label"] = df["uplift_model"].map(lambda m: labels.get(m, m))

    fig = go.Figure()
    fig.add_bar(name="Restricted Lagrangian (BUOPLR stage-2)", x=df["label"], y=df["lagrangian_uplift"],
                marker_color="#B8860B")
    fig.add_bar(name="Exact LP (upper bound, not scalable)", x=df["label"], y=df["lp_uplift"],
                marker_color="#9AA5B1")
    fig.update_layout(barmode="group", template="plotly_white",
                       title=f"Assignment quality gap at n={int(df['n_users'].iloc[0])} users: Lagrangian relaxation vs exact LP",
                       yaxis_title="Total realized main-outcome uplift", height=440,
                       margin=dict(l=60, r=40, t=70, b=50))
    fig.write_html(str(out_html), include_plotlyjs="cdn")

    fig_mpl, ax = plt.subplots(figsize=(6.5, 4.0), dpi=200)
    x = range(len(df))
    w = 0.35
    ax.bar([i - w / 2 for i in x], df["lagrangian_uplift"], width=w, label="Restricted Lagrangian\n(BUOPLR stage-2)", color="#B8860B")
    ax.bar([i + w / 2 for i in x], df["lp_uplift"], width=w, label="Exact LP\n(not scalable)", color="#9AA5B1")
    ax.set_xticks(list(x))
    ax.set_xticklabels(df["label"], fontsize=8)
    ax.set_ylabel("Realized main-outcome uplift")
    ax.set_title(f"Lagrangian vs exact LP @ n={int(df['n_users'].iloc[0])}")
    ax.legend(fontsize=8)
    plt.tight_layout()
    fig_mpl.savefig(out_png, dpi=200)
    plt.close(fig_mpl)


def fig_scaling(scaling_csv: Path, out_html: Path, out_png: Path):
    df = pd.read_csv(scaling_csv).sort_values("n_users")

    fig = go.Figure()
    fig.add_trace(go.Scatter(x=df["n_users"], y=df["lagrangian_assign_s"], mode="lines+markers",
                              name="BUOPLR restricted Lagrangian", line=dict(color="#B8860B")))
    fig.add_trace(go.Scatter(x=df["n_users"], y=df["greedy_assign_s"], mode="lines+markers",
                              name="Greedy myopic", line=dict(color="#6E7B8B")))
    fig.update_layout(template="plotly_white", title="Assignment-stage runtime vs. number of users",
                       xaxis_title="Users", yaxis_title="Assignment runtime (s)",
                       xaxis_type="log", height=440, margin=dict(l=60, r=40, t=60, b=50))
    fig.write_html(str(out_html), include_plotlyjs="cdn")

    fig_mpl, ax = plt.subplots(figsize=(6.5, 4.0), dpi=200)
    ax.plot(df["n_users"], df["lagrangian_assign_s"], "o-", color="#B8860B", label="BUOPLR Lagrangian")
    ax.plot(df["n_users"], df["greedy_assign_s"], "o-", color="#6E7B8B", label="Greedy myopic")
    ax.set_xscale("log")
    ax.set_xlabel("Users (log scale)")
    ax.set_ylabel("Assignment runtime (s)")
    ax.set_title("BUOPLR assignment scales near-linearly")
    ax.legend(fontsize=9)
    plt.tight_layout()
    fig_mpl.savefig(out_png, dpi=200)
    plt.close(fig_mpl)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--in-dir", type=str, required=True)
    ap.add_argument("--out-dir", type=str, required=True)
    args = ap.parse_args()
    in_dir, out_dir = Path(args.in_dir), Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    fig_pipeline_comparison(in_dir / "results.csv", out_dir / "pipeline_comparison.html",
                             out_dir / "pipeline_comparison.png")
    fig_lp_quality_gap(in_dir / "lp_quality_gap.csv", out_dir / "lp_quality_gap.html",
                        out_dir / "lp_quality_gap.png")
    fig_scaling(in_dir / "scaling.csv", out_dir / "scaling.html", out_dir / "scaling.png")
    print(f"Wrote figures to {out_dir}")


if __name__ == "__main__":
    main()