diff --git a/desc/coils.py b/desc/coils.py index be954c115..ea0da348a 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -1,5 +1,6 @@ """Classes for magnetic field coils.""" +import functools import numbers from abc import ABC from collections.abc import MutableSequence @@ -28,7 +29,7 @@ from desc.grid import LinearGrid from desc.magnetic_fields import _MagneticField from desc.optimizable import Optimizable, OptimizableCollection, optimizable_parameter -from desc.utils import equals, errorif, flatten_list, safenorm, warnif +from desc.utils import cross, dot, equals, errorif, flatten_list, safenorm, warnif @jit @@ -245,7 +246,7 @@ def num_coils(self): """int: Number of coils.""" return 1 - def _compute_position(self, params=None, grid=None, **kwargs): + def _compute_position(self, params=None, grid=None, dx1=False, **kwargs): """Compute coil positions accounting for stellarator symmetry. Parameters @@ -255,18 +256,31 @@ def _compute_position(self, params=None, grid=None, **kwargs): grid : Grid or int, optional Grid of coordinates to evaluate at. Defaults to a Linear grid. If an integer, uses that many equally spaced points. + dx1 : bool + If True, also return dx/ds for the curve. Returns ------- x : ndarray, shape(len(self),source_grid.num_nodes,3) Coil positions, in [R,phi,Z] or [X,Y,Z] coordinates. + x_s : ndarray, shape(len(self),source_grid.num_nodes,3) + Coil position derivatives, in [R,phi,Z] or [X,Y,Z] coordinates. + Only returned if dx1=True. """ - x = self.compute("x", grid=grid, params=params, **kwargs)["x"] - x = jnp.transpose(jnp.atleast_3d(x), [2, 0, 1]) # shape=(1,num_nodes,3) - basis = kwargs.pop("basis", "xyz") + kwargs.setdefault("basis", "xyz") + keys = ["x", "x_s"] if dx1 else ["x"] + data = self.compute(keys, grid=grid, params=params, **kwargs) + x = jnp.transpose(jnp.atleast_3d(data["x"]), [2, 0, 1]) # shape=(1,num_nodes,3) + if dx1: + x_s = jnp.transpose( + jnp.atleast_3d(data["x_s"]), [2, 0, 1] + ) # shape=(1,num_nodes,3) + basis = kwargs.get("basis", "xyz") if basis.lower() == "rpz": x = x.at[:, :, 1].set(jnp.mod(x[:, :, 1], 2 * jnp.pi)) + if dx1: + return x, x_s return x def _compute_A_or_B( @@ -1359,7 +1373,7 @@ def flip(self, *args, **kwargs): """Flip the coils across a plane.""" [coil.flip(*args, **kwargs) for coil in self.coils] - def _compute_position(self, params=None, grid=None, **kwargs): + def _compute_position(self, params=None, grid=None, dx1=False, **kwargs): """Compute coil positions accounting for stellarator symmetry. Parameters @@ -1369,25 +1383,35 @@ def _compute_position(self, params=None, grid=None, **kwargs): grid : Grid or int, optional Grid of coordinates to evaluate at. Defaults to a Linear grid. If an integer, uses that many equally spaced points. + dx1 : bool + If True, also return dx/ds for each curve. Returns ------- x : ndarray, shape(len(self),source_grid.num_nodes,3) Coil positions, in [R,phi,Z] or [X,Y,Z] coordinates. + x_s : ndarray, shape(len(self),source_grid.num_nodes,3) + Coil position derivatives, in [R,phi,Z] or [X,Y,Z] coordinates. + Only returned if dx1=True. """ basis = kwargs.pop("basis", "xyz") + keys = ["x", "x_s"] if dx1 else ["x"] if params is None: - params = [get_params("x", coil, basis=basis) for coil in self] - data = self.compute("x", grid=grid, params=params, basis=basis, **kwargs) + params = [get_params(keys, coil, basis=basis) for coil in self] + data = self.compute(keys, grid=grid, params=params, basis=basis, **kwargs) data = tree_leaves(data, is_leaf=lambda x: isinstance(x, dict)) x = jnp.dstack([d["x"].T for d in data]).T # shape=(ncoils,num_nodes,3) - + if dx1: + x_s = jnp.dstack([d["x_s"].T for d in data]).T # shape=(ncoils,num_nodes,3) # stellarator symmetry is easiest in [X,Y,Z] coordinates - if basis.lower() == "rpz": - xyz = rpz2xyz(x) - else: - xyz = x + xyz = rpz2xyz(x) if basis.lower() == "rpz" else x + if dx1: + xyz_s = ( + rpz2xyz_vec(x_s, xyz[:, :, 0], xyz[:, :, 1]) + if basis.lower() == "rpz" + else x_s + ) # if stellarator symmetric, add reflected coils from the other half field period if self.sym: @@ -1396,27 +1420,64 @@ def _compute_position(self, params=None, grid=None, **kwargs): ) xyz_sym = xyz @ reflection_matrix(normal).T @ reflection_matrix([0, 0, 1]).T xyz = jnp.vstack((xyz, jnp.flipud(xyz_sym))) + if dx1: + xyz_s_sym = ( + xyz_s @ reflection_matrix(normal).T @ reflection_matrix([0, 0, 1]).T + ) + xyz_s = jnp.vstack((xyz_s, jnp.flipud(xyz_s_sym))) # field period rotation is easiest in [R,phi,Z] coordinates rpz = xyz2rpz(xyz) + if dx1: + rpz_s = xyz2rpz_vec(xyz_s, xyz[:, :, 0], xyz[:, :, 1]) # if field period symmetry, add rotated coils from other field periods - if self.NFP > 1: - rpz0 = rpz - for k in range(1, self.NFP): - rpz = jnp.vstack( - (rpz, rpz0 + jnp.array([0, 2 * jnp.pi * k / self.NFP, 0])) - ) + rpz0 = rpz + for k in range(1, self.NFP): + rpz = jnp.vstack((rpz, rpz0 + jnp.array([0, 2 * jnp.pi * k / self.NFP, 0]))) + if dx1: + rpz_s = jnp.tile(rpz_s, (self.NFP, 1, 1)) # ensure phi in [0, 2pi) rpz = rpz.at[:, :, 1].set(jnp.mod(rpz[:, :, 1], 2 * jnp.pi)) - if basis.lower() == "xyz": - x = rpz2xyz(rpz) - else: - x = rpz + x = rpz2xyz(rpz) if basis.lower() == "xyz" else rpz + if dx1: + x_s = ( + rpz2xyz_vec(rpz_s, phi=rpz[:, :, 1]) + if basis.lower() == "xyz" + else rpz_s + ) + return x, x_s return x + def _compute_linking_number(self, params=None, grid=None): + """Calculate linking numbers for coils in the coilset. + + Parameters + ---------- + params : dict or array-like of dict, optional + Parameters to pass to coils, either the same for all coils or one for each. + grid : Grid or int, optional + Grid of coordinates to evaluate at. Defaults to a Linear grid. + If an integer, uses that many equally spaced points. + + Returns + ------- + link : ndarray, shape(num_coils, num_coils) + Linking number of each coil with each other coil. link=0 means they are not + linked, +/- 1 means the coils link each other in one direction or another. + + """ + if grid is None: + grid = LinearGrid(N=50) + dx = grid.spacing[:, 2] + x, x_s = self._compute_position(params, grid, dx1=True, basis="xyz") + link = _linking_number( + x[:, None], x[None, :], x_s[:, None], x_s[None, :], dx, dx + ) + return link / (4 * jnp.pi) + def _compute_A_or_B( self, coords, @@ -2307,7 +2368,7 @@ def compute( ) ] - def _compute_position(self, params=None, grid=None, **kwargs): + def _compute_position(self, params=None, grid=None, dx1=False, **kwargs): """Compute coil positions accounting for stellarator symmetry. Parameters @@ -2318,11 +2379,16 @@ def _compute_position(self, params=None, grid=None, **kwargs): Grid of coordinates to evaluate at. Defaults to a Linear grid. If an integer, uses that many equally spaced points. If array-like, should be 1 value per coil. + dx1 : bool + If True, also return dx/ds for each curve. Returns ------- x : ndarray, shape(len(self),source_grid.num_nodes,3) Coil positions, in [R,phi,Z] or [X,Y,Z] coordinates. + x_s : ndarray, shape(len(self),source_grid.num_nodes,3) + Coil position derivatives, in [R,phi,Z] or [X,Y,Z] coordinates. + Only returned if dx1=True. """ errorif( @@ -2331,15 +2397,17 @@ def _compute_position(self, params=None, grid=None, **kwargs): "grid must be supplied to MixedCoilSet._compute_position, since the " + "default grid for each coil could have a different number of nodes.", ) + kwargs.setdefault("basis", "xyz") params = self._make_arraylike(params) grid = self._make_arraylike(grid) - x = jnp.vstack( - [ - coil._compute_position(par, grd, **kwargs) - for coil, par, grd in zip(self.coils, params, grid) - ] - ) - return x + out = [] + for coil, par, grd in zip(self.coils, params, grid): + out.append(coil._compute_position(par, grd, dx1, **kwargs)) + if dx1: + x = jnp.vstack([foo[0] for foo in out]) + x_s = jnp.vstack([foo[1] for foo in out]) + return x, x_s + return jnp.vstack(out) def _compute_A_or_B( self, @@ -2795,3 +2863,18 @@ def flatten_coils(coilset): if ignore_groups: cset = cls(*flatten_coils(cset), check_intersection=check_intersection) return cset + + +@functools.partial(jnp.vectorize, signature="(m,3),(n,3),(m,3),(n,3),(m),(n)->()") +def _linking_number(x1, x2, x1_s, x2_s, dx1, dx2): + """Linking number between curves x1 and x2 with tangents x1_s, x2_s.""" + x1_s *= dx1[:, None] + x2_s *= dx2[:, None] + dx = x1[:, None, :] - x2[None, :, :] # shape(m,n,3) + dx_norm = safenorm(dx, axis=-1) # shape(m,n) + den = dx_norm**3 + dr1xdr2 = cross(x1_s[:, None, :], x2_s[None, :, :], axis=-1) # shape(m,n,3) + num = dot(dx, dr1xdr2, axis=-1) # shape(m,n) + small = dx_norm < jnp.finfo(x1.dtype).eps + ratio = jnp.where(small, 0.0, num / jnp.where(small, 1.0, den)) + return ratio.sum() diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py index 78903879e..c089e8a33 100644 --- a/desc/objectives/__init__.py +++ b/desc/objectives/__init__.py @@ -5,6 +5,7 @@ CoilCurrentLength, CoilCurvature, CoilLength, + CoilSetLinkingNumber, CoilSetMinDistance, CoilTorsion, PlasmaCoilSetMinDistance, diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index ef2e5d3e8..2ec2b39cf 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -1365,6 +1365,121 @@ def compute(self, field_params=None, constants=None): return Psi +class CoilSetLinkingNumber(_Objective): + """Prevents coils from becoming interlinked. + + The linking number of 2 curves is (approximately) 0 if they are not linked, and + (approximately) +/-1 if they are (with the sign indicating the helicity of the + linking). + + This objective returns a single value for each coil in the coilset, with that number + being the sum of the absolute value of the linking numbers of that coil with every + other coil in the coilset, approximating the number of other coils that are linked + + Parameters + ---------- + coil : CoilSet + Coil(s) that are to be optimized. + grid : Grid, list, optional + Collocation grid used to discretize each coil. Defaults to + ``LinearGrid(N=50)`` + + """ + + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", + bounds_default="``target=0``.", + coil=True, + ) + + _scalar = False + _units = "(dimensionless)" + _print_value_fmt = "Coil linking number: " + + def __init__( + self, + coil, + grid=None, + target=None, + bounds=None, + weight=1, + normalize=True, + normalize_target=True, + loss_function=None, + deriv_mode="auto", + jac_chunk_size=None, + name="coil-coil linking number", + ): + from desc.coils import CoilSet + + if target is None and bounds is None: + target = 0 + self._grid = grid + errorif( + not isinstance(coil, CoilSet), + ValueError, + "coil must be of type CoilSet, not an individual Coil", + ) + super().__init__( + things=coil, + target=target, + bounds=bounds, + weight=weight, + normalize=normalize, + normalize_target=normalize_target, + loss_function=loss_function, + deriv_mode=deriv_mode, + jac_chunk_size=jac_chunk_size, + name=name, + ) + + def build(self, use_jit=True, verbose=1): + """Build constant arrays. + + Parameters + ---------- + use_jit : bool, optional + Whether to just-in-time compile the objective and derivatives. + verbose : int, optional + Level of output. + + """ + coilset = self.things[0] + grid = self._grid or LinearGrid(N=50) + + self._dim_f = coilset.num_coils + self._constants = {"coilset": coilset, "grid": grid, "quad_weights": 1.0} + + super().build(use_jit=use_jit, verbose=verbose) + + def compute(self, params, constants=None): + """Compute linking numbers between coils. + + Parameters + ---------- + params : dict + Dictionary of coilset degrees of freedom, eg CoilSet.params_dict + constants : dict + Dictionary of constant data, eg transforms, profiles etc. + Defaults to self._constants. + + Returns + ------- + f : array of floats + For each coil, the sum of the absolute value of the linking numbers between + that coil and every other coil in the coilset, which approximates the + number of coils linked with that coil. + + """ + if constants is None: + constants = self.constants + link = constants["coilset"]._compute_linking_number( + params=params, grid=constants["grid"] + ) + + return jnp.abs(link).sum(axis=0) + + class SurfaceCurrentRegularization(_Objective): """Target the surface current magnitude. @@ -1384,51 +1499,21 @@ class SurfaceCurrentRegularization(_Objective): i.e. a CurrentPotentialField Intended to be used with a QuadraticFlux objective, to form - the REGCOIL algorithm described in [1]_ (if used with a - ``FourierCurrentPotentialField``). + a problem similar to the REGCOIL algorithm described in [1]_ (if used with a + ``FourierCurrentPotentialField``, is equivalent to the ``simple`` + regularization of the ``run_regcoil`` method). Parameters ---------- surface_current_field : CurrentPotentialField Surface current which is producing the magnetic field, the parameters of this will be optimized to minimize the objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. - Defaults to zero. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - Defaults to target=0 - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - When used with QuadraticFlux objective, this acts as the regularization - parameter (with w^2 = lambda), with 0 corresponding to no regularization. - The larger this parameter is, the less complex the surface current will be, - but the worse the normal field. - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target should be normalized before comparing to computed values. - if `normalize` is `True` and the target is in physical units, this should also - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. + source_grid : Grid, optional Collocation grid containing the nodes to evaluate current source at on the winding surface. If used in conjunction with the QuadraticFlux objective, with its ``field_grid`` matching this ``source_grid``, this replicates the REGCOIL algorithm described in [1]_. - name : str, optional - Name of the objective function. References ---------- @@ -1437,6 +1522,22 @@ class SurfaceCurrentRegularization(_Objective): """ + weight_str = ( + "weight : {float, ndarray}, optional" + "\n\tWeighting to apply to the Objective, relative to other Objectives." + "\n\tMust be broadcastable to to Objective.dim_f" + "\n\tWhen used with QuadraticFlux objective, this acts as the regularization" + "\n\tparameter (with w^2 = lambda), with 0 corresponding to no regularization." + "\n\tThe larger this parameter is, the less complex the surface current will " + "be,\n\tbut the worse the normal field." + ) + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", + bounds_default="``target=0``.", + loss_detail=" Note: has no effect for this objective.", + overwrite={"weight": weight_str}, + ) + _coordinates = "tz" _units = "A/m" _print_value_fmt = "Surface Current Regularization: " diff --git a/docs/api.rst b/docs/api.rst index 5c4ce4a29..7d06b9356 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -164,6 +164,7 @@ Objective Functions desc.objectives.CoilCurrentLength desc.objectives.CoilCurvature desc.objectives.CoilLength + desc.objectives.CoilSetLinkingNumber desc.objectives.CoilSetMinDistance desc.objectives.CoilTorsion desc.objectives.CurrentDensity diff --git a/docs/api_objectives.rst b/docs/api_objectives.rst index 7fdd2e1a8..529cc5072 100644 --- a/docs/api_objectives.rst +++ b/docs/api_objectives.rst @@ -100,6 +100,7 @@ Coil Optimization desc.objectives.CoilLength desc.objectives.CoilCurvature desc.objectives.CoilTorsion + desc.objectives.CoilSetLinkingNumber desc.objectives.CoilSetMinDistance desc.objectives.PlasmaCoilSetMinDistance desc.objectives.CoilCurrentLength diff --git a/tests/test_coils.py b/tests/test_coils.py index 71127660d..368a2fd12 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -1282,3 +1282,26 @@ def test_repr(): coils.name = "MyCoils" assert "MyCoils" in str(coils) + + +@pytest.mark.unit +def test_linking_number(): + """Test calculation of linking number.""" + coil = FourierPlanarCoil(center=[10, 1, 0]) + grid = LinearGrid(N=25) + # regular modular coilset from symmetry, so that there are 10 coils, half going + # one way and half going the other way + coilset = CoilSet(coil, NFP=5, sym=True) + coil2 = FourierRZCoil() + # add a coil along the axis that links all the other coils + coilset2 = MixedCoilSet(coilset, coil2) + link = coilset2._compute_linking_number(grid=grid) + + # modular coils don't link each other + np.testing.assert_allclose(link[:-1, :-1], 0, atol=1e-14) + # axis coil doesn't link itself + np.testing.assert_allclose(link[-1, -1], 0, atol=1e-14) + # we expect the axis coil to link all the modular coils, with alternating sign + # due to alternating orientation of the coils due to symmetry. + expected = [1, -1] * 5 + np.testing.assert_allclose(link[-1, :-1], expected, rtol=1e-3) diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 0ad70b88e..8a25eb63f 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -44,6 +44,7 @@ CoilCurrentLength, CoilCurvature, CoilLength, + CoilSetLinkingNumber, CoilSetMinDistance, CoilTorsion, Elongation, @@ -1195,6 +1196,25 @@ def test(eq, field, correct_value, rtol=1e-14, grid=None): # test on field with no vector potential test(eq, PoloidalMagneticField(1, 1, 1), 0.0) + @pytest.mark.unit + def test_coil_linking_number(self): + """Test for linking number objective.""" + coil = FourierPlanarCoil(center=[10, 1, 0]) + # regular modular coilset from symmetry, so that there are 10 coils, half going + # one way and half going the other way + coilset = CoilSet.from_symmetry(coil, NFP=5, sym=True) + coil2 = FourierRZCoil() + # add a coil along the axis that links all the other coils + coilset2 = MixedCoilSet(coilset, coil2) + + obj = CoilSetLinkingNumber(coilset2) + obj.build() + out = obj.compute_scaled_error(coilset2.params_dict) + # the modular coils all link 1 other coil (the axis) + # while the axis links all 10 modular coils + expected = np.array([1] * 10 + [10]) + np.testing.assert_allclose(out, expected, rtol=1e-3) + @pytest.mark.unit def test_signed_plasma_vessel_distance(self): """Test calculation of signed distance from plasma to vessel.""" @@ -2246,6 +2266,7 @@ class TestComputeScalarResolution: CoilCurrentLength, CoilCurvature, CoilLength, + CoilSetLinkingNumber, CoilSetMinDistance, CoilTorsion, FusionPower, @@ -2627,7 +2648,14 @@ def test_compute_scalar_resolution_others(self, objective): @pytest.mark.regression @pytest.mark.parametrize( "objective", - [CoilLength, CoilTorsion, CoilCurvature, CoilCurrentLength, CoilSetMinDistance], + [ + CoilLength, + CoilTorsion, + CoilCurvature, + CoilCurrentLength, + CoilSetMinDistance, + CoilSetLinkingNumber, + ], ) def test_compute_scalar_resolution_coils(self, objective): """Coil objectives.""" @@ -2663,6 +2691,7 @@ class TestObjectiveNaNGrad: CoilLength, CoilCurrentLength, CoilCurvature, + CoilSetLinkingNumber, CoilSetMinDistance, CoilTorsion, ForceBalanceAnisotropic, @@ -2865,7 +2894,14 @@ def test_objective_no_nangrad(self, objective): @pytest.mark.unit @pytest.mark.parametrize( "objective", - [CoilLength, CoilTorsion, CoilCurvature, CoilCurrentLength, CoilSetMinDistance], + [ + CoilLength, + CoilTorsion, + CoilCurvature, + CoilCurrentLength, + CoilSetMinDistance, + CoilSetLinkingNumber, + ], ) def test_objective_no_nangrad_coils(self, objective): """Coil objectives."""