From b6be06dbed03fa9246d44df20bb97681b762f47a Mon Sep 17 00:00:00 2001 From: geograham Date: Tue, 19 Mar 2024 11:53:25 +0000 Subject: [PATCH] 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 | 13 +++-- .../harmonics/harmonics_approx_functions.py | 54 ++++++++++++++++--- ...Spherical_Approximation_how_it_works.ex.py | 2 +- tests/equilibria/test_harmonics.py | 19 +++++-- 4 files changed, 72 insertions(+), 16 deletions(-) diff --git a/bluemira/equilibria/find.py b/bluemira/equilibria/find.py index 46ad2d9bfe..88fb1f815e 100644 --- a/bluemira/equilibria/find.py +++ b/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 diff --git a/bluemira/equilibria/optimisation/harmonics/harmonics_approx_functions.py b/bluemira/equilibria/optimisation/harmonics/harmonics_approx_functions.py index e254cf1bd3..bad4a13995 100644 --- a/bluemira/equilibria/optimisation/harmonics/harmonics_approx_functions.py +++ b/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 diff --git a/examples/equilibria/Spherical_Approximation_how_it_works.ex.py b/examples/equilibria/Spherical_Approximation_how_it_works.ex.py index c65f1d1da2..82b210a66e 100644 --- a/examples/equilibria/Spherical_Approximation_how_it_works.ex.py +++ b/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, ) diff --git a/tests/equilibria/test_harmonics.py b/tests/equilibria/test_harmonics.py index d65e3ae123..212c808915 100644 --- a/tests/equilibria/test_harmonics.py +++ b/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,16 +119,19 @@ 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) @@ -135,6 +139,9 @@ def test_collocation_points(): 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)