SparseWake/sparsewake

1from __future__ import annotations2 3from pathlib import Path4 5import matplotlib.pyplot as plt6import pandas as pd7 8 9COLORS = {"wake+potential": "#325AA8", "wake only": "#D55E00", "potential only": "#009E73"}10 11 12def save_pdf_svg(fig: plt.Figure, out: Path, name: str) -> None:13    out.mkdir(parents=True, exist_ok=True)14    fig.savefig(out / f"{name}.pdf", bbox_inches="tight")15    fig.savefig(out / f"{name}.svg", bbox_inches="tight")16    plt.close(fig)17 18 19def plot_history_sensor_noise(results_dir: Path, out: Path) -> None:20    main = pd.read_csv(results_dir / "main_numbers.csv")21    hist = main[main["experiment"] == "history_sweep"].copy()22    sensor = pd.read_csv(results_dir / "sensor_ablation.csv")23    noise = pd.read_csv(results_dir / "component_noise.csv")24    fig, axes = plt.subplots(1, 3, figsize=(10.5, 3.1))25    axes[0].errorbar(hist["condition"].str.extract(r"H=(\d+)")[0].astype(int), hist["close_rmse_mean"], yerr=hist["close_rmse_std"], marker="o", label="close wake")26    axes[0].errorbar(hist["condition"].str.extract(r"H=(\d+)")[0].astype(int), hist["global_rmse_mean"], yerr=hist["global_rmse_std"], marker="s", label="global")27    axes[0].axhline(0.1, color="0.6", ls="--", lw=1)28    axes[0].set_xlabel("History length H")29    axes[0].set_ylabel("Position RMSE (L)")30    axes[0].legend(frameon=False, fontsize=8)31    axes[1].bar(sensor["condition"], sensor["close_rmse_mean"], yerr=sensor["close_rmse_std"], color="#6C8EBF")32    axes[1].axhline(0.1, color="0.6", ls="--", lw=1)33    axes[1].set_ylabel("Close-wake RMSE (L)")34    axes[1].tick_params(axis="x", rotation=25)35    for comp, sub in noise.groupby("component"):36        axes[2].errorbar(sub["noise"], sub["close_rmse_mean"], yerr=sub["close_rmse_std"], marker="o", label=comp)37    axes[2].axhline(0.1, color="0.6", ls="--", lw=1)38    axes[2].set_xlabel("Raw-noise multiplier")39    axes[2].set_ylabel("Close-wake RMSE (L)")40    axes[2].legend(frameon=False, fontsize=7)41    fig.tight_layout()42    save_pdf_svg(fig, out, "fig02_main_results")43 44 45def plot_self_signal(results_dir: Path, out: Path) -> None:46    rand = pd.read_csv(results_dir / "self_signal_randomized.csv")47    fig, ax = plt.subplots(figsize=(4.5, 3.2))48    x = range(len(rand))49    ax.bar([i - 0.18 for i in x], rand["close_rmse_mean"], width=0.36, label="position RMSE")50    ax.bar([i + 0.18 for i in x], rand["close_theta_mae_mean"] / 10.0, width=0.36, label="theta MAE / 10")51    ax.axhline(0.1, color="0.6", ls="--", lw=1)52    ax.set_xticks(list(x), rand["condition"], rotation=20)53    ax.set_ylabel("Close-wake metric")54    ax.legend(frameon=False, fontsize=8)55    fig.tight_layout()56    save_pdf_svg(fig, out, "fig03_self_signal_control")57 58 59def plot_supp(results_dir: Path, out: Path) -> None:60    model = pd.read_csv(results_dir / "model_screen.csv")61    sample = pd.read_csv(results_dir / "sample_size_convergence.csv")62    fig, ax = plt.subplots(figsize=(5.5, 3.2))63    available = model[model["status"] == "available"]64    ax.bar(available["model"], available["global_rmse_mean"], yerr=available["global_rmse_std"], color="#8DA0CB")65    ax.set_ylabel("Global position RMSE (L)")66    ax.tick_params(axis="x", rotation=30)67    fig.tight_layout()68    save_pdf_svg(fig, out, "figS4_architecture_screen")69    fig, ax = plt.subplots(figsize=(4.5, 3.2))70    ax.errorbar(sample["training_poses"], sample["global_rmse_mean"], yerr=sample["global_rmse_std"], marker="o", label="global")71    ax.errorbar(sample["training_poses"], sample["close_rmse_mean"], yerr=sample["close_rmse_std"], marker="s", label="close wake")72    ax.axhline(0.1, color="0.6", ls="--", lw=1)73    ax.set_xlabel("Training poses")74    ax.set_ylabel("Position RMSE (L)")75    ax.legend(frameon=False, fontsize=8)76    fig.tight_layout()77    save_pdf_svg(fig, out, "figS5_sample_size_convergence")78 79