Dp/signed distance (#817)

- Add ``use_signed_distance`` flag to ``PlasmaVesselDistance`` which
will use a signed distance as the target, which is positive when the
plasma is inside of the vessel surface and negative if the plasma is
outside of the vessel surface, to allow optimizer to distinguish if the
equilbrium surface exits the vessel surface and guard against it by
targeting a positive signed distance.
- Deprecates `alpha` in favor of `softmin_alpha` to avoid confusion with
other quantities similarly-named in the code

Resolves #1070 
- [x] vectorize the signed distance algorithm
- [x] use signed dist with softmin
- [x] remove redundant or unneeded tests
- [x] decide on keeping Circular dist objective

CHANGELOG.md

@@ -1,6 +1,10 @@
 Changelog
 =========
 
+New Features
+
+- Add ``use_signed_distance`` flag to ``PlasmaVesselDistance`` which will use a signed distance as the target, which is positive when the plasma is inside of the vessel surface and negative if the plasma is outside of the vessel surface, to allow optimizer to distinguish if the equilbrium surface exits the vessel surface and guard against it by targeting a positive signed distance.
+
 v0.12.1
 -------
 

desc/objectives/_geometry.py

@@ -2,16 +2,16 @@
 
 import numpy as np
 
-from desc.backend import jnp
-from desc.compute import get_profiles, get_transforms, rpz2xyz
+from desc.backend import jnp, vmap
+from desc.compute import get_profiles, get_transforms, rpz2xyz, xyz2rpz
 from desc.compute.utils import _compute as compute_fun
 from desc.compute.utils import safenorm
 from desc.grid import LinearGrid, QuadratureGrid
-from desc.utils import Timer, warnif
+from desc.utils import Timer, errorif, parse_argname_change, warnif
 
 from .normalization import compute_scaling_factors
 from .objective_funs import _Objective
-from .utils import softmin
+from .utils import check_if_points_are_inside_perimeter, softmin
 
 
 class AspectRatio(_Objective):
@@ -519,10 +519,10 @@ class PlasmaVesselDistance(_Objective):
     points on surface corresponding to the grid that the plasma-vessel distance
     is evaluated at, which can cause cusps or regions of very large curvature.
 
-    NOTE: When use_softmin=True, ensures that alpha*values passed in is
+    NOTE: When use_softmin=True, ensures that softmin_alpha*values passed in is
     at least >1, otherwise the softmin will return inaccurate approximations
     of the minimum. Will automatically multiply array values by 2 / min_val if the min
-    of alpha*array is <1. This is to avoid inaccuracies that arise when values <1
+    of softmin_alpha*array is <1. This is to avoid inaccuracies when values <1
     are present in the softmin, which can cause inaccurate mins or even incorrect
     signs of the softmin versus the actual min.
 
@@ -565,20 +565,26 @@ class PlasmaVesselDistance(_Objective):
         Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``.
     use_softmin: bool, optional
         Use softmin or hard min.
-    surface_fixed: bool, optional
-        Whether the surface the distance from the plasma is computed to
-        is fixed or not. If True, the surface is fixed and its coordinates are
-        precomputed, which saves on computation time during optimization, and
-        self.things = [eq] only.
-        If False, the surface coordinates are computed at every iteration.
+    use_signed_distance: bool, optional
+        Whether to use absolute value of distance or a signed distance, with d
+        being positive if the plasma is inside of the bounding surface, and
+        negative if outside of the bounding surface.
+        NOTE: ``plasma_grid`` and ``surface_grid`` must have the same
+        toroidal angle values for signed distance to be used.
+    eq_fixed, surface_fixed: bool, optional
+        Whether the eq/surface is fixed or not. If True, the eq/surface is fixed
+        and its coordinates are precomputed, which saves on computation time during
+        optimization, and self.things = [surface]/[eq] only.
+        If False, the eq/surface coordinates are computed at every iteration.
         False by default, so that self.things = [eq, surface]
-    alpha: float, optional
-        Parameter used for softmin. The larger alpha, the closer the softmin
-        approximates the hardmin. softmin -> hardmin as alpha -> infinity.
-        if alpha*array < 1, the underlying softmin will automatically multiply
-        the array by 2/min_val to ensure that alpha*array>1. Making alpha larger
-        than this minimum value will make the softmin a more accurate approximation
-        of the true min.
+        Both cannot be True.
+    softmin_alpha: float, optional
+        Parameter used for softmin. The larger softmin_alpha, the closer the softmin
+        approximates the hardmin. softmin -> hardmin as softmin_alpha -> infinity.
+        if softmin_alpha*array < 1, the underlying softmin will automatically multiply
+        the array by 2/min_val to ensure that softmin_alpha*array>1. Making
+        softmin_alpha larger than this minimum value will make the softmin a
+        more accurate approximation of the true min.
     name : str, optional
         Name of the objective function.
     """
@@ -601,20 +607,42 @@ def __init__(
         surface_grid=None,
         plasma_grid=None,
         use_softmin=False,
+        eq_fixed=False,
         surface_fixed=False,
-        alpha=1.0,
+        softmin_alpha=1.0,
         name="plasma-vessel distance",
+        use_signed_distance=False,
+        **kwargs,
     ):
         if target is None and bounds is None:
             bounds = (1, np.inf)
         self._surface = surface
         self._surface_grid = surface_grid
         self._plasma_grid = plasma_grid
         self._use_softmin = use_softmin
+        self._use_signed_distance = use_signed_distance
         self._surface_fixed = surface_fixed
-        self._alpha = alpha
+        self._eq_fixed = eq_fixed
+        self._eq = eq
+        errorif(
+            eq_fixed and surface_fixed, ValueError, "Cannot fix both eq and surface"
+        )
+
+        self._softmin_alpha = parse_argname_change(
+            softmin_alpha, kwargs, "alpha", "softmin_alpha"
+        )
+        errorif(
+            len(kwargs) != 0,
+            AssertionError,
+            f"PlasmaVesselDistance got unexpected keyword argument: {kwargs.keys()}",
+        )
+        things = []
+        if not eq_fixed:
+            things.append(eq)
+        if not surface_fixed:
+            things.append(surface)
         super().__init__(
-            things=[eq, self._surface] if not surface_fixed else [eq],
+            things=things,
             target=target,
             bounds=bounds,
             weight=weight,
@@ -636,11 +664,15 @@ def build(self, use_jit=True, verbose=1):
             Level of output.
 
         """
-        eq = self.things[0]
-        surface = self._surface if self._surface_fixed else self.things[1]
-        # if things[1] is different than self._surface, update self._surface
-        if surface != self._surface:
-            self._surface = surface
+        if self._eq_fixed:
+            eq = self._eq
+            surface = self.things[0]
+        elif self._surface_fixed:
+            eq = self.things[0]
+            surface = self._surface
+        else:
+            eq = self.things[0]
+            surface = self.things[1]
         if self._surface_grid is None:
             surface_grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP)
         else:
