make_physics_evolution_figures.py

#!/usr/bin/env python
"""Generate manuscript figures and exact source-data tables."""

from __future__ import annotations

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


COLORS = {
    "Random": "#8c8c8c",
    "Greedy": "#D55E00",
    "Endpoint beam": "#0072B2",
    "Path beam": "#009E73",
    "Path beam + U": "#CC79A7",
    "Additive": "#E69F00",
    "Pairwise": "#56B4E9",
}


def main():
    out = Path("manuscript_physics_evolution/figures")
    source = out / "source_data"
    source.mkdir(parents=True, exist_ok=True)
    ablation = pd.read_csv("outputs/bnab_planner_ablations/summary.csv")
    noise = pd.read_csv("outputs/bnab_oracle_noise_robustness/summary.csv")
    learned = pd.read_csv("outputs/bnab_learned_state_potentials/shell5_prediction_metrics.csv")
    risk = pd.read_csv("outputs/global_oracle_esm650m_ensemble_audit/global_risk_coverage.csv")
    oracle = pd.read_csv("outputs/global_oracle_esm650m_ensemble_audit/matched_predictions.csv")

    methods = {
        "random": "Random",
        "greedy": "Greedy",
        "beam_endpoint_b20": "Endpoint beam",
        "beam_path_b20_lambda1": "Path beam",
    }
    panel_a = ablation[ablation.method.isin(methods)].copy()
    panel_a["display_method"] = panel_a.method.map(methods)
    panel_a.to_csv(source / "figure1a_horizon.csv", index=False)
    panel_b = ablation[(ablation.horizon == 5) & ablation.method.isin([
        "random", "greedy", "beam_endpoint_b5", "beam_endpoint_b20",
        "beam_endpoint_b50", "beam_path_b20_lambda1", "beam_path_b20_lambda2",
    ])].copy()
    panel_b.to_csv(source / "figure1b_width_penalty.csv", index=False)
    noise.to_csv(source / "figure1c_noise.csv", index=False)
    learned.to_csv(source / "figure1d_shell5.csv", index=False)

    plt.rcParams.update({
        "font.size": 8,
        "axes.titlesize": 9,
        "axes.labelsize": 8,
        "legend.fontsize": 7,
        "pdf.fonttype": 42,
        "ps.fonttype": 42,
    })
    fig, axes = plt.subplots(2, 2, figsize=(7.2, 5.6), constrained_layout=True)

    ax = axes[0, 0]
    for label in ("Random", "Greedy", "Endpoint beam", "Path beam"):
        data = panel_a[panel_a.display_method == label].sort_values("horizon")
        ax.plot(data.horizon, data.fully_monotonic_rate, marker="o", lw=1.6,
                color=COLORS[label], label=label)
    ax.set(xticks=[1, 3, 5], xlabel="Mutation horizon", ylabel="Fully monotonic trajectory rate")
    ax.set_ylim(-0.02, 1.03)
    ax.legend(frameon=False, ncol=2, loc="lower left")
    ax.set_title("a  Look-ahead prevents long-horizon failures", loc="left", fontweight="bold")

    ax = axes[0, 1]
    order = ["random", "greedy", "beam_endpoint_b5", "beam_endpoint_b20",
             "beam_endpoint_b50", "beam_path_b20_lambda1", "beam_path_b20_lambda2"]
    labels = ["Random", "Greedy", "Beam\nB=5", "Beam\nB=20", "Beam\nB=50",
              "Path\nλ=1", "Path\nλ=2"]
    data = panel_b.set_index("method").loc[order]
    colors = [COLORS["Random"], COLORS["Greedy"], COLORS["Endpoint beam"],
              COLORS["Endpoint beam"], COLORS["Endpoint beam"], COLORS["Path beam"], COLORS["Path beam"]]
    ax.bar(np.arange(len(data)), data.fully_monotonic_rate, color=colors, width=0.75)
    ax.set_xticks(np.arange(len(data)), labels, rotation=0)
    ax.set(ylabel="Fully monotonic rate", ylim=(0, 1.03))
    ax.set_title("b  Width and path-penalty ablation (T=5)", loc="left", fontweight="bold")

    ax = axes[1, 0]
    noise_methods = {
        "greedy": "Greedy",
        "beam_endpoint_b20": "Endpoint beam",
        "beam_path_b20": "Path beam",
        "beam_path_uncertainty_b20": "Path beam + U",
    }
    for method, label in noise_methods.items():
        data = noise[noise.method == method].sort_values("sigma")
        ax.plot(data.oracle_mae, data.fully_monotonic_rate, marker="o", lw=1.6,
                color=COLORS[label], label=label)
    ax.set(xlabel="Realized oracle MAE (log10 KD)", ylabel="True fully monotonic rate", ylim=(0.15, 1.02))
    ax.legend(frameon=False, ncol=2)
    ax.set_title("c  Robustness to heteroscedastic oracle noise", loc="left", fontweight="bold")

    ax = axes[1, 1]
    pivot = learned.pivot_table(index=["antibody", "antigen"], columns="model", values="mae")
    x = np.arange(len(pivot))
    for i, (_, row) in enumerate(pivot.iterrows()):
        color = COLORS["Path beam"] if row.pairwise < row.additive else COLORS["Greedy"]
        ax.plot([i, i], [row.additive, row.pairwise], color=color, lw=1.2, alpha=0.8)
    ax.scatter(x - 0.08, pivot.additive, color=COLORS["Additive"], s=18, label="Additive", zorder=3)
    ax.scatter(x + 0.08, pivot.pairwise, color=COLORS["Pairwise"], s=18, label="Pairwise", zorder=3)
    ax.set_xticks(x, [f"{a}\n{g}" for a, g in pivot.index], rotation=45, ha="right")
    ax.set(ylabel="Weight-5 shell MAE", ylim=(0, max(pivot.max()) * 1.12))
    ax.legend(frameon=False)
    ax.set_title("d  Pairwise epistasis helps 3/4 antibodies", loc="left", fontweight="bold")

    fig.savefig(out / "figure1_main.png", dpi=300)
    fig.savefig(out / "figure1_main.pdf")
    plt.close(fig)

    risk.to_csv(source / "figure2a_risk_coverage.csv", index=False)
    oracle[["target", "global_mean", "local_mean", "global_std", "complex_id"]].to_csv(
        source / "figure2b_oracle_predictions.csv", index=False
    )
    fig, axes = plt.subplots(1, 2, figsize=(7.2, 2.8), constrained_layout=True)
    ax = axes[0]
    ax.plot(risk.coverage, risk.mae, marker="o", color=COLORS["Path beam"], label="Low-U selection")
    ax.plot(risk.coverage, risk.random_expected_mae, ls="--", color=COLORS["Random"], label="No confidence ranking")
    ax.set(xlabel="Coverage", ylabel="Global-oracle MAE", xlim=(0.08, 1.02))
    ax.legend(frameon=False)
    ax.set_title("a  Ensemble uncertainty detects error", loc="left", fontweight="bold")
    ax = axes[1]
    ax.scatter(oracle.target, oracle.global_mean, s=7, alpha=0.3, color=COLORS["Path beam"], label="Global Φ")
    ax.scatter(oracle.target, oracle.local_mean, s=7, alpha=0.22, color=COLORS["Endpoint beam"], label="Local comparator")
    limits = [min(oracle.target.min(), oracle.global_mean.min()), max(oracle.target.max(), oracle.global_mean.max())]
    ax.plot(limits, limits, color="black", lw=0.8, ls="--")
    ax.set(xlabel="Experimental ΔΔG", ylabel="Predicted ΔΔG")
    ax.legend(frameon=False)
    ax.set_title("b  Consistency–precision trade-off", loc="left", fontweight="bold")
    fig.savefig(out / "figure2_oracle.png", dpi=300)
    fig.savefig(out / "figure2_oracle.pdf")
    plt.close(fig)
    print(out / "figure1_main.pdf")
    print(out / "figure2_oracle.pdf")


if __name__ == "__main__":
    main()