Merge branch 'master' into rg/NFP_fac

.pre-commit-config.yaml

@@ -28,3 +28,12 @@ repos:
     rev: v3.2.0
     hooks:
       - id: pyupgrade
+  - repo: local
+    hooks:
+      - id: check_unmarked_tests
+        name: check_unmarked_tests
+        entry: devtools/check_unmarked_tests.sh
+        language: script
+        files: ^tests/
+        types: [python]
+        pass_filenames: true

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,

devtools/check_unmarked_tests.sh

@@ -0,0 +1,22 @@
+#!/bin/sh
+
+# Start the timer using date (in seconds since epoch)
+start_time=$(date +%s)
+
+echo "Files to check: $@"
+# Collect unmarked tests for the specific file and suppress errors
+unmarked=$(pytest "$@" --collect-only -m "not unit and not regression" -q 2> /dev/null | head -n -2)
+
+# Count the number of unmarked tests found, ignoring empty lines
+num_unmarked=$(echo "$unmarked" | sed '/^\s*$/d' | wc -l)
+
+# If there are any unmarked tests, print them and exit with status 1
+if [ "$num_unmarked" -gt 0 ]; then
+    echo "----found unmarked tests----"
+    echo "$unmarked"
+    # Calculate the elapsed time and print with a newline
+    end_time=$(date +%s)
+    elapsed_time=$((end_time - start_time))
+    printf "\nTime taken: %d seconds" "$elapsed_time"
+    exit 1
+fi

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_curves.py

@@ -262,6 +262,7 @@ def test_from_input_file(self):
         np.testing.assert_allclose(curve3.NFP, curve4.NFP)
         np.testing.assert_allclose(curve3.sym, curve4.sym)
 
+    @pytest.mark.unit
     def test_to_FourierRZCurve(self):
         """Test conversion to FourierRZCurve."""
         xyz = FourierXYZCurve(modes=[-1, 1], X_n=[0, 10], Y_n=[10, 0], Z_n=[0, 0])

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,
@@ -914,6 +915,7 @@ def test(coil, grid=None):
         test(mixed_coils)
         test(nested_coils, grid=grid)
 
+    @pytest.mark.unit
     def test_coil_type_error(self):
         """Tests error when objective is not passed a coil."""
         curve = FourierPlanarCurve(r_n=2, basis="rpz")
@@ -1104,6 +1106,7 @@ def test(
 
         # TODO: add more complex test case with a stellarator and/or MixedCoilSet
 
+    @pytest.mark.unit
     def test_quadratic_flux(self):
         """Test calculation of quadratic flux on the boundary."""
         t_field = ToroidalMagneticField(1, 1)
@@ -2274,6 +2277,7 @@ class TestComputeScalarResolution:
         # these require special logic
         BootstrapRedlConsistency,
         BoundaryError,
+        CoilArclengthVariance,
         CoilCurrentLength,
         CoilCurvature,
         CoilLength,
@@ -2643,12 +2647,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):
@@ -2682,6 +2687,7 @@ class TestObjectiveNaNGrad:
         BallooningStability,
         BootstrapRedlConsistency,
         BoundaryError,
+        CoilArclengthVariance,
         CoilLength,
         CoilCurrentLength,
         CoilCurvature,
@@ -2881,12 +2887,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):
@@ -2970,6 +2977,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__()
@@ -2998,6 +3006,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()