Add mirror ratio objective

CHANGELOG.md

@@ -6,6 +6,7 @@ New Feature
 - Adds a new profile class ``PowerProfile`` for raising profiles to a power.
 - Adds an option ``scaled_termination`` (defaults to True) to all of the desc optimizers to measure the norms for ``xtol`` and ``gtol`` in the scaled norm provided by ``x_scale`` (which defaults to using an adaptive scaling based on the Jacobian or Hessian). This should make things more robust when optimizing parameters with widely different magnitudes. The old behavior can be recovered by passing ``options={"scaled_termination": False}``.
 - ``desc.objectives.Omnigenity`` is now vectorized and able to optimize multiple surfaces at the same time. Previously it was required to use a different objective for each surface.
+- Adds a new objective ``desc.objectives.MirrorRatio`` for targeting a particular mirror ratio on each flux surface, for either an ``Equilibrium`` or ``OmnigenousField``.
 
 Bug Fixes
 

desc/compute/_field.py

@@ -3394,6 +3394,10 @@ def _B_dot_gradB_z(params, transforms, profiles, data, **kwargs):
     coordinates="r",
     data=["|B|"],
     resolution_requirement="tz",
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.magnetic_fields._core.OmnigenousField",
+    ],
 )
 def _min_tz_modB(params, transforms, profiles, data, **kwargs):
     data["min_tz |B|"] = surface_min(transforms["grid"], data["|B|"])
@@ -3413,6 +3417,10 @@ def _min_tz_modB(params, transforms, profiles, data, **kwargs):
     coordinates="r",
     data=["|B|"],
     resolution_requirement="tz",
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.magnetic_fields._core.OmnigenousField",
+    ],
 )
 def _max_tz_modB(params, transforms, profiles, data, **kwargs):
     data["max_tz |B|"] = surface_max(transforms["grid"], data["|B|"])