@@ -660,6 +692,18 @@ def build(self, use_jit=True, verbose=1):
             "Plasma grid includes interior points, should be rho=1.",
         )
 
+        # TODO: How to use with generalized toroidal angle?
+        errorif(
+            self._use_signed_distance
+            and not np.allclose(
+                plasma_grid.nodes[plasma_grid.unique_zeta_idx, 2],
+                surface_grid.nodes[surface_grid.unique_zeta_idx, 2],
+            ),
+            ValueError,
+            "Plasma grid and surface grid must contain points only at the "
+            "same zeta values in order to use signed distance",
+        )
+
         self._dim_f = surface_grid.num_nodes
         self._equil_data_keys = ["R", "phi", "Z"]
         self._surface_data_keys = ["x"]
@@ -714,7 +758,15 @@ def build(self, use_jit=True, verbose=1):
             )["x"]
             surface_coords = rpz2xyz(surface_coords)
             self._constants["surface_coords"] = surface_coords
-
+        elif self._eq_fixed:
+            data_eq = compute_fun(
+                self._eq,
+                self._equil_data_keys,
+                params=self._eq.params_dict,
+                transforms=equil_transforms,
+                profiles=equil_profiles,
+            )
+            self._constants["data_equil"] = data_eq
         timer.stop("Precomputing transforms")
         if verbose > 1:
             timer.disp("Precomputing transforms")
