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

13
gridlock/read.py
View File

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