[read] add option to visualize on preexisting axes

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)
 
+        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')