FixSumCoilCurrent objective

desc/objectives/__init__.py

@@ -77,6 +77,7 @@
     FixPressure,
     FixPsi,
     FixSheetCurrent,
+    FixSumCoilCurrent,
     FixSumModesLambda,
     FixSumModesR,
     FixSumModesZ,

desc/objectives/linear_objectives.py

@@ -151,7 +151,7 @@
         weight=1,
         normalize=True,
         normalize_target=True,
-        name="Fixed parameters",
+        name="fixed parameters",
     ):
         self._params = params
         super().__init__(
@@ -2726,6 +2726,153 @@
         super().build(use_jit=use_jit, verbose=verbose)
 
 
+class FixSumCoilCurrent(FixCoilCurrent):
+    """Fixes the sum of coil current(s) in a Coil or CoilSet.
+
+    NOTE: When using this objective, take care in knowing the signs of the current in
+    the coils and the orientations of the coils. It is possible for coils with the same
+    signs of their current to have currents flowing in differing directions in physical
+    space due to the orientation of the coils.
+
+    Parameters
+    ----------
+    coil : Coil
+        Coil(s) that will be optimized to satisfy 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.
+        Default is the objective value for the coil.
+    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.
+        Default is to use the target instead.
+    weight : {float, ndarray}, optional
+        Weighting to apply to the Objective, relative to other Objectives.
+        Must be broadcastable to to Objective.dim_f
+    normalize : bool, optional
+        Whether to compute the error in physical units or non-dimensionalize.
+    normalize_target : bool, optional
+        Whether target and bounds 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.
+    indices : nested list of bool, optional
+        Pytree of bool specifying which coil currents to sum together.
+        See the example for how to use this on a mixed coil set.
+        If True/False sums all/none of the coil currents.
+    name : str, optional
+        Name of the objective function.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        import numpy as np
+        from desc.coils import (
+            CoilSet, FourierPlanarCoil, FourierRZCoil, FourierXYZCoil, MixedCoilSet
+        )
+        from desc.objectives import FixSumCoilCurrent
+
+        # toroidal field coil set with 4 coils
+        tf_coil = FourierPlanarCoil(
+            current=3, center=[2, 0, 0], normal=[0, 1, 0], r_n=[1]
+        )
+        tf_coilset = CoilSet.linspaced_angular(tf_coil, n=4)
+        # vertical field coil set with 3 coils
+        vf_coil = FourierRZCoil(current=-1, R_n=3, Z_n=-1)
+        vf_coilset = CoilSet.linspaced_linear(
+            vf_coil, displacement=[0, 0, 2], n=3, endpoint=True
+        )
+        # another single coil
+        xyz_coil = FourierXYZCoil(current=2)
+        # full coil set with TF coils, VF coils, and other single coil
+        full_coilset = MixedCoilSet((tf_coilset, vf_coilset, xyz_coil))
+
+        # equilibrium G(rho=1) determines the necessary net poloidal current through
+        # the coils (as dictated by Ampere's law)
+        grid_at_surf = LinearGrid(rho=1.0, M=eq.M_grid, N=eq.N_grid)
+        G_tot = eq.compute("G", grid=grid_at_surf)["G"][0] / (mu_0 * 2 * jnp.pi)
+
+        # to use this objective to satisfy Ampere's law for the targeted equilibrium,
+        # only coils that link the equilibrium poloidally should be included in the sum,
+        # which is the TF coil set and the FourierXYZ coil, but not the VF coil set
+        obj = FixSumCoilCurrent(full_coilset, indices=[True, False, True], target=G_tot)
+
+    """
+
+    _scalar = True
+    _linear = True
+    _fixed = False
+    _units = "(A)"
+    _print_value_fmt = "Summed coil current error: {:10.3e} "
+
+    def __init__(
+        self,
+        coil,
+        target=None,
+        bounds=None,
+        weight=1,
+        normalize=True,
+        normalize_target=True,
+        indices=True,
+        name="summed coil current",
+    ):
+        self._default_target = False
+        if target is None and bounds is None:
+            self._default_target = True
+        super().__init__(
+            coil=coil,
+            target=target,
+            bounds=bounds,
+            weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
+            indices=indices,
+            name=name,
+        )
+
+    def build(self, use_jit=False, 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.
+
+        """
+        super().build(use_jit=use_jit, verbose=verbose)
+        self._dim_f = 1
+        if self._default_target:
+            self.update_target(thing=self.things[0])
+
+    def compute(self, params, constants=None):
+        """Compute sum of coil currents.
+
+        Parameters
+        ----------
+        params : list of dict
+            List of dictionaries of degrees of freedom, eg CoilSet.params_dict
+        constants : dict
+            Dictionary of constant data, eg transforms, profiles etc. Defaults to
+            self.constants
+
+        Returns
+        -------
+        f : ndarray
+            Sum of coil currents.
+
+        """
+        return jnp.sum(
+            jnp.concatenate(
+                [
+                    jnp.atleast_1d(param[idx])
+                    for param, idx in zip(tree_leaves(params), self._indices)
+                ]
+            )
+        )
+
+
 class FixOmniWell(FixParameters):
     """Fixes OmnigenousField.B_lm coefficients.
 

docs/api.rst

@@ -187,6 +187,8 @@ Objective Functions
     desc.objectives.FixParameters
     desc.objectives.FixPressure
     desc.objectives.FixPsi
+    desc.objectives.FixSumCoilCurrent
+    desc.objectives.FixSumModesLambda
     desc.objectives.FixSumModesR
     desc.objectives.FixSumModesZ
     desc.objectives.FixThetaSFL

tests/test_linear_objectives.py

@@ -6,6 +6,7 @@
 from qsc import Qsc
 
 import desc.examples
+from desc.backend import jnp
 from desc.equilibrium import Equilibrium
 from desc.geometry import FourierRZToroidalSurface
 from desc.grid import LinearGrid
@@ -37,11 +38,13 @@
     FixParameters,
     FixPressure,
     FixPsi,
+    FixSumCoilCurrent,
     FixSumModesLambda,
     FixSumModesR,
     FixSumModesZ,
     FixThetaSFL,
     GenericObjective,
+    LinearObjectiveFromUser,
     ObjectiveFunction,
     get_equilibrium_objective,
     get_fixed_axis_constraints,
@@ -1006,3 +1009,51 @@ def test_fix_coil_current(DummyMixedCoilSet):
     obj.build()
     assert obj.dim_f == 4
     np.testing.assert_allclose(obj.target, [3, -1, -1, 2])
+
+
+@pytest.mark.unit
+def test_fix_sum_coil_current(DummyMixedCoilSet):
+    """Tests FixSumCoilCurrent."""
+    coilset = load(load_from=str(DummyMixedCoilSet["output_path"]), file_format="hdf5")
+
+    # sum a single coil current
+    obj = FixSumCoilCurrent(coil=coilset.coils[1].coils[0])
+    obj.build()
+    params = coilset.coils[1].coils[0].params_dict
+    np.testing.assert_allclose(obj.compute(params), -1)
+
+    # sum all coil currents
+    obj = FixSumCoilCurrent(coil=coilset)
+    obj.build()
+    params = coilset.params_dict
+    np.testing.assert_allclose(obj.compute(params), 3)
+    # the default target should be the original sum
+    np.testing.assert_allclose(obj.compute_scaled_error(params), 0)
+
+    # only sum currents of some coils in the coil set
+    obj = FixSumCoilCurrent(
+        coil=coilset, indices=[[True], [True, False, True], True, False]
+    )
+    obj.build()
+    np.testing.assert_allclose(obj.compute(params), 3)
+
+
+@pytest.mark.unit
+def test_linear_objective_from_user_on_collection(DummyCoilSet):
+    """Test LinearObjectiveFromUser on an OptimizableCollection."""
+    # test that LinearObjectiveFromUser can be used for the same functionality as
+    # FixSumCoilCurrent to sum all currents in a CoilSet
+
+    coilset = load(load_from=str(DummyCoilSet["output_path_asym"]), file_format="hdf5")
+    params = coilset.params_dict
+
+    obj1 = FixSumCoilCurrent(coil=coilset)
+    obj1.build()
+
+    obj2 = LinearObjectiveFromUser(
+        lambda params: jnp.sum(jnp.array([param["current"] for param in params])),
+        coilset,
+    )
+    obj2.build()
+
+    np.testing.assert_allclose(obj1.compute(params), obj2.compute(params))