LinkingCurrent objective

desc/objectives/__init__.py

@@ -8,6 +8,7 @@
     CoilSetLinkingNumber,
     CoilSetMinDistance,
     CoilTorsion,
+    LinkingCurrent,
     PlasmaCoilSetMinDistance,
     QuadraticFlux,
     ToroidalFlux,

desc/objectives/_coils.py

@@ -1,6 +1,7 @@
 import numbers
 
 import numpy as np
+from scipy.constants import mu_0
 
 from desc.backend import (
     fori_loop,
@@ -14,7 +15,7 @@
 from desc.compute.utils import _compute as compute_fun
 from desc.grid import LinearGrid, _Grid
 from desc.integrals import compute_B_plasma
-from desc.utils import Timer, errorif, safenorm, warnif
+from desc.utils import Timer, broadcast_tree, errorif, safenorm, warnif
 
 from .normalization import compute_scaling_factors
 from .objective_funs import _Objective, collect_docs
@@ -1364,6 +1365,168 @@
         return Psi
 
 
+class LinkingCurrent(_Objective):
+    """Target the self-consistent poloidal linking current between the plasma and coils.
+
+    Assumes the coil topology does not change (ie the linking number with the plasma
+    is fixed).
+
+    Parameters
+    ----------
+    eq : Equilibrium
+        Equilibrium that will be optimized to satisfy the Objective.
+    coil : CoilSet
+        Coil(s) that are to be optimized.
+    grid : Grid, optional
+        Collocation grid containing the nodes to evaluate plasma current at.
+        Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym)``.
+
+    """
+
+    __doc__ = __doc__.rstrip() + collect_docs(
+        target_default="``target=0``.",
+        bounds_default="``target=0``.",
+    )
+
+    _scalar = True
+    _units = "(A)"
+    _print_value_fmt = "Linking current: "
+
+    def __init__(
+        self,
+        eq,
+        coil,
+        target=None,
+        bounds=None,
+        weight=1,
+        normalize=True,
+        normalize_target=True,
+        loss_function=None,
+        deriv_mode="auto",
+        grid=None,
+        jac_chunk_size=None,
+        name="linking current",
+    ):
+        if target is None and bounds is None:
+            target = 0
+        self._grid = grid
+        super().__init__(
+            things=[eq, 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.
+
+        """
+        eq = self.things[0]
+        coil = self.things[1]
+        grid = self._grid or LinearGrid(
+            M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym
+        )
+        warnif(
+            not np.allclose(grid.nodes[:, 0], 1),
+            UserWarning,
+            "grid includes interior points, should be rho=1.",
+        )
+
+        self._dim_f = 1
+        self._data_keys = ["G"]
+
+        all_params = tree_map(lambda dim: np.arange(dim), coil.dimensions)
+        current_params = tree_map(lambda idx: {"current": idx}, True)
+        # indices of coil currents
+        self._indices = tree_leaves(broadcast_tree(current_params, all_params))
+        self._num_coils = coil.num_coils
+
+        profiles = get_profiles(self._data_keys, obj=eq, grid=grid)
+        transforms = get_transforms(self._data_keys, obj=eq, grid=grid)
+
+        # compute linking number of coils with plasma. To do this we add a fake "coil"
+        # along the magnetic axis and compute the linking number of that coilset
+        from desc.coils import FourierRZCoil, MixedCoilSet
+
+        axis_coil = FourierRZCoil(
+            1.0,
+            eq.axis.R_n,
+            eq.axis.Z_n,
+            eq.axis.R_basis.modes[:, 2],
+            eq.axis.Z_basis.modes[:, 2],
+            eq.axis.NFP,
+        )
+        dummy_coilset = MixedCoilSet(axis_coil, coil)
+        # linking number for coils with axis
+        link = np.round(dummy_coilset._compute_linking_number())[0, 1:]
+
+        self._constants = {
+            "profiles": profiles,
+            "transforms": transforms,
+            "quad_weights": 1.0,
+            "link": link,
+        }
+
+        if self._normalize:
+            params = tree_leaves(
+                coil.params_dict, is_leaf=lambda x: isinstance(x, dict)
+            )
+            self._normalization = np.sum([np.abs(param["current"]) for param in params])
+
+        super().build(use_jit=use_jit, verbose=verbose)
+
+    def compute(self, eq_params, coil_params, constants=None):
+        """Compute linking current error.
+
+        Parameters
+        ----------
+        eq_params : dict
+            Dictionary of equilibrium degrees of freedom, eg ``Equilibrium.params_dict``
+        coil_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
+            Linking current error.
+
+        """
+        if constants is None:
+            constants = self.constants
+        data = compute_fun(
+            "desc.equilibrium.equilibrium.Equilibrium",
+            self._data_keys,
+            params=eq_params,
+            transforms=constants["transforms"],
+            profiles=constants["profiles"],
+        )
+        eq_linking_current = 2 * jnp.pi * data["G"][0] / mu_0
+        coil_currents = jnp.concatenate(
+            [
+                jnp.atleast_1d(param[idx])
+                for param, idx in zip(tree_leaves(coil_params), self._indices)
+            ]
+        )
+        coil_linking_current = jnp.sum(constants["link"] * coil_currents)
+        return eq_linking_current - coil_linking_current
+
+
 class CoilSetLinkingNumber(_Objective):
     """Prevents coils from becoming interlinked.
 

tests/test_objective_funs.py

@@ -56,6 +56,7 @@
     HeatingPowerISS04,
     Isodynamicity,
     LinearObjectiveFromUser,
+    LinkingCurrent,
     MagneticWell,
     MeanCurvature,
     MercierStability,
@@ -1297,6 +1298,18 @@ def test_signed_plasma_vessel_distance(self):
         assert abs(d.max() - (-a_s)) < 1e-14
         assert abs(d.min() - (-a_s)) < grid.spacing[0, 1] * a_s
 
+    @pytest.mark.unit
+    def test_linking_current(self):
+        """Test calculation of signed linking current from coils to plasma."""
+        eq = Equilibrium()
+        G = eq.compute("G", grid=LinearGrid(rho=1.0))["G"] * 2 * jnp.pi / mu_0
+        coil = FourierPlanarCoil(current=G / 8, center=[10, 1, 0])
+        coilset = CoilSet.from_symmetry(coil, NFP=4, sym=True)
+        obj = LinkingCurrent(eq, coilset)
+        obj.build()
+        f = obj.compute(eq.params_dict, coilset.params_dict)
+        np.testing.assert_allclose(f, 0)
+
 
 @pytest.mark.regression
 def test_derivative_modes():
@@ -2271,6 +2284,7 @@ class TestComputeScalarResolution:
         FusionPower,
         GenericObjective,
         HeatingPowerISS04,
+        LinkingCurrent,
         Omnigenity,
         PlasmaCoilSetMinDistance,
         PlasmaVesselDistance,
@@ -2650,6 +2664,26 @@ def test_compute_scalar_resolution_coils(self, objective):
             f[i] = obj.compute_scalar(obj.x())
         np.testing.assert_allclose(f, f[-1], rtol=1e-2, atol=1e-12)
 
+    @pytest.mark.unit
+    def test_compute_scalar_resolution_linking_current(self):
+        """LinkingCurrent."""
+        coil = FourierPlanarCoil(center=[10, 1, 0])
+        eq = Equilibrium()
+        coilset = CoilSet.from_symmetry(coil, NFP=4, sym=True)
+        f = np.zeros_like(self.res_array, dtype=float)
+        for i, res in enumerate(self.res_array):
+            obj = ObjectiveFunction(
+                LinkingCurrent(
+                    eq,
+                    coilset,
+                    grid=LinearGrid(M=int(eq.M_grid * res), N=int(eq.N_grid * res)),
+                ),
+                use_jit=False,
+            )
+            obj.build(verbose=0)
+            f[i] = obj.compute_scalar(obj.x())
+        np.testing.assert_allclose(f, f[-1], rtol=1e-2, atol=1e-12)
+
 
 class TestObjectiveNaNGrad:
     """Make sure reverse mode AD works correctly for all objectives."""
@@ -2677,6 +2711,7 @@ class TestObjectiveNaNGrad:
         ForceBalanceAnisotropic,
         FusionPower,
         HeatingPowerISS04,
+        LinkingCurrent,
         Omnigenity,
         PlasmaCoilSetMinDistance,
         PlasmaVesselDistance,
@@ -2919,6 +2954,17 @@ def test_objective_no_nangrad_ballooning(self):
         g = obj.grad(obj.x())
         assert not np.any(np.isnan(g))
 
+    @pytest.mark.unit
+    def test_objective_no_nangrad_linking_current(self):
+        """LinkingCurrent."""
+        coil = FourierPlanarCoil(center=[10, 1, 0])
+        coilset = CoilSet.from_symmetry(coil, NFP=4, sym=True)
+        eq = Equilibrium()
+        obj = ObjectiveFunction(LinkingCurrent(eq, coilset))
+        obj.build()
+        g = obj.grad(obj.x())
+        assert not np.any(np.isnan(g))
+
 
 @pytest.mark.unit
 def test_asymmetric_normalization():
@@ -2943,6 +2989,7 @@ def test_asymmetric_normalization():
         assert np.all(np.isfinite(val))
 
 
+@pytest.mark.unit
 def test_objective_print_widths():
     """Test that the objective's name is shorter than max."""
     subclasses = _Objective.__subclasses__()
@@ -2971,6 +3018,7 @@ def test_objective_print_widths():
                 )
 
 
+@pytest.mark.unit
 def test_objective_docstring():
     """Test that the objective docstring and collect_docs are consistent."""
     objective_docs = _Objective.__doc__.rstrip()