Add Coil Arclength Differential Variance Objective (#1113)

- Adds objective similar to the one detailed in [Wechsung
2021](https://www.pnas.org/doi/full/10.1073/pnas.2202084119#supplementary-materials)
to penalize the variance of the arclength differential along a curve,
which was claimed to be useful for avoiding issues with parameterization
non-uniqueness in optimization.
- [x] add a test for this
- [x] Add simplified version, this is equal to the [simsopt
implementation](https://github.com/hiddenSymmetries/simsopt/blob/988e15c1b742689fedfbd9a070a1d24d2c900bca/src/simsopt/geo/curveobjectives.py#L366)
with `nintervals="full"`

From reference:

![image](https://github.com/PlasmaControl/DESC/assets/37969854/7e4e6731-cd0b-49cd-836f-8645b66037a2)

![image](https://github.com/PlasmaControl/DESC/assets/37969854/77b20fcf-133d-4a4d-87a7-b6ed3e63c5c8)

desc/compute/_curve.py

@@ -571,8 +571,8 @@ def _x_FourierXYZCurve(params, transforms, profiles, data, **kwargs):
     dim=3,
     params=["X_n", "Y_n", "Z_n", "rotmat"],
     transforms={
-        "X": [[0, 0, 0], [0, 0, 1]],
-        "Y": [[0, 0, 0], [0, 0, 1]],
+        "X": [[0, 0, 1]],
+        "Y": [[0, 0, 1]],
         "Z": [[0, 0, 1]],
     },
     profiles=[],

desc/objectives/__init__.py

@@ -2,6 +2,7 @@
 
 from ._bootstrap import BootstrapRedlConsistency
 from ._coils import (
+    CoilArclengthVariance,
     CoilCurrentLength,
     CoilCurvature,
     CoilLength,

desc/objectives/_coils.py

@@ -177,7 +177,7 @@ def compute(self, params, constants=None):
         Returns
         -------
         f : float or array of floats
-            Coil length.
+            Coil objective value(s).
 
         """
         if constants is None:
@@ -958,6 +958,128 @@ def body(k):
         return min_dist_per_coil
 
 
+class CoilArclengthVariance(_CoilObjective):
+    """Variance of ||dx/ds|| along the curve.
+
+    This objective is meant to combat any issues corresponding to non-uniqueness of
+    the representation of a curve, in that the same physical curve can be represented
+    by different parametrizations by changing the curve parameter [1]_. Note that this
+    objective has no effect for ``FourierRZCoil`` and ``FourierPlanarCoil`` which have
+    a single unique parameterization (the objective will always return 0 for these
+    types).
+
+    References
+    ----------
+    .. [1] Wechsung, et al. "Precise stellarator quasi-symmetry can be achieved
+       with electromagnetic coils." PNAS (2022)
+
+    Parameters
+    ----------
+    coil : CoilSet or Coil
+        Coil(s) that are to be optimized
+    grid : Grid, optional
+        Collocation grid containing the nodes to evaluate at.
+        Defaults to ``LinearGrid(N=2 * coil.N + 5)``
+
+    """
+
+    __doc__ = __doc__.rstrip() + collect_docs(
+        target_default="``target=0``.",
+        bounds_default="``target=0``.",
+        coil=True,
+    )
+
+    _scalar = False  # Not always a scalar, if a coilset is passed in
+    _units = "(m^2)"
+    _print_value_fmt = "Coil Arclength Variance: "
+
+    def __init__(
+        self,
+        coils,
+        target=None,
+        bounds=None,
+        weight=1,
+        normalize=True,
+        normalize_target=True,
+        loss_function=None,
+        deriv_mode="auto",
+        grid=None,
+        name="coil arclength variance",
+    ):
+        if target is None and bounds is None:
+            target = 0
+
+        super().__init__(
+            coils,
+            ["x_s"],
+            target=target,
+            bounds=bounds,
+            weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
+            loss_function=loss_function,
+            deriv_mode=deriv_mode,
+            grid=grid,
+            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.
+
+        """
+        super().build(use_jit=use_jit, verbose=verbose)
+
+        self._dim_f = self._num_coils
+        self._constants["quad_weights"] = 1
+
+        coilset = self.things[0]
+        # local import to avoid circular import
+        from desc.coils import CoilSet, FourierXYZCoil, SplineXYZCoil, _Coil
+
+        def _is_single_coil(c):
+            return isinstance(c, _Coil) and not isinstance(c, CoilSet)
+
+        coils = tree_leaves(coilset, is_leaf=_is_single_coil)
+        self._constants["mask"] = np.array(
+            [int(isinstance(coil, (FourierXYZCoil, SplineXYZCoil))) for coil in coils]
+        )
+
+        if self._normalize:
+            self._normalization = np.mean([scale["a"] ** 2 for scale in self._scales])
+
+        _Objective.build(self, use_jit=use_jit, verbose=verbose)
+
+    def compute(self, params, constants=None):
+        """Compute coil arclength variance.
+
+        Parameters
+        ----------
+        params : dict
+            Dictionary of the coil's degrees of freedom.
+        constants : dict
+            Dictionary of constant data, eg transforms, profiles etc. Defaults to
+            self._constants.
+
+        Returns
+        -------
+        f : float or array of floats
+            Coil arclength variance.
+        """
+        if constants is None:
+            constants = self.constants
+        data = super().compute(params, constants=constants)
+        data = tree_leaves(data, is_leaf=lambda x: isinstance(x, dict))
+        out = jnp.array([jnp.var(jnp.linalg.norm(dat["x_s"], axis=1)) for dat in data])
+        return out * constants["mask"]
+
+
 class QuadraticFlux(_Objective):
     """Target B*n = 0 on LCFS.
 

desc/transform.py

@@ -43,6 +43,7 @@ class Transform(IOAble):
     """
 
     _io_attrs_ = ["_grid", "_basis", "_derivatives", "_rcond", "_method"]
+    _static_attrs = ["_derivatives"]
 
     def __init__(
         self,

docs/api.rst

@@ -160,6 +160,7 @@ Objective Functions
     desc.objectives.BootstrapRedlConsistency
     desc.objectives.BoundaryError
     desc.objectives.BScaleLength
+    desc.objectives.CoilArclengthVariance
     desc.objectives.CoilCurrentLength
     desc.objectives.CoilCurvature
     desc.objectives.CoilLength

docs/api_objectives.rst

@@ -104,6 +104,7 @@ Coil Optimization
     desc.objectives.CoilSetMinDistance
     desc.objectives.PlasmaCoilSetMinDistance
     desc.objectives.CoilCurrentLength
+    desc.objectives.CoilArclengthVariance
     desc.objectives.ToroidalFlux
 
 

tests/test_examples.py

@@ -34,6 +34,7 @@
     AspectRatio,
     BallooningStability,
     BoundaryError,
+    CoilArclengthVariance,
     CoilCurvature,
     CoilLength,
     CoilSetMinDistance,
@@ -1568,6 +1569,45 @@ def circle_constraint(params):
     )
 
 
+@pytest.mark.unit
+@pytest.mark.optimize
+def test_coil_arclength_optimization():
+    """Test coil arclength variance optimization."""
+    c1 = FourierXYZCoil()
+    c1.change_resolution(N=5)
+    target_length = 2 * c1.compute("length")["length"]
+    obj = ObjectiveFunction(
+        (
+            CoilLength(c1, target=target_length),
+            CoilCurvature(c1, target=1, weight=1e-2),
+        )
+    )
+    obj2 = ObjectiveFunction(
+        (
+            CoilLength(c1, target=target_length),
+            CoilCurvature(c1, target=1, weight=1e-2),
+            CoilArclengthVariance(c1, target=0, weight=100),
+        )
+    )
+    opt = Optimizer("lsq-exact")
+    (coil_opt_without_arc_obj,), _ = opt.optimize(
+        c1, objective=obj, verbose=3, copy=True, ftol=1e-6
+    )
+    (coil_opt_with_arc_obj,), _ = opt.optimize(
+        c1, objective=obj2, verbose=3, copy=True, ftol=1e-6, maxiter=200
+    )
+    xs1 = coil_opt_with_arc_obj.compute("x_s")["x_s"]
+    xs2 = coil_opt_without_arc_obj.compute("x_s")["x_s"]
+    np.testing.assert_allclose(
+        coil_opt_without_arc_obj.compute("length")["length"], target_length, rtol=1e-4
+    )
+    np.testing.assert_allclose(
+        coil_opt_with_arc_obj.compute("length")["length"], target_length, rtol=1e-4
+    )
+    np.testing.assert_allclose(np.var(np.linalg.norm(xs1, axis=1)), 0, atol=1e-5)
+    assert np.var(np.linalg.norm(xs1, axis=1)) < np.var(np.linalg.norm(xs2, axis=1))
+
+
 @pytest.mark.regression
 def test_ballooning_stability_opt():
     """Perform ballooning stability optimization with DESC."""

tests/test_objective_funs.py

@@ -41,6 +41,7 @@
     BootstrapRedlConsistency,
     BoundaryError,
     BScaleLength,
+    CoilArclengthVariance,
     CoilCurrentLength,
     CoilCurvature,
     CoilLength,
@@ -2264,6 +2265,7 @@ class TestComputeScalarResolution:
         # these require special logic
         BootstrapRedlConsistency,
         BoundaryError,
+        CoilArclengthVariance,
         CoilCurrentLength,
         CoilCurvature,
         CoilLength,
@@ -2631,12 +2633,13 @@ def test_compute_scalar_resolution_others(self, objective):
     @pytest.mark.parametrize(
         "objective",
         [
+            CoilArclengthVariance,
+            CoilCurrentLength,
+            CoilCurvature,
             CoilLength,
             CoilTorsion,
-            CoilCurvature,
-            CoilCurrentLength,
-            CoilSetMinDistance,
             CoilSetLinkingNumber,
+            CoilSetMinDistance,
         ],
     )
     def test_compute_scalar_resolution_coils(self, objective):
@@ -2670,6 +2673,7 @@ class TestObjectiveNaNGrad:
         BallooningStability,
         BootstrapRedlConsistency,
         BoundaryError,
+        CoilArclengthVariance,
         CoilLength,
         CoilCurrentLength,
         CoilCurvature,
@@ -2866,12 +2870,13 @@ def test_objective_no_nangrad(self, objective):
     @pytest.mark.parametrize(
         "objective",
         [
+            CoilArclengthVariance,
+            CoilCurrentLength,
+            CoilCurvature,
             CoilLength,
             CoilTorsion,
-            CoilCurvature,
-            CoilCurrentLength,
-            CoilSetMinDistance,
             CoilSetLinkingNumber,
+            CoilSetMinDistance,
         ],
     )
     def test_objective_no_nangrad_coils(self, objective):