Add a View menu, improve grid plotting

doc/grid-file.rst

@@ -138,6 +138,27 @@ BoutMesh writes three poloidal coordinates to the grid file:
        full poloidal turn around a closed flux surface. Not calculated on open
        flux surfaces.
 
+Wall location
++++++++++++++
+
+If a wall is defined in the equilibrium then the coordinates of the
+closed wall is saved:
+
+.. list-table::
+   :widths: 30 70
+
+   * - ``closed_wall_R``
+
+     - Major radius locations [in meters] of points on the wall. This
+       array forms a closed loop so the last element is the same as
+       the first.
+
+   * - ``closed_wall_Z``
+
+     - Height locations [in meters] of points on the wall, the same
+       number of points as ``closed_wall_R``. This array forms a
+       closed loop so the last element is the same as the first.
+
 2D arrays
 ---------
 
@@ -211,6 +232,21 @@ Magnetic field quantities
 
      - Total magnetic field.
 
+Boundary quantities
++++++++++++++++++++
+
+.. list-table::
+   :widths: 30 70
+
+   * - ``penalty_mask``
+
+     - A 2D mask indicating whether a cell is inside or outside the
+       wall. It's value is 1 for cells entirely outside the wall; 0
+       for cells entirely inside the wall. Cells that cross the wall
+       are given a penalty proportional to the fraction of the cell
+       poloidal length that is inside the wall.
+
+
 Integral quantities
 +++++++++++++++++++
 

doc/whats-new.md

@@ -1,6 +1,20 @@
 Release history
 ===============
 
+### Bug fixes
+
+
+### New features
+
+- Radial grid line construction can recover from failure and generate
+  a rough grid to help visual inspection when option
+  `follow_perpendicular_recover` is set to True
+- A `View` menu enables the grid plot to be customised, with cell edges, corners,
+  grid lines and other components.
+- `penalty_mask` is calculated and written to the grid file, based on intersection
+  of the grid with the wall. This enables immersed boundary conditions.
+- Wall coordinates are written to output grid as `closed_wall_R` and `closed_wall_Z`
+
 0.5.2 (13th March 2023)
 -------------------------
 

hypnotoad/core/mesh.py

