jan/gridlock
2
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

[read] add option to visualize on preexisting axes

Browse Source
This commit is contained in:
Jan Petykiewicz 2025-09-22 19:18:49 -07:00
parent 6802e57fa9
commit 21304f0dbf
1 changed files with 10 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
gridlock/read.py
View File

@ -87,6 +87,7 @@ class GridReadMixin(GridPosMixin):
sample_period: int = 1,
finalize: bool = True,
pcolormesh_args: dict[str, Any] | None = None,
ax: 'matplotlib.axes.Axes' | None = None,
) -> tuple['matplotlib.figure.Figure', 'matplotlib.axes.Axes']:
"""
Visualize a slice of a grid.
@ -99,6 +100,7 @@ class GridReadMixin(GridPosMixin):
sample_period: Period for down-sampling the image. Default 1 (disabled)
finalize: Whether to call `pyplot.show()` after constructing the plot. Default `True`
pcolormesh_args: Args passed through to matplotlib `pcolormesh()`
ax: If provided, plot to these axes (instead of creating a new figure & axes)
Returns:
(Figure, Axes)
@ -124,7 +126,10 @@ class GridReadMixin(GridPosMixin):
xmesh, ymesh = numpy.meshgrid(x, y, indexing='ij')
x_label, y_label = ('xyz'[a] for a in surface)
if ax is None:
fig, ax = pyplot.subplots()
else:
fig = ax.figure
mappable = ax.pcolormesh(xmesh, ymesh, grid_slice, **pcolormesh_args)
fig.colorbar(mappable)
ax.set_aspect('equal', adjustable='box')