| """Render a result's plot series to a PNG.""" |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import matplotlib |
| import numpy as np |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
|
|
| from zenith.engines.base import Result |
|
|
|
|
| def save_plot(result: Result, path: Path) -> None: |
| items = list(result.series.items()) |
| if not items: |
| raise ValueError(f"{result.sim}/{result.op} has no plot series") |
| fig, axes = plt.subplots(1, len(items), figsize=(5 * len(items), 4), squeeze=False) |
| for ax, (name, s) in zip(axes[0], items): |
| if "z" in s: |
| z = np.array([[np.nan if v is None else v for v in row] for row in s["z"]], dtype=float) |
| im = ax.imshow(z, cmap="magma", extent=(s["x"][0], s["x"][-1], s["y"][-1], s["y"][0]), aspect="auto") |
| fig.colorbar(im, ax=ax, label=s["z_label"]) |
| else: |
| ax.plot(s["x"], s["y"], marker="o" if len(s["x"]) == 1 else None) |
| ax.set_title(name.replace("_", " ")) |
| ax.set_xlabel(s["x_label"]) |
| ax.set_ylabel(s["y_label"]) |
| if name.startswith("orbit_xy"): |
| ax.set_aspect("equal") |
| ax.plot([0], [0], "*", color="orange", markersize=12) |
| fig.suptitle(f"{result.sim} / {result.op}") |
| fig.tight_layout() |
| fig.savefig(path, dpi=100) |
| plt.close(fig) |
|
|
