File size: 1,387 Bytes
0cf8154 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | """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 # noqa: E402
from zenith.engines.base import Result # noqa: E402
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: # map grid: rows run north to south
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)
|