File size: 701 Bytes
4e2940e | 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 | """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)
|