diff --git a/README.md b/README.md index b557f996..5ed846ae 100644 --- a/README.md +++ b/README.md @@ -111,6 +111,27 @@ 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` 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 ---------- 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/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 bff5924f..c8bb084e 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,117 @@ 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) + + 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] + 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) + + 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 new file mode 100755 index 00000000..ae705be0 --- /dev/null +++ b/hypnotoad/scripts/hypnotoad_plot_equilibrium.py @@ -0,0 +1,127 @@ +#!/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", + 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" + ) + 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( + "--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", + 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.highlight_region is not None: + eq.plotHighlightRegion(args.highlight_region, color=args.highlight_color) + + 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 = "