From 21304f0dbfce8b089a44a7dfb13c072b57839bfa Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Mon, 22 Sep 2025 19:18:49 -0700 Subject: [PATCH] [read] add option to visualize on preexisting axes --- gridlock/read.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/gridlock/read.py b/gridlock/read.py index cfc8f3d..b5840e7 100644 --- a/gridlock/read.py +++ b/gridlock/read.py @@ -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) @@ -112,10 +114,10 @@ class GridReadMixin(GridPosMixin): pcolormesh_args = {} grid_slice = self.get_slice( - cell_data=cell_data, - plane=plane, - which_shifts=which_shifts, - sample_period=sample_period, + cell_data = cell_data, + plane = plane, + which_shifts = which_shifts, + sample_period = sample_period, ) surface = numpy.delete(range(3), plane.axis) @@ -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) - fig, ax = pyplot.subplots() + 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')