use mask

bluemira/equilibria/find.py

@@ -932,6 +932,7 @@ def in_plasma(
     psi: np.ndarray,
     o_points: Optional[List[Opoint]] = None,
     x_points: Optional[List[Xpoint]] = None,
+    include_edges: bool = False,
 ) -> np.ndarray:
     """
     Get a psi-shaped mask of psi where 1 is inside the plasma, 0 outside.
@@ -956,10 +957,10 @@ def in_plasma(
     """
     mask = np.zeros_like(psi)
     lcfs, _ = find_LCFS_separatrix(x, z, psi, o_points=o_points, x_points=x_points)
-    return _in_plasma(x, z, mask, lcfs.xz.T)
+    return _in_plasma(x, z, mask, lcfs.xz.T, include_edges)
 
 
-def in_zone(x: np.ndarray, z: np.ndarray, zone: np.ndarray):
+def in_zone(x: np.ndarray, z: np.ndarray, zone: np.ndarray, include_edges: bool = False):
     """
     Get a masking matrix for a specified zone.
 
@@ -977,12 +978,16 @@ def in_zone(x: np.ndarray, z: np.ndarray, zone: np.ndarray):
     The masking array where 1 denotes inside the zone, and 0 outside
     """
     mask = np.zeros_like(x)
-    return _in_plasma(x, z, mask, zone)
+    return _in_plasma(x, z, mask, zone, include_edges)
 
 
 @nb.jit(nopython=True, cache=True)
 def _in_plasma(
-    x: np.ndarray, z: np.ndarray, mask: np.ndarray, sep: np.ndarray
+    x: np.ndarray,
+    z: np.ndarray,
+    mask: np.ndarray,
+    sep: np.ndarray,
+    include_edges: bool = False,
 ) -> np.ndarray:
     """
     Get a masking matrix for a specified zone. JIT compilation utility.
@@ -1005,6 +1010,6 @@ def _in_plasma(
     n, m = x.shape
     for i in range(n):
         for j in range(m):
-            if in_polygon(x[i, j], z[i, j], sep):
+            if in_polygon(x[i, j], z[i, j], sep, include_edges):
                 mask[i, j] = 1
     return mask

bluemira/equilibria/optimisation/harmonics/harmonics_approx_functions.py

@@ -23,13 +23,13 @@
 from bluemira.base.look_and_feel import bluemira_debug, bluemira_print
 from bluemira.equilibria.coils import CoilSet
 from bluemira.equilibria.equilibrium import Equilibrium
+from bluemira.equilibria.find import in_zone
 from bluemira.equilibria.grid import Grid
 from bluemira.equilibria.plotting import PLOT_DEFAULTS
 from bluemira.geometry.coordinates import (
     Coordinates,
     get_area_2d,
     get_intersect,
-    in_polygon,
     polygon_in_polygon,
 )
 from bluemira.geometry.face import BluemiraFace
@@ -307,19 +307,11 @@ def collocation_points(
         )
 
         # Only use grid points that are within LCFS
-        collocation_x = np.array([])
-        collocation_z = np.array([])
-        for i in np.arange(grid_num_x):
-            for j in np.arange(grid_num_x):
-                in_poly = in_polygon(
-                    x=rect_grid.x[i, j],
-                    z=rect_grid.z[i, j],
-                    poly=plasma_boundary.xz.T,
-                    include_edges=True,
-                )
-                if in_poly:
-                    collocation_x = np.append(collocation_x, rect_grid.x[i, j])
-                    collocation_z = np.append(collocation_z, rect_grid.z[i, j])
+        mask = in_zone(
+            rect_grid.x, rect_grid.z, plasma_boundary.xz.T, include_edges=True
+        )
+        collocation_x = rect_grid.x[mask == 1]
+        collocation_z = rect_grid.z[mask == 1]
 
         # Spherical coordinates
         collocation_r = np.sqrt(collocation_x**2 + collocation_z**2)