Add Coil Arclength Differential Variance Objective

desc/objectives/_coils.py

@@ -607,6 +607,142 @@ def compute(self, params, constants=None):
         return out
 
 
+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.
+
+    See Wechsung 2021 supporting information for more detail.
+
+    Parameters
+    ----------
+    coil : CoilSet or Coil
+        Coil(s) that are to be optimized
+    target : float, ndarray, optional
+        Target value(s) of the objective. Only used if bounds is None.
+        Must be broadcastable to Objective.dim_f. If array, it has to
+        be flattened according to the number of inputs. Defaults to ``target=2*np.pi``.
+    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=2*np.pi``.
+    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.
+        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. Operates over all coils, not each individual coil.
+    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.
+    grid : Grid, optional
+        Collocation grid containing the nodes to evaluate at.
+        Defaults to ``LinearGrid(N=2 * coil.N + 5)``
+    name : str, optional
+        Name of the objective function.
+    """
+
+    _scalar = False  # Not always a scalar, if a coilset is passed in
+    _units = "(m^2)"
+    _print_value_fmt = "Coil Arclength Variance: {:10.3e} "
+
+    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",
+    ):
+        self._coils = coils
+        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.
+
+        """
+        from desc.coils import CoilSet, _Coil
+
+        super().build(use_jit=use_jit, verbose=verbose)
+
+        if self._normalize:
+            self._normalization = self._scales["a"] ** 2
+
+        # TODO: repeated code but maybe it's fine
+        flattened_coils = tree_flatten(
+            self._coils,
+            is_leaf=lambda x: isinstance(x, _Coil) and not isinstance(x, CoilSet),
+        )[0]
+        flattened_coils = (
+            [flattened_coils[0]]
+            if not isinstance(self._coils, CoilSet)
+            else flattened_coils
+        )
+        self._dim_f = len(flattened_coils)
+        self._constants["quad_weights"] = 1
+
+    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.
+        """
+        data = super().compute(params, constants=constants)
+        data = tree_flatten(data, is_leaf=lambda x: isinstance(x, dict))[0]
+        out = jnp.array([jnp.var(dat["length"]) for dat in data])
+        return out
+
+
 class QuadraticFlux(_Objective):
     """Target B*n = 0 on LCFS.