import os, sys, matplotlib
if ( not os.environ.get("DISPLAY") ) and ( "ipykernel" not in sys.modules ):
matplotlib.use("Agg")
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,
):
"""
General filled-contour plot for any 2-D NumPy array *Z* defined on the
rectangular mesh (X, Y).
All keyword arguments after the literal ``*`` are the same as in the wrapper
routines; only *x_label* and *y_label* are required so the caller can
supply axis text appropriate to its own context.
"""
# -------- 1. colour limits & levels
vmin = Z.min() if lower is None else lower
vmax = Z.max() if upper is None else upper
if isinstance(levels, int):
levels = np.linspace(vmin, vmax, levels + 1)
# diverging scale centred at 0 when data span both signs
if (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
fig, ax = plt.subplots(figsize=(6, 5))
cf = ax.contourf(X, Y, Z,
levels=levels,
cmap="seismic",
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)
plt.tight_layout()
if filename:
fig.savefig(filename, bbox_inches="tight")
elif ( not os.environ.get("DISPLAY") ) and ( "ipykernel" not in sys.modules ):
fig.savefig("img.svg", bbox_inches="tight")
else:
plt.show()
return fig, ax
# ------------------------------------------------------------------