diff --git a/CHANGELOG.md b/CHANGELOG.md index 705e6c7e10..aab80a4173 100644 --- a/CHANGELOG.md +++ b/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 ------- diff --git a/desc/compute/__init__.py b/desc/compute/__init__.py index 416825cc55..87f03068bb 100644 --- a/desc/compute/__init__.py +++ b/desc/compute/__init__.py @@ -94,3 +94,30 @@ def _build_data_index(): _build_data_index() + + +def set_tier(name, p): + """Determine how deep in the dependency tree a given name is. + + tier of 0 means no dependencies on other data, + tier of 1 means it depends on only tier 0 stuff, + tier of 2 means it depends on tier 0 and tier 1, etc etc. + + Designed such that if you compute things in the order determined by tiers, + all dependencies will always be computed in the correct order. + """ + if "tier" in data_index[p][name]: + return + if len(data_index[p][name]["full_with_axis_dependencies"]["data"]) == 0: + data_index[p][name]["tier"] = 0 + else: + thistier = 0 + for name1 in data_index[p][name]["full_with_axis_dependencies"]["data"]: + set_tier(name1, p) + thistier = max(thistier, data_index[p][name1]["tier"]) + data_index[p][name]["tier"] = thistier + 1 + + +for par in data_index.keys(): + for name in data_index[par]: + set_tier(name, par) diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py index 70bc1fa6e5..8fca1346d2 100644 --- a/desc/compute/_basis_vectors.py +++ b/desc/compute/_basis_vectors.py @@ -1431,7 +1431,6 @@ def _e_sub_rho(params, transforms, profiles, data, **kwargs): # At the magnetic axis, this function returns the multivalued map whose # image is the set { 𝐞ᵨ | ρ=0 }. data["e_rho"] = jnp.array([data["R_r"], data["R"] * data["omega_r"], data["Z_r"]]).T - return data diff --git a/desc/compute/_core.py b/desc/compute/_core.py index e0e8312b9b..7bc6e8acb3 100644 --- a/desc/compute/_core.py +++ b/desc/compute/_core.py @@ -22,12 +22,37 @@ "desc.geometry.core.Surface", "desc.geometry.core.Curve", ], + aliases=["rho_t", "rho_z", "theta_r", "theta_z", "zeta_r", "zeta_t"], ) def _0(params, transforms, profiles, data, **kwargs): data["0"] = jnp.zeros(transforms["grid"].num_nodes) return data +@register_compute_fun( + name="1", + label="1", + units="~", + units_long="None", + description="Ones", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="rtz", + data=[], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + "desc.geometry.core.Curve", + ], + aliases=["rho_r", "theta_t", "zeta_z"], +) +def _1(params, transforms, profiles, data, **kwargs): + data["1"] = jnp.ones(transforms["grid"].num_nodes) + return data + + @register_compute_fun( name="R", label="R", @@ -2711,6 +2736,10 @@ def _phi(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["omega_r"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _phi_r(params, transforms, profiles, data, **kwargs): data["phi_r"] = data["omega_r"] @@ -2785,6 +2814,10 @@ def _phi_rz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["omega_t"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _phi_t(params, transforms, profiles, data, **kwargs): data["phi_t"] = data["omega_t"] @@ -2841,6 +2874,10 @@ def _phi_tz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["omega_z"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _phi_z(params, transforms, profiles, data, **kwargs): data["phi_z"] = 1 + data["omega_z"] @@ -2890,75 +2927,6 @@ def _rho(params, transforms, profiles, data, **kwargs): return data -@register_compute_fun( - name="rho_r", - label="\\partial_{\\rho} \\rho", - units="~", - units_long="None", - description="Radial coordinate, proportional to the square root " - + "of the toroidal flux, derivative wrt radial coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="r", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], -) -def _rho_r(params, transforms, profiles, data, **kwargs): - data["rho_r"] = jnp.ones_like(data["0"]) - return data - - -@register_compute_fun( - name="rho_t", - label="\\partial_{\\theta} \\rho", - units="~", - units_long="None", - description="Radial coordinate, proportional to the square root " - "of the toroidal flux, derivative wrt poloidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="r", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], -) -def _rho_t(params, transforms, profiles, data, **kwargs): - data["rho_t"] = data["0"] - return data - - -@register_compute_fun( - name="rho_z", - label="\\partial_{\\zeta} \\rho", - units="~", - units_long="None", - description="Radial coordinate, proportional to the square root " - "of the toroidal flux, derivative wrt toroidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="r", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], -) -def _rho_z(params, transforms, profiles, data, **kwargs): - data["rho_z"] = data["0"] - return data - - @register_compute_fun( name="theta", label="\\theta", @@ -3113,75 +3081,6 @@ def _theta_PEST_zz(params, transforms, profiles, data, **kwargs): return data -@register_compute_fun( - name="theta_r", - label="\\partial_{\\rho} \\theta", - units="rad", - units_long="radians", - description="Poloidal angular coordinate (geometric, not magnetic), " - "derivative wrt radial coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="t", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], -) -def _theta_r(params, transforms, profiles, data, **kwargs): - data["theta_r"] = data["0"] - return data - - -@register_compute_fun( - name="theta_t", - label="\\partial_{\\theta} \\theta", - units="rad", - units_long="radians", - description="Poloidal angular coordinate (geometric, not magnetic), " - "derivative wrt poloidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="t", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], -) -def _theta_t(params, transforms, profiles, data, **kwargs): - data["theta_t"] = jnp.ones_like(data["0"]) - return data - - -@register_compute_fun( - name="theta_z", - label="\\partial_{\\zeta} \\theta", - units="rad", - units_long="radians", - description="Poloidal angular coordinate (geometric, not magnetic), " - "derivative wrt toroidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="t", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], -) -def _theta_z(params, transforms, profiles, data, **kwargs): - data["theta_z"] = data["0"] - return data - - @register_compute_fun( name="zeta", label="\\zeta", @@ -3202,69 +3101,3 @@ def _theta_z(params, transforms, profiles, data, **kwargs): def _zeta(params, transforms, profiles, data, **kwargs): data["zeta"] = transforms["grid"].nodes[:, 2] return data - - -@register_compute_fun( - name="zeta_r", - label="\\partial_{\\rho} \\zeta", - units="rad", - units_long="radians", - description="Toroidal angular coordinate derivative, wrt radial coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="z", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], -) -def _zeta_r(params, transforms, profiles, data, **kwargs): - data["zeta_r"] = data["0"] - return data - - -@register_compute_fun( - name="zeta_t", - label="\\partial_{\\theta} \\zeta", - units="rad", - units_long="radians", - description="Toroidal angular coordinate, derivative wrt poloidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="z", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], -) -def _zeta_t(params, transforms, profiles, data, **kwargs): - data["zeta_t"] = data["0"] - return data - - -@register_compute_fun( - name="zeta_z", - label="\\partial_{\\zeta} \\zeta", - units="rad", - units_long="radians", - description="Toroidal angular coordinate, derivative wrt toroidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="z", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], -) -def _zeta_z(params, transforms, profiles, data, **kwargs): - data["zeta_z"] = jnp.ones_like(data["0"]) - return data diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index 818da18092..2e96e7a767 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -31,7 +31,10 @@ def _s(params, transforms, profiles, data, **kwargs): label="ds", units="~", units_long="None", - description="Spacing of curve parameter", + description=( + "Quadrature weights for integration along the curve," + + " i.e. an alias for ``grid.spacing[:,2]``" + ), dim=1, params=[], transforms={"grid": []}, diff --git a/desc/compute/_field.py b/desc/compute/_field.py index eef1354e3d..e53f033464 100644 --- a/desc/compute/_field.py +++ b/desc/compute/_field.py @@ -2935,7 +2935,7 @@ def _gradB2(params, transforms, profiles, data, **kwargs): @register_compute_fun( name="|grad(|B|^2)|/2mu0", - label="|\\nabla |B|^{2}/(2\\mu_0)|", + label="|\\nabla |B|^{2}|/(2\\mu_0)", units="N \\cdot m^{-3}", units_long="Newton / cubic meter", description="Magnitude of magnetic pressure gradient", diff --git a/desc/compute/_profiles.py b/desc/compute/_profiles.py index ccb701f072..1cfbc2f23f 100644 --- a/desc/compute/_profiles.py +++ b/desc/compute/_profiles.py @@ -82,10 +82,10 @@ def _psi_r(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="r", - data=["rho"], + data=["1"], ) def _psi_rr(params, transforms, profiles, data, **kwargs): - data["psi_rr"] = params["Psi"] * jnp.ones_like(data["rho"]) / jnp.pi + data["psi_rr"] = data["1"] * params["Psi"] / jnp.pi return data diff --git a/desc/compute/_surface.py b/desc/compute/_surface.py index 5fb84fa060..d37f82337e 100644 --- a/desc/compute/_surface.py +++ b/desc/compute/_surface.py @@ -118,66 +118,6 @@ def _phi_Surface(params, transforms, profiles, data, **kwargs): return data -@register_compute_fun( - name="phi_r", - label="\\partial_{\\rho} \\phi", - units="rad", - units_long="radians", - description="Toroidal angle in lab frame, derivative wrt radial coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["e_rho"], - parameterization="desc.geometry.core.Surface", -) -def _phi_r_Surface(params, transforms, profiles, data, **kwargs): - coords_r = data["e_rho"] - data["phi_r"] = coords_r[:, 1] - return data - - -@register_compute_fun( - name="phi_t", - label="\\partial_{\\theta} \\phi", - units="rad", - units_long="radians", - description="Toroidal angle in lab frame, derivative wrt poloidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["e_theta"], - parameterization="desc.geometry.core.Surface", -) -def _phi_t_Surface(params, transforms, profiles, data, **kwargs): - coords_t = data["e_theta"] - data["phi_t"] = coords_t[:, 1] - return data - - -@register_compute_fun( - name="phi_z", - label="\\partial_{\\zeta} \\phi", - units="rad", - units_long="radians", - description="Toroidal angle in lab frame, derivative wrt toroidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["e_zeta"], - parameterization="desc.geometry.core.Surface", -) -def _phi_z_Surface(params, transforms, profiles, data, **kwargs): - coords_z = data["e_zeta"] - data["phi_z"] = coords_z[:, 1] - return data - - @register_compute_fun( name="Z", label="Z", diff --git a/desc/compute/data_index.py b/desc/compute/data_index.py index de4ecf9314..26341ec587 100644 --- a/desc/compute/data_index.py +++ b/desc/compute/data_index.py @@ -132,6 +132,8 @@ def register_compute_fun( # noqa: C901 source_grid_requirement = {} if not isinstance(parameterization, (tuple, list)): parameterization = [parameterization] + if not isinstance(aliases, (tuple, list)): + aliases = [aliases] deps = { "params": params, diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index 7ae828d95e..f72e9f7605 100644 --- a/desc/objectives/_geometry.py +++ b/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,9 +607,12 @@ 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) @@ -611,10 +620,29 @@ def __init__( 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): diff --git a/desc/objectives/utils.py b/desc/objectives/utils.py index 924e8cfb43..27a8d30b37 100644 --- a/desc/objectives/utils.py +++ b/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 diff --git a/desc/optimize/tr_subproblems.py b/desc/optimize/tr_subproblems.py index c9149d2bb3..b107bca0a4 100644 --- a/desc/optimize/tr_subproblems.py +++ b/desc/optimize/tr_subproblems.py @@ -421,7 +421,7 @@ def trust_region_step_exact_qr( p_newton = solve_triangular_regularized(R, -Q.T @ f) else: Q, R = qr(J.T, mode="economic") - p_newton = Q @ solve_triangular_regularized(R.T, f, lower=True) + p_newton = Q @ solve_triangular_regularized(R.T, -f, lower=True) def truefun(*_): return p_newton, False, 0.0 diff --git a/desc/plotting.py b/desc/plotting.py index ba822a5b2b..fdcb9a393c 100644 --- a/desc/plotting.py +++ b/desc/plotting.py @@ -16,7 +16,7 @@ from desc.backend import sign from desc.basis import fourier, zernike_radial_poly -from desc.coils import CoilSet +from desc.coils import CoilSet, _Coil from desc.compute import data_index, get_transforms from desc.compute.utils import _parse_parameterization, surface_averages_map from desc.equilibrium.coords import map_coordinates @@ -2394,6 +2394,11 @@ def plot_coils(coils, grid=None, fig=None, return_data=False, **kwargs): ValueError, f"plot_coils got unexpected keyword argument: {kwargs.keys()}", ) + errorif( + not isinstance(coils, _Coil), + ValueError, + "Expected `coils` to be of type `_Coil`, instead got type" f" {type(coils)}", + ) if not isinstance(lw, (list, tuple)): lw = [lw] diff --git a/tests/inputs/master_compute_data_rpz.pkl b/tests/inputs/master_compute_data_rpz.pkl index 25ba750819..8a9086d237 100644 Binary files a/tests/inputs/master_compute_data_rpz.pkl and b/tests/inputs/master_compute_data_rpz.pkl differ diff --git a/tests/test_equilibrium.py b/tests/test_equilibrium.py index 834775ecae..52958a6a9b 100644 --- a/tests/test_equilibrium.py +++ b/tests/test_equilibrium.py @@ -15,6 +15,7 @@ from desc.io import InputReader from desc.objectives import ForceBalance, ObjectiveFunction, get_equilibrium_objective from desc.profiles import PowerSeriesProfile +from desc.utils import errorif from .utils import area_difference, compute_coords @@ -55,16 +56,27 @@ def test_map_coordinates(): iota = grid.compress(eq.compute("iota", grid=grid)["iota"]) iota = np.broadcast_to(iota, shape=(n,)) - out_1 = eq.map_coordinates(coords, inbasis=["rho", "alpha", "zeta"], iota=iota) + tol = 1e-5 + out_1 = eq.map_coordinates( + coords, inbasis=["rho", "alpha", "zeta"], iota=iota, tol=tol + ) assert np.isfinite(out_1).all() out_2 = eq.map_coordinates( coords, inbasis=["rho", "alpha", "zeta"], period=(np.inf, 2 * np.pi, np.inf), + tol=tol, ) assert np.isfinite(out_2).all() - diff = (out_1 - out_2) % (2 * np.pi) - assert np.all(np.isclose(diff, 0) | np.isclose(np.abs(diff), 2 * np.pi)) + diff = out_1 - out_2 + errorif( + not np.all( + np.isclose(diff, 0, atol=2 * tol) + | np.isclose(np.abs(diff), 2 * np.pi, atol=2 * tol) + ), + AssertionError, + f"diff: {diff}", + ) eq = get("DSHAPE") diff --git a/tests/test_examples.py b/tests/test_examples.py index a8bff8a538..6bae9dd16b 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -440,7 +440,6 @@ def test_NAE_QSC_solve(): orig_Rax_val = eq.axis.R_n orig_Zax_val = eq.axis.Z_n - eq_fit = eq.copy() eq_lambda_fixed_0th_order = eq.copy() eq_lambda_fixed_1st_order = eq.copy() @@ -472,7 +471,6 @@ def test_NAE_QSC_solve(): xtol=1e-6, constraints=constraints, ) - grid = LinearGrid(L=10, M=20, N=20, NFP=eq.NFP, sym=True, axis=False) grid_axis = LinearGrid(rho=0.0, theta=0.0, N=eq.N_grid, NFP=eq.NFP) # Make sure axis is same for eqq, string in zip( @@ -482,24 +480,15 @@ def test_NAE_QSC_solve(): np.testing.assert_array_almost_equal(orig_Rax_val, eqq.axis.R_n, err_msg=string) np.testing.assert_array_almost_equal(orig_Zax_val, eqq.axis.Z_n, err_msg=string) - # Make sure surfaces of solved equilibrium are similar near axis as QSC - rho_err, theta_err = area_difference_desc(eqq, eq_fit) + # Make sure iota of solved equilibrium is same on-axis as QSC - np.testing.assert_allclose(rho_err[:, 0:-6], 0, atol=1.5e-2, err_msg=string) - np.testing.assert_allclose(theta_err[:, 0:-6], 0, atol=1e-3, err_msg=string) - - # Make sure iota of solved equilibrium is same near axis as QSC - - iota = grid.compress(eqq.compute("iota", grid=grid)["iota"]) + iota = eqq.compute("iota", grid=LinearGrid(rho=0.0))["iota"] np.testing.assert_allclose(iota[0], qsc.iota, atol=1e-5, err_msg=string) - np.testing.assert_allclose(iota[1:10], qsc.iota, atol=1e-3, err_msg=string) - # check lambda to match near axis - # Evaluate lambda near the axis + # check lambda to match on-axis data_nae = eqq.compute(["lambda", "|B|"], grid=grid_axis) lam_nae = data_nae["lambda"] - # Reshape to form grids on theta and phi phi = np.squeeze(grid_axis.nodes[:, 2]) np.testing.assert_allclose( @@ -526,7 +515,6 @@ def test_NAE_QSC_solve_near_axis_based_off_eq(): orig_Rax_val = eq.axis.R_n orig_Zax_val = eq.axis.Z_n - eq_fit = eq.copy() eq_lambda_fixed_0th_order = eq.copy() eq_lambda_fixed_1st_order = eq.copy() @@ -558,7 +546,6 @@ def test_NAE_QSC_solve_near_axis_based_off_eq(): xtol=1e-6, constraints=constraints, ) - grid = LinearGrid(L=10, M=20, N=20, NFP=eq.NFP, sym=True, axis=False) grid_axis = LinearGrid(rho=0.0, theta=0.0, N=eq.N_grid, NFP=eq.NFP) # Make sure axis is same for eqq, string in zip( @@ -568,18 +555,11 @@ def test_NAE_QSC_solve_near_axis_based_off_eq(): np.testing.assert_array_almost_equal(orig_Rax_val, eqq.axis.R_n, err_msg=string) np.testing.assert_array_almost_equal(orig_Zax_val, eqq.axis.Z_n, err_msg=string) - # Make sure surfaces of solved equilibrium are similar near axis as QSC - rho_err, theta_err = area_difference_desc(eqq, eq_fit) - - np.testing.assert_allclose(rho_err[:, 0:-6], 0, atol=1.5e-2, err_msg=string) - np.testing.assert_allclose(theta_err[:, 0:-6], 0, atol=1e-3, err_msg=string) + # Make sure iota of solved equilibrium is same on axis as QSC - # Make sure iota of solved equilibrium is same near axis as QSC - - iota = grid.compress(eqq.compute("iota", grid=grid)["iota"]) + iota = eqq.compute("iota", grid=LinearGrid(rho=0.0))["iota"] np.testing.assert_allclose(iota[0], qsc.iota, atol=1e-5, err_msg=string) - np.testing.assert_allclose(iota[1:10], qsc.iota, rtol=5e-3, err_msg=string) ### check lambda to match on axis # Evaluate lambda on the axis @@ -603,18 +583,16 @@ def test_NAE_QSC_solve_near_axis_based_off_eq(): def test_NAE_QIC_solve(): """Test O(rho) NAE QIC constraints solve.""" # get Qic example - qic = Qic.from_paper("QI NFP2 r2", nphi=301, order="r1") + qic = Qic.from_paper("QI NFP2 r2", nphi=199, order="r1") qic.lasym = False # don't need to consider stellarator asym for order 1 constraints ntheta = 75 r = 0.01 - N = 11 + N = 9 eq = Equilibrium.from_near_axis(qic, r=r, L=7, M=7, N=N, ntheta=ntheta) orig_Rax_val = eq.axis.R_n orig_Zax_val = eq.axis.Z_n - eq_fit = eq.copy() - # this has all the constraints we need, cs = get_NAE_constraints(eq, qic, order=1) @@ -629,75 +607,24 @@ def test_NAE_QIC_solve(): np.testing.assert_array_almost_equal(orig_Rax_val, eq.axis.R_n) np.testing.assert_array_almost_equal(orig_Zax_val, eq.axis.Z_n) - # Make sure surfaces of solved equilibrium are similar near axis as QIC - rho_err, theta_err = area_difference_desc(eq, eq_fit) - - np.testing.assert_allclose(rho_err[:, 0:3], 0, atol=5e-2) - # theta error isn't really an indicator of near axis behavior - # since it's computed over the full radius, but just indicates that - # eq is similar to eq_fit - np.testing.assert_allclose(theta_err, 0, atol=5e-2) - # Make sure iota of solved equilibrium is same near axis as QIC - grid = LinearGrid(L=10, M=20, N=20, NFP=eq.NFP, sym=True, axis=False) - iota = grid.compress(eq.compute("iota", grid=grid)["iota"]) + iota = eq.compute("iota", grid=LinearGrid(rho=0.0))["iota"] np.testing.assert_allclose(iota[0], qic.iota, rtol=1e-5) - np.testing.assert_allclose(iota[1:10], qic.iota, rtol=1e-3) - - # check lambda to match near axis - grid_2d_05 = LinearGrid(rho=np.array(1e-6), M=50, N=50, NFP=eq.NFP, endpoint=True) - # Evaluate lambda near the axis - data_nae = eq.compute("lambda", grid=grid_2d_05) - lam_nae = data_nae["lambda"] + grid_axis = LinearGrid(rho=0.0, theta=0.0, zeta=qic.phi, NFP=eq.NFP) + phi = grid_axis.nodes[:, 2].squeeze() - # Reshape to form grids on theta and phi - zeta = ( - grid_2d_05.nodes[:, 2] - .reshape( - (grid_2d_05.num_theta, grid_2d_05.num_rho, grid_2d_05.num_zeta), order="F" - ) - .squeeze() - ) - - lam_nae = lam_nae.reshape( - (grid_2d_05.num_theta, grid_2d_05.num_rho, grid_2d_05.num_zeta), order="F" - ) - - phi = np.squeeze(zeta[0, :]) - lam_nae = np.squeeze(lam_nae[:, 0, :]) + # check lambda to match on-axis + lam_nae = eq.compute("lambda", grid=grid_axis)["lambda"] - lam_av_nae = np.mean(lam_nae, axis=0) np.testing.assert_allclose( - lam_av_nae, -qic.iota * qic.nu_spline(phi), atol=1e-4, rtol=1e-2 + lam_nae, -qic.iota * qic.nu_spline(phi), atol=1e-4, rtol=1e-2 ) # check |B| on axis - - grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, rho=np.array(1e-6)) - # Evaluate B modes near the axis - data_nae = eq.compute(["|B|_mn", "B modes"], grid=grid) - modes = data_nae["B modes"] - B_mn_nae = data_nae["|B|_mn"] - # Evaluate B on an angular grid - theta = np.linspace(0, 2 * np.pi, 150) - phi = np.linspace(0, 2 * np.pi, qic.nphi) - th, ph = np.meshgrid(theta, phi) - B_nae = np.zeros((qic.nphi, 150)) - - for i, (l, m, n) in enumerate(modes): - if m >= 0 and n >= 0: - B_nae += B_mn_nae[i] * np.cos(m * th) * np.cos(n * ph) - elif m >= 0 > n: - B_nae += -B_mn_nae[i] * np.cos(m * th) * np.sin(n * ph) - elif m < 0 <= n: - B_nae += -B_mn_nae[i] * np.sin(m * th) * np.cos(n * ph) - elif m < 0 and n < 0: - B_nae += B_mn_nae[i] * np.sin(m * th) * np.sin(n * ph) - # Eliminate the poloidal angle to focus on the toroidal behavior - B_av_nae = np.mean(B_nae, axis=1) - np.testing.assert_allclose(B_av_nae, np.ones(np.size(phi)) * qic.B0, atol=2e-2) + B_nae = eq.compute(["|B|"], grid=grid_axis)["|B|"] + np.testing.assert_allclose(B_nae, qic.B0, atol=1e-3) @pytest.mark.unit @@ -1058,7 +985,7 @@ def test_non_eq_optimization(): use_softmin=True, surface_grid=grid, plasma_grid=grid, - alpha=5000, + softmin_alpha=5000, ) objective = ObjectiveFunction((obj,)) optimizer = Optimizer("lsq-auglag") @@ -1630,3 +1557,102 @@ def circle_constraint(params): abs(surf_opt.Z_lmn[surf_opt.Z_basis.get_idx(M=-1, N=0)]) + offset, rtol=2e-2, ) + + +@pytest.mark.slow +@pytest.mark.regression +@pytest.mark.optimize +def test_signed_PlasmaVesselDistance(): + """Tests that signed distance works with surface optimization.""" + eq = get("HELIOTRON") + eq.change_resolution(M=2, N=2) + + surf = eq.surface.copy() + surf.change_resolution(M=1, N=1) + + target_dist = -0.25 + + grid = LinearGrid(M=10, N=4, NFP=eq.NFP) + obj = PlasmaVesselDistance( + surface=surf, + eq=eq, + target=target_dist, + surface_grid=grid, + plasma_grid=grid, + use_signed_distance=True, + eq_fixed=True, + ) + objective = ObjectiveFunction((obj,)) + + optimizer = Optimizer("lsq-exact") + (surf,), _ = optimizer.optimize( + (surf,), objective, verbose=3, maxiter=60, ftol=1e-8, xtol=1e-9 + ) + + np.testing.assert_allclose( + obj.compute(*obj.xs(surf)), target_dist, atol=1e-2, err_msg="Using hardmin" + ) + + # with softmin + surf = eq.surface.copy() + surf.change_resolution(M=1, N=1) + obj = PlasmaVesselDistance( + surface=surf, + eq=eq, + target=target_dist, + surface_grid=grid, + plasma_grid=grid, + use_signed_distance=True, + use_softmin=True, + softmin_alpha=100, + eq_fixed=True, + ) + objective = ObjectiveFunction((obj,)) + + optimizer = Optimizer("lsq-exact") + (surf,), _ = optimizer.optimize( + (surf,), + objective, + verbose=3, + maxiter=60, + ftol=1e-8, + xtol=1e-9, + ) + + np.testing.assert_allclose( + obj.compute(*obj.xs(surf)), target_dist, atol=1e-2, err_msg="Using softmin" + ) + + # with changing eq + eq = Equilibrium(M=1, N=1) + surf = eq.surface.copy() + surf.change_resolution(M=1, N=1) + grid = LinearGrid(M=20, N=8, NFP=eq.NFP) + + obj = PlasmaVesselDistance( + surface=surf, + eq=eq, + target=target_dist, + surface_grid=grid, + plasma_grid=grid, + use_signed_distance=True, + ) + objective = ObjectiveFunction(obj) + + optimizer = Optimizer("lsq-exact") + (eq, surf), _ = optimizer.optimize( + (eq, surf), + objective, + constraints=(FixParameters(surf),), + verbose=3, + maxiter=60, + ftol=1e-8, + xtol=1e-9, + ) + + np.testing.assert_allclose( + obj.compute(*obj.xs(eq, surf)), + target_dist, + atol=1e-2, + err_msg="allowing eq to change", + ) diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 9fcc12634b..d6f67a5ffc 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -640,6 +640,7 @@ def test_plasma_vessel_distance(self): ) obj.build() d = obj.compute_unscaled(*obj.xs(eq, surface)) + assert d.size == obj.dim_f assert abs(d.min() - (a_s - a_p)) < 1e-14 assert abs(d.max() - (a_s - a_p)) < surf_grid.spacing[0, 1] * a_p @@ -679,6 +680,7 @@ def test_plasma_vessel_distance(self): ) obj.build() d = obj.compute_unscaled(*obj.xs(eq, surface)) + assert d.size == obj.dim_f assert np.all(np.abs(d) < a_s - a_p) # for large enough alpha, should be same as actual min @@ -688,7 +690,7 @@ def test_plasma_vessel_distance(self): surface_grid=surf_grid, surface=surface, use_softmin=True, - alpha=100, + softmin_alpha=100, ) obj.build() d = obj.compute_unscaled(*obj.xs(eq, surface)) @@ -1137,6 +1139,89 @@ def test(eq, field, correct_value, rtol=1e-14, grid=None): test(eq, field, psi_from_field) + @pytest.mark.unit + def test_signed_plasma_vessel_distance(self): + """Test calculation of signed distance from plasma to vessel.""" + R0 = 10.0 + a_p = 1.0 + a_s = 2.0 + # default eq has R0=10, a=1 + eq = Equilibrium(M=3, N=2) + # surface with same R0, a=2, so true d=1 for all pts + surface = FourierRZToroidalSurface( + R_lmn=[R0, a_s], Z_lmn=[-a_s], modes_R=[[0, 0], [1, 0]], modes_Z=[[-1, 0]] + ) + grid = LinearGrid(M=5, N=6) + obj = PlasmaVesselDistance( + eq=eq, + surface_grid=grid, + plasma_grid=grid, + surface=surface, + use_signed_distance=True, + ) + obj.build() + d = obj.compute_unscaled(*obj.xs(eq, surface)) + assert obj.dim_f == d.size + np.testing.assert_allclose(d, a_s - a_p) + + # ensure that it works (dimension-wise) when compute_scaled is called + _ = obj.compute_scaled(*obj.xs(eq, surface)) + + # For plasma outside surface, should get signed distance + a_s = 0.5 * a_p + surface = FourierRZToroidalSurface( + R_lmn=[R0, a_s], + Z_lmn=[-a_s], + modes_R=[[0, 0], [1, 0]], + modes_Z=[[-1, 0]], + ) + grid = LinearGrid(M=5, N=6) + obj = PlasmaVesselDistance( + eq=eq, + surface_grid=grid, + plasma_grid=grid, + surface=surface, + use_signed_distance=True, + ) + obj.build() + d = obj.compute_unscaled(*obj.xs(eq, surface)) + assert obj.dim_f == d.size + np.testing.assert_allclose(d, a_s - a_p) + + # ensure it works with different sized grids (poloidal resolution different) + grid = LinearGrid(M=5, N=6) + obj = PlasmaVesselDistance( + eq=eq, + surface_grid=grid, + plasma_grid=LinearGrid(M=10, N=6), + surface=surface, + use_signed_distance=True, + ) + obj.build() + d = obj.compute_unscaled(*obj.xs(eq, surface)) + assert obj.dim_f == d.size + assert abs(d.max() - (-a_s)) < 1e-14 + assert abs(d.min() - (-a_s)) < grid.spacing[0, 1] * a_s + + # ensure it works with different sized grids (poloidal resolution different) + # and using softmin (with deprecated name alpha) + grid = LinearGrid(M=5, N=6) + with pytest.raises(FutureWarning): + obj = PlasmaVesselDistance( + eq=eq, + surface_grid=grid, + plasma_grid=LinearGrid(M=10, N=6), + surface=surface, + use_signed_distance=True, + use_softmin=True, + alpha=4000, + ) + obj.build() + d = obj.compute_unscaled(*obj.xs(eq, surface)) + assert obj.dim_f == d.size + assert abs(d.max() - (-a_s)) < 1e-14 + assert abs(d.min() - (-a_s)) < grid.spacing[0, 1] * a_s + @pytest.mark.regression def test_derivative_modes(): diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 8252d349d8..1351b46c93 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -801,6 +801,8 @@ def test_plot_coils(): coil.rotate(angle=np.pi / N) coils = CoilSet.linspaced_angular(coil, I, [0, 0, 1], np.pi / NFP, N // NFP // 2) coils2 = MixedCoilSet.from_symmetry(coils, NFP, True) + with pytest.raises(ValueError, match="Expected `coils`"): + plot_coils("not a coil") fig, data = plot_coils(coils2, return_data=True) def flatten_coils(coilset):