From ce898c8f16c6383fe707a16a5c25506d03334f88 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 28 Nov 2022 16:15:03 +0000 Subject: [PATCH 1/6] Script/executable for plotting an equilibrium --- doc/whats-new.md | 2 + hypnotoad/core/equilibrium.py | 75 ++++++++++-- .../scripts/hypnotoad_plot_equilibrium.py | 108 ++++++++++++++++++ setup.py | 2 + 4 files changed, 179 insertions(+), 8 deletions(-) create mode 100755 hypnotoad/scripts/hypnotoad_plot_equilibrium.py diff --git a/doc/whats-new.md b/doc/whats-new.md index 1981d8cb..de81ed4c 100644 --- a/doc/whats-new.md +++ b/doc/whats-new.md @@ -39,6 +39,8 @@ What's new ### New features +- Convenience script for making a plot of an equilibrium from a geqdsk file (#139)\ + By [John Omotani](https://github.com/johnomotani) - If an `int` (or other `Number`) literal is passed to a float-type option, it is converted implicitly instead of causing an error. (#134)\ By [John Omotani](https://github.com/johnomotani) diff --git a/hypnotoad/core/equilibrium.py b/hypnotoad/core/equilibrium.py index bff5924f..1a0cd925 100644 --- a/hypnotoad/core/equilibrium.py +++ b/hypnotoad/core/equilibrium.py @@ -4818,7 +4818,7 @@ def plotPotential( return axis - def plotWall(self, axis=None): + def plotWall(self, axis=None, *, color="k", linestyle="-", linewidth=2, **kwargs): if self.wall: wall_R = [p.R for p in self.wall] wall_Z = [p.Z for p in self.wall] @@ -4830,16 +4830,75 @@ def plotWall(self, axis=None): if axis is None: from matplotlib import pyplot - axis = pyplot.plot(wall_R, wall_Z, "k-", linewidth=2) + axis = pyplot.plot( + wall_R, + wall_Z, + color=color, + linestyle=linestyle, + linewidth=linewidth, + **kwargs, + ) else: - axis.plot(wall_R, wall_Z, "k-", linewidth=2) + axis.plot( + wall_R, + wall_Z, + color=color, + linestyle=linestyle, + linewidth=linewidth, + **kwargs, + ) return axis - def plotSeparatrix(self): + def plotSeparatrix( + self, + *, + scatter=True, + separate_contours=False, + npoints=100, + marker="x", + **kwargs, + ): + """ + Plot the separatrix contour(s) + + Parameters + ---------- + scatter : bool, default True + If `True`, make a scatter plot of the points on the EquilibriumRegion + contours in `self.regions`. If `False`, make a line plot. Only used when + `separate_contours=False`. + separate_contours : bool, default False + If `False`, plot the EquilibriumRegion contours from `self.regions`. If + `True`, make a contour plot of the psi values in `self.psi_sep` - this is + useful for disconnected double-null equilibria to plot the true + separatrices. + npoints : int, default 100 + When `separate_contours=True`, number of points used in the grid + discretizing psi. + marker : default "x" + Argument passed to `marker` argument of `pyplot.scatter()`. + **kwargs + Extra keyword arguments passed to `pyplot.scatter()` or `pyplot.plot()`. + """ from matplotlib import pyplot - for region in self.regions.values(): - R = [p.R for p in region] - Z = [p.Z for p in region] - pyplot.scatter(R, Z, marker="x", label=region.name) + if separate_contours: + R = numpy.linspace(self.Rmin, self.Rmax, npoints) + Z = numpy.linspace(self.Zmin, self.Zmax, npoints) + + pyplot.contour( + R, + Z, + self.psi(R[:, numpy.newaxis], Z[numpy.newaxis, :]).T, + self.psi_sep, + **kwargs, + ) + else: + for region in self.regions.values(): + R = [p.R for p in region] + Z = [p.Z for p in region] + if scatter: + pyplot.scatter(R, Z, marker=marker, label=region.name, **kwargs) + else: + pyplot.plot(R, Z, label=region.name, **kwargs) diff --git a/hypnotoad/scripts/hypnotoad_plot_equilibrium.py b/hypnotoad/scripts/hypnotoad_plot_equilibrium.py new file mode 100755 index 00000000..eae30b26 --- /dev/null +++ b/hypnotoad/scripts/hypnotoad_plot_equilibrium.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 + + +def main(): + from argparse import ArgumentParser + + parser = ArgumentParser("Plot the equilibrium stored in a geqdsk file") + parser.add_argument( + "equilibrium_file", help="Path to equilibrium file in geqdsk format" + ) + parser.add_argument( + "--wall", + action="store_false", + default=True, + help="Plot the wall contour if present", + ) + parser.add_argument( + "--wall-width", + default=2.0, + type=float, + help="Width of the wall contour and separatrix contour", + ) + parser.add_argument( + "--separatrix", + action="store_false", + default=True, + help="Plot the separatrix. Note that for a disconnected double-null " + "equilibrium, the thing plotted by default is not really the separatrix, as " + "around the core the contour is a weighted average of the two separatrices " + "constructed so that it joins the two X-points, while in the divertor legs it " + "is the separatrix connected to the nearest X-point.", + ) + parser.add_argument( + "--separate-separatrices", + action="store_true", + default=False, + help="Plot the two separatrices separately for a disconnected double-null " + "equilibrium.", + ) + parser.add_argument( + "--separatrix-color", default="blue", help="Color of the separatrix contour" + ) + parser.add_argument( + "--psi-labels", action="store_true", default=False, help="Label psi contours" + ) + parser.add_argument( + "--n-contours", default=40, type=int, help="Number of psi contours to plot" + ) + parser.add_argument( + "--color-contours", + action="store_true", + default=False, + help="Color psi contours", + ) + parser.add_argument( + "--show", + action="store_false", + default=True, + help="Show plot in interactive window", + ) + parser.add_argument( + "--output", + default=None, + help="Name for output file. Suffix determines file format", + ) + args = parser.parse_args() + + from ..cases import tokamak + from matplotlib import pyplot as plt + + with open(args.equilibrium_file, "rt") as fh: + eq = tokamak.read_geqdsk(fh) + + # Work out sensible aspect ratio for figure + figwidth = 4.0 + figheight = figwidth * (eq.Zmax - eq.Zmin) / (eq.Rmax - eq.Rmin) + plt.figure(figsize=(figwidth, figheight)) + + if args.color_contours: + colors = None + else: + colors = "grey" + eq.plotPotential( + ncontours=args.n_contours, labels=args.psi_labels, colors=colors, linestyles="-" + ) + + if args.wall: + eq.plotWall(linewidth=args.wall_width) + + if args.separatrix: + eq.plotSeparatrix( + scatter=False, + separate_contours=args.separate_separatrices, + linewidth=args.wall_width, + linewidths=args.wall_width, + color=args.separatrix_color, + colors=args.separatrix_color, + ) + + if args.output is not None: + plt.savefig(args.output, bbox_inches="tight") + + if args.show: + plt.show() + + +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py index 85a11051..6797afdf 100644 --- a/setup.py +++ b/setup.py @@ -41,6 +41,8 @@ "hypnotoad_circular = hypnotoad.scripts.hypnotoad_circular:main", "hypnotoad_geqdsk = hypnotoad.scripts.hypnotoad_geqdsk:main", "hypnotoad_torpex = hypnotoad.scripts.hypnotoad_torpex:main", + "hypnotoad-plot-equilibrium = " + "hypnotoad.scripts.hypnotoad_plot_equilibrium:main", "hypnotoad-plot-grid-cells = " "hypnotoad.scripts.hypnotoad_plot_grid_cells:main", "hypnotoad-recreate-inputs = " From aed148d18a3e27eb5761410dc651c19de8d61fd6 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 28 Nov 2022 17:24:38 +0000 Subject: [PATCH 2/6] Add list of executables provided by the hypnotoad package in README.md --- README.md | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/README.md b/README.md index b557f996..4a067c11 100644 --- a/README.md +++ b/README.md @@ -111,6 +111,28 @@ accuracy. The following may be a good starting point: points, so the structure should be very similar. +Utilities +--------- + +Hypnotoad provides several executables for working with equilibria and grid files: +- `hypnotoad-gui` is the main GUI interface +- `hypnotoad_geqdsk` is a command line interface for creating tokamak + equilibria from geqdsk equilibrium files +- `hypnotoad_circular` is a command line interface for creating grid files for + concentric, circular flux surfaces +- `hypnotoad_torpex` is a command line interface for creating grid files for + TORPEX X-point configurations +- `hypnotoad-plot-equilibrium` is a command line tool for creating plots of the + equilibrium (flux surfaces, wall and separatrix) from a geqdsk file +- `hypnotoad-plot-grid-cells` creates a plot of the grid cells from a grid file + generated by hypnotoad +- `hypnotoad-recreate-inputs` writes extracts from a grid file copies of the + input YAML file and geqdsk file that were used to create the grid file + originally + +For more information, pass the `--help` flag to any of these commands. + + Developing ---------- From 5b0e58fcced5801f310c1a3b100b3056cc9b733c Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 28 Nov 2022 18:07:46 +0000 Subject: [PATCH 3/6] Function to plot a highlighted region between two psiN values May be useful to show the area covered by a grid as an alternative to plotting the grid points. --- hypnotoad/cases/tokamak.py | 36 +++++++++--------- hypnotoad/core/equilibrium.py | 38 +++++++++++++++++++ .../scripts/hypnotoad_plot_equilibrium.py | 16 ++++++++ 3 files changed, 73 insertions(+), 17 deletions(-) diff --git a/hypnotoad/cases/tokamak.py b/hypnotoad/cases/tokamak.py index e1355bee..14fabbe0 100644 --- a/hypnotoad/cases/tokamak.py +++ b/hypnotoad/cases/tokamak.py @@ -642,6 +642,17 @@ def dpos_dl(distance, pos): leg_lines = leg_lines[::-1] return {"inner": leg_lines[0], "outer": leg_lines[1]} + # psi values + def _psinorm_to_psi(self, psinorm): + if psinorm is None: + return None + return self.psi_axis + psinorm * (self.psi_sep[0] - self.psi_axis) + + def _psi_to_psinorm(self, psi): + if psi is None: + return None + return (psi - self.psi_axis) / (self.psi_sep[0] - self.psi_axis) + def makeRegions(self): """Main region generation function. Regions are logically rectangular ranges in poloidal angle; segments are @@ -680,34 +691,25 @@ def makeRegions(self): if self.psi_axis is None: raise ValueError("psi_axis has not been set") - # psi values - def psinorm_to_psi(psinorm): - if psinorm is None: - return None - return self.psi_axis + psinorm * (self.psi_sep[0] - self.psi_axis) - - def psi_to_psinorm(psi): - if psi is None: - return None - return (psi - self.psi_axis) / (self.psi_sep[0] - self.psi_axis) - self.psi_core = with_default( - self.user_options.psi_core, psinorm_to_psi(self.user_options.psinorm_core) + self.user_options.psi_core, + self._psinorm_to_psi(self.user_options.psinorm_core), ) self.psi_sol = with_default( - self.user_options.psi_sol, psinorm_to_psi(self.user_options.psinorm_sol) + self.user_options.psi_sol, + self._psinorm_to_psi(self.user_options.psinorm_sol), ) self.psi_sol_inner = with_default( self.user_options.psi_sol_inner, - psinorm_to_psi(self.user_options.psinorm_sol_inner), + self._psinorm_to_psi(self.user_options.psinorm_sol_inner), ) self.psi_pf_lower = with_default( self.user_options.psi_pf_lower, - psinorm_to_psi(self.user_options.psinorm_pf_lower), + self._psinorm_to_psi(self.user_options.psinorm_pf_lower), ) self.psi_pf_upper = with_default( self.user_options.psi_pf_upper, - psinorm_to_psi(self.user_options.psinorm_pf_upper), + self._psinorm_to_psi(self.user_options.psinorm_pf_upper), ) self.poloidal_spacing_delta_psi = with_default( @@ -721,7 +723,7 @@ def psi_to_psinorm(psi): *( (psi, xpoint) for psi, xpoint in zip(self.psi_sep, self.x_points) - if psi_to_psinorm(psi) < self.user_options.psinorm_sol + if self._psi_to_psinorm(psi) < self.user_options.psinorm_sol ) ) diff --git a/hypnotoad/core/equilibrium.py b/hypnotoad/core/equilibrium.py index 1a0cd925..4ea169da 100644 --- a/hypnotoad/core/equilibrium.py +++ b/hypnotoad/core/equilibrium.py @@ -4902,3 +4902,41 @@ def plotSeparatrix( pyplot.scatter(R, Z, marker=marker, label=region.name, **kwargs) else: pyplot.plot(R, Z, label=region.name, **kwargs) + + def plotHighlightRegion( + self, psiN_bounds, *, color="orange", alpha=0.5, npoints=100, **kwargs + ): + """ + Highlight a region between given psiN values, may be useful for example to show + the extent of a simulation grid without plotting all the grid points. + + Parameters + ---------- + psiN_bounds : (float, float) + Inner and outer values of psiN to highlight between. + color : str, default "orange" + Color to use for highlight. + alpha : float, default 0.5 + Transparency level for the highlighted region. + npoints : int, default 100 + When `separate_contours=True`, number of points used in the grid + discretizing psi. + **kwargs + Extra keyword arguments passed to `pyplot.contourf()`. + """ + from matplotlib import pyplot + + psi_bounds = tuple(self._psinorm_to_psi(x) for x in psiN_bounds) + + R = numpy.linspace(self.Rmin, self.Rmax, npoints) + Z = numpy.linspace(self.Zmin, self.Zmax, npoints) + + pyplot.contourf( + R, + Z, + self.psi(R[:, numpy.newaxis], Z[numpy.newaxis, :]).T, + psi_bounds, + colors=color, + alpha=alpha, + **kwargs, + ) diff --git a/hypnotoad/scripts/hypnotoad_plot_equilibrium.py b/hypnotoad/scripts/hypnotoad_plot_equilibrium.py index eae30b26..b6c90778 100755 --- a/hypnotoad/scripts/hypnotoad_plot_equilibrium.py +++ b/hypnotoad/scripts/hypnotoad_plot_equilibrium.py @@ -52,6 +52,19 @@ def main(): default=False, help="Color psi contours", ) + parser.add_argument( + "--highlight-region", + nargs=2, + type=float, + default=None, + help="Highlight a region between the two values of normalised psi given by the " + "arguments to this flag.", + ) + parser.add_argument( + "--highlight-color", + default="orange", + help="Color to use for region highlighted by `--hilight-region`", + ) parser.add_argument( "--show", action="store_false", @@ -84,6 +97,9 @@ def main(): ncontours=args.n_contours, labels=args.psi_labels, colors=colors, linestyles="-" ) + if args.highlight_region is not None: + eq.plotHighlightRegion(args.highlight_region, color=args.highlight_color) + if args.wall: eq.plotWall(linewidth=args.wall_width) From 03d92a6549c050e75646f093e882fc33ae0862fc Mon Sep 17 00:00:00 2001 From: John Omotani Date: Wed, 30 Nov 2022 12:35:27 +0000 Subject: [PATCH 4/6] Allow list of colors for separatrices in hypnotoad-plot-equilibrium --- hypnotoad/scripts/hypnotoad_plot_equilibrium.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/hypnotoad/scripts/hypnotoad_plot_equilibrium.py b/hypnotoad/scripts/hypnotoad_plot_equilibrium.py index b6c90778..ae705be0 100755 --- a/hypnotoad/scripts/hypnotoad_plot_equilibrium.py +++ b/hypnotoad/scripts/hypnotoad_plot_equilibrium.py @@ -38,7 +38,10 @@ def main(): "equilibrium.", ) parser.add_argument( - "--separatrix-color", default="blue", help="Color of the separatrix contour" + "--separatrix-color", + default="blue", + nargs="*", + help="Color (or list of colors) for the separatrix contour(s)", ) parser.add_argument( "--psi-labels", action="store_true", default=False, help="Label psi contours" From 73dbad6f80b371b28cb3d75ea3efcf74f76b4e9a Mon Sep 17 00:00:00 2001 From: John Omotani Date: Wed, 30 Nov 2022 12:50:45 +0000 Subject: [PATCH 5/6] Ensure primary separatrix plotted on top in Equilibrium.plotSeparatrix() --- hypnotoad/core/equilibrium.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/hypnotoad/core/equilibrium.py b/hypnotoad/core/equilibrium.py index 4ea169da..c8bb084e 100644 --- a/hypnotoad/core/equilibrium.py +++ b/hypnotoad/core/equilibrium.py @@ -4887,13 +4887,17 @@ def plotSeparatrix( R = numpy.linspace(self.Rmin, self.Rmax, npoints) Z = numpy.linspace(self.Zmin, self.Zmax, npoints) - pyplot.contour( - R, - Z, - self.psi(R[:, numpy.newaxis], Z[numpy.newaxis, :]).T, - self.psi_sep, - **kwargs, - ) + for i, psi_val in reversed(tuple(enumerate(self.psi_sep))): + this_kwargs = { + k: v[i] if isinstance(v, Sequence) else v for k, v in kwargs.items() + } + pyplot.contour( + R, + Z, + self.psi(R[:, numpy.newaxis], Z[numpy.newaxis, :]).T, + (psi_val,), + **this_kwargs, + ) else: for region in self.regions.values(): R = [p.R for p in region] From 01441e0039a72ce2d4eef47854658041a4a24945 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Wed, 30 Nov 2022 19:57:34 +0000 Subject: [PATCH 6/6] Fix typo --- README.md | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 4a067c11..5ed846ae 100644 --- a/README.md +++ b/README.md @@ -126,9 +126,8 @@ Hypnotoad provides several executables for working with equilibria and grid file equilibrium (flux surfaces, wall and separatrix) from a geqdsk file - `hypnotoad-plot-grid-cells` creates a plot of the grid cells from a grid file generated by hypnotoad -- `hypnotoad-recreate-inputs` writes extracts from a grid file copies of the - input YAML file and geqdsk file that were used to create the grid file - originally +- `hypnotoad-recreate-inputs` extracts from a grid file copies of the input + YAML file and geqdsk file that were used to create the grid file originally For more information, pass the `--help` flag to any of these commands.