@@ -725,14 +777,15 @@ def build(self, use_jit=True, verbose=1):
 
         super().build(use_jit=use_jit, verbose=verbose)
 
-    def compute(self, equil_params, surface_params=None, constants=None):
+    def compute(self, params_1, params_2=None, constants=None):
         """Compute plasma-surface distance.
 
         Parameters
         ----------
-        equil_params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        surface_params : dict
+        params_1 : dict
+            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict,
+            if eq_fixed is False, else the surface degrees of freedom
+        params_2 : dict
             Dictionary of surface degrees of freedom, eg Surface.params_dict
             Only needed if self._surface_fixed = False
         constants : dict
@@ -747,14 +800,25 @@ def compute(self, equil_params, surface_params=None, constants=None):
         """
         if constants is None:
             constants = self.constants
-        data = compute_fun(
-            "desc.equilibrium.equilibrium.Equilibrium",
-            self._equil_data_keys,
-            params=equil_params,
-            transforms=constants["equil_transforms"],
-            profiles=constants["equil_profiles"],
-        )
-        plasma_coords = rpz2xyz(jnp.array([data["R"], data["phi"], data["Z"]]).T)
+        if self._eq_fixed:
+            surface_params = params_1
+        elif self._surface_fixed:
+            equil_params = params_1
+        else:
+            equil_params = params_1
+            surface_params = params_2
+        if not self._eq_fixed:
+            data = compute_fun(
+                "desc.equilibrium.equilibrium.Equilibrium",
+                self._equil_data_keys,
+                params=equil_params,
+                transforms=constants["equil_transforms"],
+                profiles=constants["equil_profiles"],
+            )
+        else:
+            data = constants["data_equil"]
+        plasma_coords_rpz = jnp.array([data["R"], data["phi"], data["Z"]]).T
+        plasma_coords = rpz2xyz(plasma_coords_rpz)
         if self._surface_fixed:
             surface_coords = constants["surface_coords"]
         else:
@@ -766,12 +830,54 @@ def compute(self, equil_params, surface_params=None, constants=None):
                 profiles={},
             )["x"]
             surface_coords = rpz2xyz(surface_coords)
-        d = safenorm(plasma_coords[:, None, :] - surface_coords[None, :, :], axis=-1)
+
+        diff_vec = plasma_coords[:, None, :] - surface_coords[None, :, :]
+        d = safenorm(diff_vec, axis=-1)
+
+        point_signs = jnp.ones(surface_coords.shape[0])
+        if self._use_signed_distance:
+            surface_coords_rpz = xyz2rpz(surface_coords)
+
+            plasma_coords_rpz = plasma_coords_rpz.reshape(
+                constants["equil_transforms"]["grid"].num_zeta,
+                constants["equil_transforms"]["grid"].num_theta,
+                3,
+            )
+            surface_coords_rpz = surface_coords_rpz.reshape(
+                constants["surface_transforms"]["grid"].num_zeta,
+                constants["surface_transforms"]["grid"].num_theta,
+                3,
+            )
+
+            # loop over zeta planes
+            def fun(plasma_pts_at_zeta_plane, surface_pts_at_zeta_plane):
+                plasma_pts_at_zeta_plane = jnp.vstack(
+                    (plasma_pts_at_zeta_plane, plasma_pts_at_zeta_plane[0, :])
+                )
+
+                pt_sign = check_if_points_are_inside_perimeter(
+                    plasma_pts_at_zeta_plane[:, 0],
+                    plasma_pts_at_zeta_plane[:, 2],
+                    surface_pts_at_zeta_plane[:, 0],
+                    surface_pts_at_zeta_plane[:, 2],
+                )
+
+                return pt_sign
+
+            point_signs = vmap(fun, in_axes=0)(
+                plasma_coords_rpz, surface_coords_rpz
+            ).flatten()
+            # at end here, point_signs is either +/- 1  with
+            # positive meaning the surface pt
+            # is outside the plasma and -1 if the surface pt is
+            # inside the plasma
 
         if self._use_softmin:  # do softmin
-            return jnp.apply_along_axis(softmin, 0, d, self._alpha)
+            return (
+                jnp.apply_along_axis(softmin, 0, d, self._softmin_alpha) * point_signs
+            )
         else:  # do hardmin
