Add gridded collocation option (#3065)

* add new collocation point option

* tidy up use of collocation points function

* use mask

* tidy up after refactor and give n_points default value

bluemira/equilibria/find.py

@@ -809,6 +809,7 @@ def in_plasma(
     psi: npt.NDArray[np.float64],
     o_points: Optional[List[Opoint]] = None,
     x_points: Optional[List[Xpoint]] = None,
+    include_edges: bool = False,
 ) -> npt.NDArray[np.float64]:
     """
     Get a psi-shaped mask of psi where 1 is inside the plasma, 0 outside.
@@ -833,11 +834,14 @@ 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: npt.NDArray[np.float64], z: npt.NDArray[np.float64], zone: npt.NDArray[np.float64]
+    x: npt.NDArray[np.float64],
+    z: npt.NDArray[np.float64],
+    zone: npt.NDArray[np.float64],
+    include_edges: bool = False,
 ):
     """
     Get a masking matrix for a specified zone.
@@ -856,7 +860,7 @@ def in_zone(
     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)
@@ -865,6 +869,7 @@ def _in_plasma(
     z: npt.NDArray[np.float64],
     mask: npt.NDArray[np.float64],
     sep: npt.NDArray[np.float64],
+    include_edges: bool = False,
 ) -> npt.NDArray[np.float64]:
     """
     Get a masking matrix for a specified zone. JIT compilation utility.
@@ -887,6 +892,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

@@ -24,6 +24,7 @@
 from bluemira.equilibria.coils import CoilSet
 from bluemira.equilibria.equilibrium import Equilibrium
 from bluemira.equilibria.error import EquilibriaError
+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 (
@@ -177,6 +178,7 @@ class PointType(Enum):
     ARC_PLUS_EXTREMA = auto()
     RANDOM = auto()
     RANDOM_PLUS_EXTREMA = auto()
+    GRID_POINTS = auto()
 
 
 @dataclass
@@ -190,10 +192,11 @@ class Collocation:
 
 
 def collocation_points(
-    n_points: int,
     plasma_boundary: Coordinates,
     point_type: PointType,
+    n_points: int = 10,
     seed: Optional[int] = None,
+    grid_num: Optional[Tuple[int, int]] = None,
 ) -> Collocation:
     """
     Create a set of collocation points for use wih spherical harmonic
@@ -204,21 +207,29 @@ def collocation_points(
     - equispaced points on an arc of fixed radius,
     - equispaced points on an arc plus extrema,
     - random points within a circle enclosed by the LCFS,
-    - random points plus extrema.
+    - random points plus extrema,
+    - a grid of points containing the LCFS.
 
     Parameters
     ----------
     n_points:
         Number of points/targets (not including extrema - these are added
-        automatically if relevant).
+        automatically if relevant). For use with point_type 'arc',
+        'arc_plus_extrema', 'random', 'random_plus_extrema', or 'grid_num'.
+        For 'grid_num' it will create an n_points by n_points grid (see
+        grid_num for a non square grid.)
     plasma_boundary:
         XZ coordinates of the plasma boundary
     point_type:
         Method for creating a set of points: 'arc', 'arc_plus_extrema',
-        'random', or 'random_plus_extrema'
+        'random', or 'random_plus_extrema', 'grid_points'
     seed:
         Seed value to use with a random point distribution, defaults
-        to `RNGSeeds.equilibria_harmonics.value`.
+        to `RNGSeeds.equilibria_harmonics.value`. For use with 'random'
+        or 'random_plus_extrema' point_type.
+    grid_num:
+        Tuple with the number of desired grid points in the x and z direction.
+        For use with 'grid_points' point_type.
 
     Returns
     -------
@@ -289,6 +300,31 @@ def collocation_points(
         collocation_r = np.sqrt(collocation_x**2 + collocation_z**2)
         collocation_theta = np.arctan2(collocation_x, collocation_z)
 
+    if point_type is PointType.GRID_POINTS:
+        # Create uniform, rectangular grid using max and min LCFS values
+        if grid_num is None:
+            grid_num = (n_points, n_points)
+        grid_num_x, grid_num_z = grid_num
+        rect_grid = Grid(
+            np.amin(x_bdry),
+            np.amax(x_bdry),
+            np.amin(z_bdry),
+            np.amax(z_bdry),
+            nx=grid_num_x,
+            nz=grid_num_z,
+        )
+
+        # Only use grid points that are within LCFS
+        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)
+        collocation_theta = np.arctan2(collocation_x, collocation_z)
+
     return Collocation(collocation_r, collocation_theta, collocation_x, collocation_z)
 
 
@@ -455,6 +491,7 @@ def spherical_harmonic_approximation(
     eq: Equilibrium,
     n_points: int = 8,
     point_type: PointType = PointType.ARC_PLUS_EXTREMA,
+    grid_num: Optional[str] = None,
     acceptable_fit_metric: float = 0.01,
     plot: bool = False,
     nlevels: int = 50,
@@ -490,6 +527,10 @@ def spherical_harmonic_approximation(
         in the x- and z-directions (4 points total),
         - 'random',
         - 'random_plus_extrema'.
+        - 'grid_points'
+    grid_num:
+        Number of points in x-direction and z-direction,
+        to use with grid point distribution.
     acceptable_fit_metric:
         Value between 0 and 1 chosen by user (default=0.01).
         If the LCFS found using the SH approximation method perfectly matches the
@@ -550,10 +591,11 @@ def spherical_harmonic_approximation(
 
     # Create the set of collocation points within the LCFS for the SH calculations
     collocation = collocation_points(
-        n_points,
         original_LCFS,
         point_type,
+        n_points,
         seed,
+        grid_num,
     )
 
     # SH amplitudes needed to produce an approximation of vacuum psi contribution

examples/equilibria/Spherical_Approximation_how_it_works.ex.py

@@ -206,9 +206,9 @@
 
 # Create the set of collocation points for the harmonics
 collocation = collocation_points(
-    n,
     original_LCFS,
     PointType.ARC_PLUS_EXTREMA,
+    n,
     seed=15,
 )
 

tests/equilibria/test_harmonics.py

@@ -109,6 +109,7 @@ def test_coil_harmonic_amplitude_matrix():
 
 def test_collocation_points():
     n_points = 8
+    grid_num = (10, 10)
 
     x = [1, 1.5, 2, 2.1, 2, 1.5, 1, 0.9, 1]
     z = [-1.8, -1.9, -1.8, 0, 1.8, 1.9, 1.8, 0, -1.8]
@@ -118,23 +119,29 @@ def test_collocation_points():
     point_type_2 = PointType.ARC_PLUS_EXTREMA
     point_type_3 = PointType.RANDOM
     point_type_4 = PointType.RANDOM_PLUS_EXTREMA
+    point_type_5 = PointType.GRID_POINTS
 
-    colloc1 = collocation_points(n_points, plasma_boundary, point_type_1)
-    colloc2 = collocation_points(n_points, plasma_boundary, point_type_2)
-    colloc3 = collocation_points(n_points, plasma_boundary, point_type_3)
-    colloc4 = collocation_points(n_points, plasma_boundary, point_type_4)
+    colloc1 = collocation_points(plasma_boundary, point_type_1, n_points=n_points)
+    colloc2 = collocation_points(plasma_boundary, point_type_2, n_points=n_points)
+    colloc3 = collocation_points(plasma_boundary, point_type_3, n_points=n_points)
+    colloc4 = collocation_points(plasma_boundary, point_type_4, n_points=n_points)
+    colloc5 = collocation_points(plasma_boundary, point_type_5, grid_num=grid_num)
 
     assert colloc1.r.shape[0] == 8
     assert colloc2.r.shape[0] == 12
     assert colloc3.r.shape[0] == 8
     assert colloc4.r.shape[0] == 12
+    assert colloc5.r.shape[0] == 64
 
     for x, z in zip(colloc2.x, colloc2.z):
         assert in_polygon(x, z, plasma_boundary.xz.T, include_edges=True)
 
     for x, z in zip(colloc4.x, colloc4.z):
         assert in_polygon(x, z, plasma_boundary.xz.T, include_edges=True)
 
+    for x, z in zip(colloc5.x, colloc5.z):
+        assert in_polygon(x, z, plasma_boundary.xz.T, include_edges=True)
+
 
 def test_coils_outside_sphere_vacuum_psi():
     eq = Equilibrium.from_eqdsk(Path(TEST_PATH, "SH_test_file.json").as_posix())
@@ -161,7 +168,9 @@ def test_get_psi_harmonic_amplitudes():
     eq = Equilibrium.from_eqdsk(Path(TEST_PATH, "SH_test_file.json").as_posix())
 
     test_colocation = collocation_points(
-        n_points=18, plasma_boundary=eq.get_LCFS(), point_type=PointType.ARC
+        plasma_boundary=eq.get_LCFS(),
+        point_type=PointType.ARC,
+        n_points=18,
     )
 
     sh_coil_names, _ = coils_outside_lcfs_sphere(eq)