From 6802e57fa9fb26d5017b55b9bd46725db38ae04b Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Mon, 22 Sep 2025 19:18:08 -0700 Subject: [PATCH 1/8] [read] add missing arg to docstring --- gridlock/read.py | 1 + 1 file changed, 1 insertion(+) diff --git a/gridlock/read.py b/gridlock/read.py index 707251a..cfc8f3d 100644 --- a/gridlock/read.py +++ b/gridlock/read.py @@ -98,6 +98,7 @@ class GridReadMixin(GridPosMixin): which_shifts: Which grid to display. Default is the first grid (0). 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()` Returns: (Figure, Axes) From 21304f0dbfce8b089a44a7dfb13c072b57839bfa Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Mon, 22 Sep 2025 19:18:49 -0700 Subject: [PATCH 2/8] [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') From 68520b871018c3d86254fa1fa87faf4838351d11 Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Mon, 22 Sep 2025 19:19:28 -0700 Subject: [PATCH 3/8] [read] add visualize_edges() --- gridlock/read.py | 74 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) diff --git a/gridlock/read.py b/gridlock/read.py index b5840e7..600227d 100644 --- a/gridlock/read.py +++ b/gridlock/read.py @@ -135,12 +135,86 @@ class GridReadMixin(GridPosMixin): ax.set_aspect('equal', adjustable='box') ax.set_xlabel(x_label) ax.set_ylabel(y_label) + if finalize: pyplot.show() return fig, ax + def visualize_edges( + self, + cell_data: NDArray, + plane: PlaneProtocol | PlaneDict, + which_shifts: int = 0, + finalize: bool = True, + contour_args: dict[str, Any] | None = None, + ax: 'matplotlib.axes.Axes' | None = None, + level_fraction: float = 0.7, + ) -> tuple['matplotlib.figure.Figure', 'matplotlib.axes.Axes']: + """ + Visualize the edges of a grid slice. + This is intended as an overlay on top of visualize_slice (e.g. showing epsilon boundaries + on an E-field plot). + + Interpolates if given a position between two grid planes. + + Args: + cell_data: Cell data to visualize + plane: Axis and position (`Plane`) of the plane to read. + which_shifts: Which grid to display. Default is the first grid (0). + 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) + level_fraction: Value between 0 and 1 which tunes how many contours are generated. + 1 indicates that every possible step should have its own contour. + + Returns: + (Figure, Axes) + """ + from matplotlib import pyplot + + if level_fraction > 1: + raise GridError(f'{level_fraction=} must be between 0 and 1') + + if isinstance(plane, dict): + plane = Plane(**plane) + + if pcolormesh_args is None: + pcolormesh_args = {} + + grid_slice = self.get_slice( + cell_data = cell_data, + plane = plane, + which_shifts = which_shifts, + sample_period = sample_period, + ) + cvals, cval_counts = numpy.unique(grid_slice, return_counts=True) + if cvals.size == 1: + levels = [cvals[0] + 1] + else: + cval_order = numpy.argsort(cval_counts)[::-1] + level_count = 2 + while cval_counts[cval_order[:level_count]].sum() < level_fraction: + level_count += 1 + ctr_levels = cvals[cval_order[:level_count]] + levels = numpy.diff(ctr_levels[::-1]) + ctr_levels[:0:-1] + + surface = numpy.delete(range(3), plane.axis) + + if ax is None: + fig, ax = pyplot.subplots() + else: + fig = ax.figure + xc, yc = (self.shifted_exyz(which_shifts)[a] for a in surface) + xcmesh, ycmesh = numpy.meshgrid(xc, yc, indexing='ij') + + mappable = ax.contour(xcmesh, ycmesh, grid_slice, levels=levels, **contour_args) + + return fig, ax + + def visualize_isosurface( self, cell_data: NDArray, From 7cac73bcb400021289350a52215083e88c337cd7 Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Mon, 22 Sep 2025 19:25:39 -0700 Subject: [PATCH 4/8] [draw] add missing code for finalize --- gridlock/read.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/gridlock/read.py b/gridlock/read.py index 600227d..4f39432 100644 --- a/gridlock/read.py +++ b/gridlock/read.py @@ -212,6 +212,9 @@ class GridReadMixin(GridPosMixin): mappable = ax.contour(xcmesh, ycmesh, grid_slice, levels=levels, **contour_args) + if finalize: + pyplot.show() + return fig, ax From 16a76e0122845d029c7a4a3c2896ab0606553576 Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Mon, 22 Sep 2025 19:26:59 -0700 Subject: [PATCH 5/8] [read] make visualize_edges more friendly for overlay by default --- gridlock/read.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gridlock/read.py b/gridlock/read.py index 4f39432..28afdd2 100644 --- a/gridlock/read.py +++ b/gridlock/read.py @@ -182,7 +182,7 @@ class GridReadMixin(GridPosMixin): plane = Plane(**plane) if pcolormesh_args is None: - pcolormesh_args = {} + pcolormesh_args = dict(alpha=0.8, colors='gray') grid_slice = self.get_slice( cell_data = cell_data, From 64752873fbde2392b91150f8b26740dfdc633e66 Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Mon, 22 Sep 2025 19:28:40 -0700 Subject: [PATCH 6/8] [read] fix type spec --- gridlock/read.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gridlock/read.py b/gridlock/read.py index 28afdd2..44f1c8a 100644 --- a/gridlock/read.py +++ b/gridlock/read.py @@ -87,7 +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, + ax: 'matplotlib.axes.Axes | None' = None, ) -> tuple['matplotlib.figure.Figure', 'matplotlib.axes.Axes']: """ Visualize a slice of a grid. @@ -149,7 +149,7 @@ class GridReadMixin(GridPosMixin): which_shifts: int = 0, finalize: bool = True, contour_args: dict[str, Any] | None = None, - ax: 'matplotlib.axes.Axes' | None = None, + ax: 'matplotlib.axes.Axes | None' = None, level_fraction: float = 0.7, ) -> tuple['matplotlib.figure.Figure', 'matplotlib.axes.Axes']: """ From f4818fd55450463a63cfa8bcc26ebbcecfce7a88 Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Mon, 22 Sep 2025 20:32:01 -0700 Subject: [PATCH 7/8] [draw] fix arg naming --- gridlock/read.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/gridlock/read.py b/gridlock/read.py index 44f1c8a..7b4de1e 100644 --- a/gridlock/read.py +++ b/gridlock/read.py @@ -165,7 +165,7 @@ class GridReadMixin(GridPosMixin): which_shifts: Which grid to display. Default is the first grid (0). 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()` + contour_args: Args passed through to matplotlib `pcolormesh()` ax: If provided, plot to these axes (instead of creating a new figure & axes) level_fraction: Value between 0 and 1 which tunes how many contours are generated. 1 indicates that every possible step should have its own contour. @@ -181,8 +181,8 @@ class GridReadMixin(GridPosMixin): if isinstance(plane, dict): plane = Plane(**plane) - if pcolormesh_args is None: - pcolormesh_args = dict(alpha=0.8, colors='gray') + if contour_args is None: + contour_args = dict(alpha=0.8, colors='gray') grid_slice = self.get_slice( cell_data = cell_data, From 32b6c207dcf703f6c695dc9fae3f7615bd5e8f15 Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Mon, 22 Sep 2025 22:24:15 -0700 Subject: [PATCH 8/8] [read] more fixup for visualize_edges --- gridlock/read.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/gridlock/read.py b/gridlock/read.py index 7b4de1e..503e996 100644 --- a/gridlock/read.py +++ b/gridlock/read.py @@ -147,6 +147,7 @@ class GridReadMixin(GridPosMixin): cell_data: NDArray, plane: PlaneProtocol | PlaneDict, which_shifts: int = 0, + sample_period: int = 1, finalize: bool = True, contour_args: dict[str, Any] | None = None, ax: 'matplotlib.axes.Axes | None' = None, @@ -207,7 +208,7 @@ class GridReadMixin(GridPosMixin): fig, ax = pyplot.subplots() else: fig = ax.figure - xc, yc = (self.shifted_exyz(which_shifts)[a] for a in surface) + xc, yc = (self.shifted_xyz(which_shifts)[a] for a in surface) xcmesh, ycmesh = numpy.meshgrid(xc, yc, indexing='ij') mappable = ax.contour(xcmesh, ycmesh, grid_slice, levels=levels, **contour_args)