| """Shared plotting utilities.""" | |
| from __future__ import annotations | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| def setup_axes( | |
| ax: plt.Axes | None = None, | |
| figsize: tuple[float, float] = (6, 4), | |
| ) -> tuple[plt.Figure, plt.Axes]: | |
| """Return a (fig, ax) pair, creating them if *ax* is None.""" | |
| if ax is None: | |
| fig, ax = plt.subplots(figsize=figsize) | |
| else: | |
| fig = ax.figure | |
| return fig, ax | |
| def save_or_show(fig: plt.Figure, save: str | None, show: bool) -> None: | |
| """Save figure to file and/or display.""" | |
| if save is not None: | |
| fig.savefig(save, bbox_inches="tight", dpi=150) | |
| if show: | |
| plt.show() | |
| else: | |
| plt.close(fig) | |