@@ -3431,6 +3439,10 @@ def _max_tz_modB(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="r",
     data=["min_tz |B|", "max_tz |B|"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.magnetic_fields._core.OmnigenousField",
+    ],
 )
 def _mirror_ratio(params, transforms, profiles, data, **kwargs):
     data["mirror ratio"] = (data["max_tz |B|"] - data["min_tz |B|"]) / (

desc/objectives/__init__.py

@@ -31,6 +31,7 @@
     Elongation,
     GoodCoordinates,
     MeanCurvature,
+    MirrorRatio,
     PlasmaVesselDistance,
     PrincipalCurvature,
     Volume,

desc/objectives/_geometry.py

@@ -1355,3 +1355,146 @@ def compute(self, params, constants=None):
         f = data["g_rr"]
 
         return jnp.concatenate([g, constants["sigma"] * f])
+
+
+class MirrorRatio(_Objective):
+    """Target a particular value mirror ratio.
+
+    The mirror ratio is defined as:
+
+    (Bₘₐₓ - Bₘᵢₙ) / (Bₘₐₓ + Bₘᵢₙ)
+
+    Where Bₘₐₓ and Bₘᵢₙ are the maximum and minimum values of ||B|| on a given surface.
+    Returns one value for each surface in ``grid``.
+
+    Parameters
+    ----------
+    eq : Equilibrium or OmnigenousField
+        Equilibrium or OmnigenousField that will be optimized to satisfy the Objective.
+    grid : Grid, optional
+        Collocation grid containing the nodes to evaluate at. Defaults to
+        ``LinearGrid(M=eq.M_grid, N=eq.N_grid)`` for ``Equilibrium``
+        or ``LinearGrid(theta=2*eq.M_B, N=2*eq.N_x)`` for ``OmnigenousField``.
+
+    """
+
+    __doc__ = __doc__.rstrip() + collect_docs(
+        target_default="``target=0.2``.",
+        bounds_default="``target=0.2``.",
+    )
+
+    _coordinates = "r"
+    _units = "(dimensionless)"
+    _print_value_fmt = "Mirror ratio: "
+
+    def __init__(
+        self,
+        eq,
+        *,
+        grid=None,
+        target=None,
+        bounds=None,
+        weight=1,
+        normalize=True,
+        normalize_target=True,
+        loss_function=None,
+        deriv_mode="auto",
+        name="mirror ratio",
+        jac_chunk_size=None,
+    ):
+        if target is None and bounds is None:
+            target = 0.2
+        self._grid = grid
+        super().__init__(
+            things=eq,
+            target=target,
+            bounds=bounds,
+            weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
+            loss_function=loss_function,
+            deriv_mode=deriv_mode,
+            name=name,
+            jac_chunk_size=jac_chunk_size,
+        )
+
+    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.
+
+        """
+        eq = self.things[0]
+        from desc.equilibrium import Equilibrium
+        from desc.magnetic_fields import OmnigenousField
+
+        if self._grid is None and isinstance(eq, Equilibrium):
+            grid = LinearGrid(
+                M=eq.M_grid,
+                N=eq.N_grid,
+                NFP=eq.NFP,
+                sym=eq.sym,
+            )
+        elif self._grid is None and isinstance(eq, OmnigenousField):
+            grid = LinearGrid(
+                theta=2 * eq.M_B,
+                N=2 * eq.N_x,
+                NFP=eq.NFP,
+            )
+        else:
+            grid = self._grid
+
+        self._dim_f = grid.num_rho
+        self._data_keys = ["mirror ratio"]
+
+        timer = Timer()
+        if verbose > 0:
+            print("Precomputing transforms")
+        timer.start("Precomputing transforms")
+
+        profiles = get_profiles(self._data_keys, obj=eq, grid=grid)
+        transforms = get_transforms(self._data_keys, obj=eq, grid=grid)
+        self._constants = {
+            "transforms": transforms,
+            "profiles": profiles,
+        }
+
+        timer.stop("Precomputing transforms")
+        if verbose > 1:
+            timer.disp("Precomputing transforms")
+
+        super().build(use_jit=use_jit, verbose=verbose)
+
+    def compute(self, params, constants=None):
+        """Compute mirror ratio.
+
+        Parameters
+        ----------
+        params : dict
+            Dictionary of equilibrium or field degrees of freedom,
+            eg Equilibrium.params_dict
+        constants : dict
+            Dictionary of constant data, eg transforms, profiles etc. Defaults to
+            self.constants
+
+        Returns
+        -------
+        M : ndarray
+            Mirror ratio on each surface.
+
+        """
+        if constants is None:
+            constants = self.constants
+        data = compute_fun(
+            self.things[0],
+            self._data_keys,
+            params=params,
+            transforms=constants["transforms"],
+            profiles=constants["profiles"],
+        )
+        return constants["transforms"]["grid"].compress(data["mirror ratio"])

docs/api.rst

@@ -211,6 +211,7 @@ Objective Functions
     desc.objectives.MagneticWell
     desc.objectives.MeanCurvature
     desc.objectives.MercierStability
+    desc.objectives.MirrorRatio
     desc.objectives.ObjectiveFromUser
     desc.objectives.ObjectiveFunction
     desc.objectives.Omnigenity

docs/api_objectives.rst

@@ -49,6 +49,7 @@ Geometry
     desc.objectives.PrincipalCurvature
     desc.objectives.PlasmaVesselDistance
     desc.objectives.BScaleLength
+    desc.objectives.MirrorRatio
     desc.objectives.GoodCoordinates
 
 

tests/test_objective_funs.py

@@ -62,6 +62,7 @@
     MagneticWell,
     MeanCurvature,
     MercierStability,
+    MirrorRatio,
     ObjectiveFromUser,
     ObjectiveFunction,
     Omnigenity,
@@ -1423,6 +1424,53 @@ def test_signed_plasma_vessel_distance(self):
             )
             obj.build()
 
+    @pytest.mark.unit
+    def test_mirror_ratio_equilibrium(self):
+        """Test mirror ratio objective for Equilibrium."""
+        # axisymmetry, no iota, so B ~ B0/R
+        eq = Equilibrium(L=8, M=8)
+        eq.solve()
+        # R0 = 10, a=1, so Bmax = B0/9, Bmin = B0/11
+        mirror_ratio = (1 / 9 - 1 / 11) / (1 / 9 + 1 / 11)
+        obj = MirrorRatio(eq)
+        obj.build()
+        f = obj.compute(eq.params_dict)
+        # not perfect agreement bc eq is low res, so B isnt exactly B0/R
+        np.testing.assert_allclose(f, mirror_ratio, rtol=3e-3)
+
+    @pytest.mark.unit
+    def test_mirror_ratio_omni_field(self):
+        """Test mirror ratio objective for OmnigenousField."""
+        field = OmnigenousField(
+            L_B=1,
+            M_B=3,
+            L_x=1,
+            M_x=1,
+            N_x=1,
+            NFP=1,
+            helicity=(0, 1),
+            B_lm=np.array(
+                [
+                    # f(r) = B0 + B1*(2r-1)
+                    # f(0) = [0.8, 1.0, 1.2]
+                    # f(1) = [1.0, 1.0, 1.0]
+                    [0.9, 1.0, 1.1],  # B0
+                    [0.1, 0.0, -0.1],  # B1
+                ]
+            ).flatten(),
+        )
+
+        mirror_ratio_axis = (1.2 - 0.8) / (1.2 + 0.8)
+        mirror_ratio_edge = 0.0
+        grid = LinearGrid(L=5, theta=6, N=2)
+        rho = grid.nodes[grid.unique_rho_idx, 0]
+        obj = MirrorRatio(field, grid=grid)
+        obj.build()
+        f = obj.compute(field.params_dict)
+        np.testing.assert_allclose(
+            f, mirror_ratio_axis * (1 - rho) + mirror_ratio_edge * rho
+        )
+
     @pytest.mark.unit
     def test_omnigenity_multiple_surfaces(self):
         """Test omnigenity transform vectorized over multiple surfaces."""