Spaces:
Running on Zero
Running on Zero
| """ | |
| Plotting helpers for OT energy evaluation and FM diagnostics. | |
| All functions are side-effect free besides writing image files. They raise on | |
| errors, so callers can handle exceptions (keeping compatibility with existing | |
| try/except blocks in scripts). | |
| """ | |
| from __future__ import annotations | |
| from typing import Iterable, Sequence | |
| def _to_numpy1d(x: Iterable) -> "object": | |
| """Convert a 1D-like iterable (list/torch/tensor/numpy) to a numpy array.""" | |
| try: | |
| import numpy as np | |
| except Exception as e: # pragma: no cover | |
| raise RuntimeError("numpy is required for plotting helpers") from e | |
| # Torch tensor support without importing torch at module level | |
| if hasattr(x, "detach") and hasattr(x, "cpu") and hasattr(x, "numpy"): | |
| try: | |
| return x.detach().cpu().numpy() | |
| except Exception: | |
| pass | |
| # Numpy array or array-like | |
| try: | |
| return np.asarray(list(x), dtype=float) | |
| except Exception: | |
| # Best-effort fallback | |
| return np.array(list(x), dtype=float) | |
| def plot_fm_loss( | |
| iterations: Sequence[int], | |
| loss_mean: Sequence[float], | |
| loss_std: Sequence[float], | |
| *, | |
| out_path: str, | |
| show_std: bool = False, | |
| sigma_k: float = 3.0, | |
| dpi: int = 150, | |
| title: str = "FM training loss vs. iteration", | |
| ): | |
| import matplotlib.pyplot as plt | |
| its = list(iterations) | |
| lmean = _to_numpy1d(loss_mean) | |
| lstd = _to_numpy1d(loss_std) | |
| fig, ax = plt.subplots(figsize=(7.5, 4.5)) | |
| ax.set_title(title) | |
| ax.set_xlabel("iteration") | |
| ax.set_ylabel("loss (mean per selected pair)") | |
| ax.plot(its, lmean, color="C0", label="loss") | |
| if show_std: | |
| ax.fill_between(its, lmean - sigma_k * lstd, lmean + sigma_k * lstd, color="C0", alpha=0.2, label=f"±{sigma_k}σ") | |
| ax.legend(loc="best") | |
| fig.tight_layout() | |
| fig.savefig(out_path, dpi=int(dpi)) | |
| plt.close(fig) | |
| def plot_fm_ediff( | |
| iterations: Sequence[int], | |
| ediff_mean: Sequence[float], | |
| ediff_std: Sequence[float], | |
| *, | |
| out_path: str, | |
| sigma_k: float = 3.0, | |
| dpi: int = 150, | |
| title: str = "E_data - E_noise vs. iteration", | |
| ): | |
| import matplotlib.pyplot as plt | |
| its = list(iterations) | |
| edmean = _to_numpy1d(ediff_mean) | |
| edstd = _to_numpy1d(ediff_std) | |
| fig, ax = plt.subplots(figsize=(7.5, 4.5)) | |
| ax.set_title(title) | |
| ax.set_xlabel("iteration") | |
| ax.set_ylabel("energy difference (mean per selected pair)") | |
| ax.plot(its, edmean, color="C1", label="ΔE mean") | |
| ax.fill_between(its, edmean - sigma_k * edstd, edmean + sigma_k * edstd, color="C1", alpha=0.2, label=f"±{sigma_k}σ") | |
| ax.axhline(0.0, color="k", linewidth=0.8, linestyle=":") | |
| ax.legend(loc="best") | |
| fig.tight_layout() | |
| fig.savefig(out_path, dpi=int(dpi)) | |
| plt.close(fig) | |
| def plot_interpolation_energy( | |
| taus: Iterable[float], | |
| energy_means: Sequence[float], | |
| energy_stds: Sequence[float], | |
| *, | |
| mu_noise: float, | |
| sd_noise: float, | |
| mu_data: float, | |
| sd_data: float, | |
| out_path: str, | |
| sigma_k: float = 3.0, | |
| dpi: int = 150, | |
| title: str = "GEM energy vs. interpolation time (t=0 in model)", | |
| ): | |
| import matplotlib.pyplot as plt | |
| t_np = _to_numpy1d(taus) | |
| m_np = _to_numpy1d(energy_means) | |
| s_np = _to_numpy1d(energy_stds) | |
| fig, ax = plt.subplots(figsize=(7.5, 4.5)) | |
| ax.set_title(title) | |
| ax.set_xlabel("interpolation time (0=noise, 1=data)") | |
| ax.set_ylabel("energy") | |
| ax.plot(t_np, m_np, label="interp mean", color="C0") | |
| ax.fill_between(t_np, m_np - sigma_k * s_np, m_np + sigma_k * s_np, color="C0", alpha=0.2, label=f"interp ±{sigma_k}σ") | |
| # Noise/data reference bands | |
| ax.axhline(mu_noise, color="C1", linestyle="--", label="noise mean") | |
| ax.fill_between([0, 1], mu_noise - sigma_k * sd_noise, mu_noise + sigma_k * sd_noise, color="C1", alpha=0.15, label=f"noise ±{sigma_k}σ") | |
| ax.axhline(mu_data, color="C2", linestyle="--", label="data mean") | |
| ax.fill_between([0, 1], mu_data - sigma_k * sd_data, mu_data + sigma_k * sd_data, color="C2", alpha=0.15, label=f"data ±{sigma_k}σ") | |
| ax.legend(loc="best") | |
| fig.tight_layout() | |
| fig.savefig(out_path, dpi=int(dpi)) | |
| plt.close(fig) | |
| def plot_interpolation_grad( | |
| taus: Iterable[float], | |
| grad_means: Sequence[float], | |
| grad_stds: Sequence[float], | |
| *, | |
| out_path: str, | |
| sigma_k: float = 3.0, | |
| dpi: int = 150, | |
| title: str = "||-∇_x E|| vs. interpolation time (t=0 in model)", | |
| ): | |
| import matplotlib.pyplot as plt | |
| t_np = _to_numpy1d(taus) | |
| gm_np = _to_numpy1d(grad_means) | |
| gs_np = _to_numpy1d(grad_stds) | |
| fig, ax = plt.subplots(figsize=(7.5, 4.5)) | |
| ax.set_title(title) | |
| ax.set_xlabel("interpolation time (0=noise, 1=data)") | |
| ax.set_ylabel("gradient strength (mean per node/edge)") | |
| ax.plot(t_np, gm_np, label="||-∇E|| mean", color="C3") | |
| ax.fill_between(t_np, gm_np - sigma_k * gs_np, gm_np + sigma_k * gs_np, color="C3", alpha=0.2, label=f"±{sigma_k}σ") | |
| ax.legend(loc="best") | |
| fig.tight_layout() | |
| fig.savefig(out_path, dpi=int(dpi)) | |
| plt.close(fig) | |