import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np
# ------------------------------------------------------------------
[docs]
def plot_grid_contour(
Z: np.ndarray,
X: np.ndarray | None = None,
Y: np.ndarray | None = None,
*,
x_label: str,
y_label: str,
var_name: str = "",
title: str = "",
filename: str | None = None,
levels: int | np.ndarray | None = 60,
lower: float | None = None,
upper: float | None = None,
ax=None,
render_as_image: bool = False,
cmap: str = "seismic",
):
"""
Draw a filled contour or image for a two-dimensional mesh array.
``Z`` holds the plotted values; ``X`` and ``Y`` are plotting coordinates.
The caller owns the returned figure. A file is written only when
``filename`` is provided. ``ax`` embeds the plot in a caller-owned figure;
``render_as_image`` uses pixel rendering instead of contours. Color and
axis limits, labels, levels, and colormap are passed explicitly.
Returns
-------
tuple[matplotlib.figure.Figure, matplotlib.axes.Axes]
Figure and axes containing the plotted values and color bar.
"""
# -------- 1. colour limits & levels
finite_values = np.ma.masked_invalid(Z).compressed()
if finite_values.size == 0:
raise ValueError("The selected R-Z interpolation grid contains no finite data")
vmin = float(finite_values.min()) if lower is None else lower
vmax = float(finite_values.max()) if upper is None else upper
if vmin > vmax:
raise ValueError("Color limits must satisfy minimum <= maximum")
if vmin == vmax:
pad = max(abs(vmin)*1e-8, 1e-12)
vmin -= pad
vmax += pad
if isinstance(levels, int):
levels = np.linspace(vmin, vmax, levels + 1)
# diverging scale centred at 0 when data span both signs
if render_as_image and not (vmin < 0.0 < vmax):
norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
elif (vmin < 0.0 and vmax > 0.0):
norm = mcolors.TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax)
elif (vmin < 0.0 and vmax < 0.0):
norm = mcolors.TwoSlopeNorm(vmin=vmin, vcenter=vmax,
vmax=vmax*(1 + 1e-8))
else: # entirely positive
norm = mcolors.TwoSlopeNorm(vmin=vmin*(1-1e-8), vcenter=vmin,
vmax=vmax)
# -------- 2. plotting
supplied_ax = ax is not None
if ax is None:
fig, ax = plt.subplots(figsize=(6, 5))
else:
fig = ax.figure
if render_as_image:
if X is None or Y is None:
raise ValueError("R-Z coordinates are required for image rendering")
cf = ax.imshow(Z, origin="lower", extent=(X.min(), X.max(), Y.min(), Y.max()),
interpolation="nearest", cmap=cmap, norm=norm)
else:
cf = ax.contourf(X, Y, Z,
levels=levels,
cmap=cmap,
norm=norm,
extend="both",
vmin=vmin, vmax=vmax)
ax.set_xlabel(x_label)
ax.set_ylabel(y_label)
ax.set_aspect("equal" if x_label.startswith("R") else "auto")
if title:
ax.set_title(title)
cbar = fig.colorbar(cf, ax=ax)
if var_name:
cbar.set_label(var_name, rotation=270, labelpad=15)
if not supplied_ax:
fig.tight_layout()
if filename:
fig.savefig(filename, bbox_inches="tight")
return fig, ax
# ------------------------------------------------------------------