-            return d.min(axis=0)
+            return d.min(axis=0) * point_signs
 
 
 class MeanCurvature(_Objective):

desc/objectives/utils.py

@@ -295,3 +295,78 @@ def _parse_callable_target_bounds(target, bounds, x):
     if bounds is not None and callable(bounds[1]):
         bounds = (bounds[0], bounds[1](x))
     return target, bounds
+
+
+def check_if_points_are_inside_perimeter(R, Z, Rcheck, Zcheck):
+    """Function to check if the given points is inside the given polyognal perimeter.
+
+    Rcheck, Zcheck are the points to check, and R, Z define the perimeter
+    in which to check. This function assumes that all points are in the same
+    plane. Function will return an array of signs (+/- 1), with positive sign meaning
+    the point is inside of the given perimeter, and a negative sign meaning the point
+    is outside of the given perimeter.
+
+    NOTE: it does not matter if the input coordinates are cylindrical (R,Z) or
+    cartesian (X,Y), these are equivalent as long as they are in the same phi plane.
+    This function will work even if points are not in the same phi plane, but the
+    input coordinates must then be the equivalent of cartesian coordinates for whatever
+    plane the points lie in.
+
+    Algorithm based off of "An Incremental Angle Point in Polygon Test",
+    K. Weiler, https://doi.org/10.1016/B978-0-12-336156-1.50012-4
+
+    Parameters
+    ----------
+    R,Z : ndarray
+        1-D arrays of coordinates of the points defining the polygonal
+        perimeter. The function will determine if the check point is inside
+        or outside of this perimeter. These should form a closed curve.
+    Rcheck, Zcheck : ndarray
+        coordinates of the points being checked if they are inside or outside of the
+        given perimeter.
+
+    Returns
+    -------
+    pt_sign : ndarray of {-1,1}
+        Integers corresponding to if the given point is inside or outside of the given
+        perimeter, with pt_sign[i]>0 meaning the point given by Rcheck[i], Zcheck[i] is
+        inside of the given perimeter, and a negative sign meaning the point is outside
+        of the given perimeter.
+
+    """
+    # R Z are the perimeter points
+    # Rcheck Zcheck are the points being checked for whether
+    # or not they are inside the check
+
+    Rbool = R[:, None] > Rcheck
+    Zbool = Z[:, None] > Zcheck
+    # these are now size (Ncheck, Nperimeter)
+    quadrants = jnp.zeros_like(Rbool)
+    quadrants = jnp.where(jnp.logical_and(jnp.logical_not(Rbool), Zbool), 1, quadrants)
+    quadrants = jnp.where(
+        jnp.logical_and(jnp.logical_not(Rbool), jnp.logical_not(Zbool)),
+        2,
+        quadrants,
+    )
+    quadrants = jnp.where(jnp.logical_and(Rbool, jnp.logical_not(Zbool)), 3, quadrants)
+    deltas = quadrants[1:, :] - quadrants[0:-1, :]
+    deltas = jnp.where(deltas == 3, -1, deltas)
+    deltas = jnp.where(deltas == -3, 1, deltas)
+    # then flip sign if the R intercept is > Rcheck and the
+    # quadrant flipped over a diagonal
+    b = (Z[1:] / R[1:] - Z[0:-1] / R[0:-1]) / (Z[1:] - Z[0:-1])
+    Rint = Rcheck[:, None] - b * (R[1:] - R[0:-1]) / (Z[1:] - Z[0:-1])
+    deltas = jnp.where(
+        jnp.logical_and(jnp.abs(deltas) == 2, Rint.T > Rcheck),
+        -deltas,
+        deltas,
+    )
+    pt_sign = jnp.sum(deltas, axis=0)
+    # positive distance if the check pt is inside the perimeter, else
+    # negative distance is assigned
+    # pt_sign = 0 : Means point is OUTSIDE of the perimeter,
+    #               assign positive distance
+    # pt_sign = +/-4: Means point is INSIDE perimeter, so
+    #                 assign negative distance
+    pt_sign = jnp.where(jnp.isclose(pt_sign, 0), 1, -1)
+    return pt_sign