@@ -42,6 +42,7 @@
 
 from .equilibrium import (
     calc_distance,
+    find_intersections,
     Equilibrium,
     EquilibriumRegion,
     Point2D,
@@ -748,6 +749,56 @@ def fillRZ(self):
             self.Rxy.corners[-1, -1] = xpoint.R
             self.Zxy.corners[-1, -1] = xpoint.Z
 
+    def calcPenaltyMask(self, equilibrium):
+        """
+        Uses self.Rxy and self.Zxy to calculate which cells are outside the wall.
+
+        Assumes that self.Rxy and self.Zxy have already been calculated by
+        self.fillRZ()
+
+        Sets and returns a penalty_mask 2D array.
+        """
+        self.penalty_mask = numpy.zeros((self.nx, self.ny))
+
+        # A point inside the wall.
+        p0 = Point2D(
+            (equilibrium.Rmax + equilibrium.Rmin) / 2,
+            (equilibrium.Zmax + equilibrium.Zmin) / 2,
+        )
+
+        for i in range(self.nx):
+            for j in range(self.ny):
+                # Check if cell Y edges are outside the wall
+                p1 = Point2D(self.Rxy.ylow[i, j], self.Zxy.ylow[i, j])
+                intersects = find_intersections(equilibrium.closed_wallarray, p0, p1)
+                p1_outside = (
+                    False if intersects is None else intersects.shape[0] % 2 == 1
+                )
+
+                p2 = Point2D(self.Rxy.ylow[i, j + 1], self.Zxy.ylow[i, j + 1])
+                intersects = find_intersections(equilibrium.closed_wallarray, p0, p2)
+                p2_outside = (
+                    False if intersects is None else intersects.shape[0] % 2 == 1
+                )
+
+                if p1_outside and p2_outside:
+                    # Both ends of the cell are outside the wall
+                    self.penalty_mask[i, j] = 1.0
+                elif p1_outside or p2_outside:
+                    # Cell crosses the wall
+                    intersects = find_intersections(
+                        equilibrium.closed_wallarray, p1, p2
+                    )
+                    if intersects is None:
+                        # Something odd going on
+                        continue
+                    pi = Point2D(
+                        intersects[0, 0], intersects[0, 1]
+                    )  # Intersection point
+                    self.penalty_mask[i, j] = calc_distance(
+                        p1 if p1_outside else p2, pi
+                    ) / calc_distance(p1, p2)
+
     def getRZBoundary(self):
         # Upper value of ylow array logically overlaps with the lower value in the upper
         # neighbour. They should be close, but aren't guaranteed to be identical already
@@ -2738,6 +2789,8 @@ def calculateRZ(self):
             region.fillRZ()
         for region in self.regions.values():
             region.getRZBoundary()
+        for region in self.regions.values():
+            region.calcPenaltyMask(self.equilibrium)
 
     def geometry(self):
         """
@@ -3019,17 +3072,86 @@ def smoothnl(self, varname):
             if change < 1.0e-3:
                 break
 
-    def plotGridLines(self, **kwargs):
+    def plotGridCellEdges(self, ax=None, **kwargs):
+        """
+        Plot lines between cell corners
+        """
         from matplotlib import pyplot
         from cycler import cycle
 
         colors = cycle(pyplot.rcParams["axes.prop_cycle"].by_key()["color"])
 
+        if ax is None:
+            _, ax = pyplot.subplots(1)
+
+        for region in self.regions.values():
+            c = next(colors)
+            label = region.myID
+            for i in range(region.nx + 1):
+                ax.plot(
+                    region.Rxy.corners[i, :],
+                    region.Zxy.corners[i, :],
+                    c=c,
+                    label=label,
+                    **kwargs,
+                )
+                label = None
+            label = region.myID
+            for j in range(region.ny + 1):
+                ax.plot(
+                    region.Rxy.corners[:, j],
+                    region.Zxy.corners[:, j],
+                    c=c,
+                    label=None,
+                    **kwargs,
+                )
+                label = None
+
+    def plotPenaltyMask(self, ax=None, **kwargs):
+        from matplotlib import pyplot
+
+        if ax is None:
+            _, ax = pyplot.subplots(1)
+
+        for region in self.regions.values():
+            for i in range(region.nx):
+                for j in range(region.ny):
+                    penalty = region.penalty_mask[i, j]
+                    if penalty > 0.01:
+                        # Add a polygon
+                        ax.fill(
+                            [
+                                region.Rxy.corners[i, j],
+                                region.Rxy.corners[i, j + 1],
+                                region.Rxy.corners[i + 1, j + 1],
+                                region.Rxy.corners[i + 1, j],
+                            ],
+                            [
+                                region.Zxy.corners[i, j],
+                                region.Zxy.corners[i, j + 1],
+                                region.Zxy.corners[i + 1, j + 1],
+                                region.Zxy.corners[i + 1, j],
+                            ],
+                            "k",
+                            alpha=penalty,
+                        )
+
+    def plotGridLines(self, ax=None, **kwargs):
+        from matplotlib import pyplot
+        from cycler import cycle
+
+        colors = cycle(pyplot.rcParams["axes.prop_cycle"].by_key()["color"])
+
+        if ax is None:
+            fig, ax = pyplot.subplots(1)
+        else:
+            fig = ax.figure
+
         for region in self.regions.values():
             c = next(colors)
             label = region.myID
             for i in range(region.nx):
-                pyplot.plot(
+                ax.plot(
                     region.Rxy.centre[i, :],
                     region.Zxy.centre[i, :],
                     c=c,
@@ -3039,19 +3161,20 @@ def plotGridLines(self, **kwargs):
                 label = None
             label = region.myID
             for j in range(region.ny):
-                pyplot.plot(
+                ax.plot(
                     region.Rxy.centre[:, j],
                     region.Zxy.centre[:, j],
                     c=c,
                     label=None,
                     **kwargs,
                 )
                 label = None
-        l = pyplot.legend()
+        l = fig.legend()
         l.set_draggable(True)
 
     def plotPoints(
         self,
+        centers=True,
         xlow=False,
         ylow=False,
         corners=False,
@@ -3092,14 +3215,15 @@ def plotPoints(
 
             if "scatter" in plot_types:
                 m = iter(markers)
-                ax.scatter(
-                    region.Rxy.centre,
-                    region.Zxy.centre,
-                    marker=next(m),
-                    c=c,
-                    label=region.myID,
-                    **kwargs,
-                )
+                if centers:
+                    ax.scatter(
+                        region.Rxy.centre,
+                        region.Zxy.centre,
+                        marker=next(m),
+                        c=c,
+                        label=region.myID,
+                        **kwargs,
+                    )
                 if xlow:
                     ax.scatter(
                         region.Rxy.xlow, region.Zxy.xlow, marker=next(m), c=c, **kwargs
@@ -3536,6 +3660,19 @@ def addFromRegionsXArray(name):
         if hasattr(next(iter(self.equilibrium.regions.values())), "pressure"):
             addFromRegions("pressure")
 
+        # Penalty mask
+        self.penalty_mask = BoutArray(
+            numpy.zeros((self.nx, self.ny)),
+            attributes={
+                "bout_type": "Field2D",
+                "description": (
+                    "1 if cell is entirely outside wall, " "0 if entirely inside wall"
+                ),
+            },
+        )
+        for region in self.regions.values():
+            self.penalty_mask[self.region_indices[region.myID]] = region.penalty_mask
+
     def writeArray(self, name, array, f):
         f.write(name, BoutArray(array.centre, attributes=array.attributes))
         f.write(
@@ -3590,10 +3727,17 @@ def writeGridfile(self, filename):
             if hasattr(self.equilibrium, "psi_bdry_gfile"):
                 f.write("psi_bdry_gfile", self.equilibrium.psi_bdry_gfile)
 
+            if hasattr(self.equilibrium, "closed_wallarray"):
+                f.write("closed_wall_R", self.equilibrium.closed_wallarray[:, 0])
+                f.write("closed_wall_Z", self.equilibrium.closed_wallarray[:, 1])
+
             # write the 2d fields
             for name in self.fields_to_output:
                 self.writeArray(name, self.__dict__[name], f)
 
+            # penalty_mask is defined for each cell
+            f.write("penalty_mask", self.penalty_mask)
+
             # write corner positions, as these may be useful for plotting, etc.
             for name in ["Rxy", "Zxy"]:
                 self.writeCorners(name, self.__dict__[name], f)