Merge pull request #139 from boutproject/plot-equilib

Convenience script for plotting an equilibrium

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
 ----------
 

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)

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

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

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()

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 = "