diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index f8e4bbeb8..1be68dfe5 100644 --- a/desc/compute/_curve.py +++ b/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=[], diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py index 00c959582..81e50c632 100644 --- a/desc/objectives/__init__.py +++ b/desc/objectives/__init__.py @@ -2,6 +2,7 @@ from ._bootstrap import BootstrapRedlConsistency from ._coils import ( + CoilArclengthVariance, CoilCurrentLength, CoilCurvature, CoilLength, diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index d8b68604e..4dc6569ec 100644 --- a/desc/objectives/_coils.py +++ b/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. diff --git a/desc/transform.py b/desc/transform.py index a5c1798ad..140b66356 100644 --- a/desc/transform.py +++ b/desc/transform.py @@ -43,6 +43,7 @@ class Transform(IOAble): """ _io_attrs_ = ["_grid", "_basis", "_derivatives", "_rcond", "_method"] + _static_attrs = ["_derivatives"] def __init__( self, diff --git a/docs/api.rst b/docs/api.rst index 85c201396..feb85514e 100644 --- a/docs/api.rst +++ b/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 diff --git a/docs/api_objectives.rst b/docs/api_objectives.rst index 2a7c443c2..9df305dc3 100644 --- a/docs/api_objectives.rst +++ b/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 diff --git a/tests/test_examples.py b/tests/test_examples.py index 93c700b36..dcd00d61b 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -34,6 +34,7 @@ AspectRatio, BallooningStability, BoundaryError, + CoilArclengthVariance, CoilCurvature, CoilLength, CoilSetMinDistance, @@ -1563,6 +1564,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 + ) # ,options={"initial_trust_ratio":1e-4}) + (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.""" diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index fccfaae38..7a67ae5bb 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -41,6 +41,7 @@ BootstrapRedlConsistency, BoundaryError, BScaleLength, + CoilArclengthVariance, CoilCurrentLength, CoilCurvature, CoilLength, @@ -2262,6 +2263,7 @@ class TestComputeScalarResolution: # these require special logic BootstrapRedlConsistency, BoundaryError, + CoilArclengthVariance, CoilCurrentLength, CoilCurvature, CoilLength, @@ -2629,12 +2631,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): @@ -2668,6 +2671,7 @@ class TestObjectiveNaNGrad: BallooningStability, BootstrapRedlConsistency, BoundaryError, + CoilArclengthVariance, CoilLength, CoilCurrentLength, CoilCurvature, @@ -2864,12 +2868,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):