Option to use proper normalization for QS objectives

CHANGELOG.md

@@ -4,6 +4,7 @@ Changelog
 New Features
 
 - Added warning for when ``deriv_mode="batched"`` is used in an ``ObjectiveFunction`` where one or more sub-objectives is using ``rev`` mode differentiation. Also adds more info about the derivative mode and Jacobian chunk sizes when building the objective with ``verbose>1``.
+- Adds ``scale_invariant`` argument to quasi-symmetry objectives (i.e. ``QuasisymmetryTwoTerm``, ``QuasisymmetryTripleProduct`` and ``QuasisymmetryBoozer``) that introduces the normalized alternatives for the objective functions with the actual evaluated magnetic field information instead of the precomputed constant normalization. For more details on these quantities, see [Basic Optimization tutorial](https://desc-docs.readthedocs.io/en/latest/notebooks/tutorials/basic_optimization.html).
 
 Performance Improvements
 

desc/objectives/_omnigenity.py

@@ -18,6 +18,17 @@
 class QuasisymmetryBoozer(_Objective):
     """Quasi-symmetry Boozer harmonics error.
 
+    Quasi-symmetry of helicity (M, N) requires the field strength in Boozer
+    coordinates to depend on the angles only through Mϑ_B - Nζ_B, so the residuals
+    are the symmetry breaking harmonics on each surface:
+
+    f_B = {B_mn(ρ) | m/n ≠ M/N}  (T)
+
+    With ``scale_invariant`` these are divided by the norm of all the harmonics on
+    that surface, so that ||f̂_B(ρ)|| ∈ [0, 1]:
+
+    f̂_B = f_B / (Σ_mn B_mn(ρ)²)^½
+
     Parameters
     ----------
     eq : Equilibrium
@@ -32,6 +43,12 @@
         Poloidal resolution of Boozer transformation. Default = 2 * eq.M.
     N_booz : int, optional
         Toroidal resolution of Boozer transformation. Default = 2 * eq.N.
+    scale_invariant : bool, optional
+        The scale_invariant version divides each surface's harmonics by the
+        norm of all the harmonics on that surface, making the output
+        dimensionless and invariant to the magnetic field strength. Then
+        the norm of the residuals on a single surface lies in [0, 1]. Default
+        is False, no normalization. See Basic Optimization tutorial for details.
     surf_batch_size: int
         Number of flux surfaces to compute simultaneously. Defaults to
         computing all flux surfaces simultaneously. Decrease to reduce
@@ -45,7 +62,11 @@
 
     _units = "(T)"
     _print_value_fmt = "Quasi-symmetry Boozer error: "
-    _static_attrs = _Objective._static_attrs + ["_helicity", "_surf_batch_size"]
+    _static_attrs = _Objective._static_attrs + [
+        "_helicity",
+        "_surf_batch_size",
+        "_scale_invariant",
+    ]
 
     def __init__(
         self,
@@ -61,6 +82,7 @@
         helicity=(1, 0),
         M_booz=None,
         N_booz=None,
+        scale_invariant=False,
         name="QS Boozer",
         jac_chunk_size=None,
         surf_batch_size=None,
@@ -72,6 +94,10 @@
         self.M_booz = M_booz
         self.N_booz = N_booz
         self._surf_batch_size = surf_batch_size
+        self._scale_invariant = scale_invariant
+        if scale_invariant:
+            normalize = False
+            self._units = "(dimensionless)"
         super().__init__(
             things=eq,
             target=target,
@@ -177,7 +203,7 @@
         Returns
         -------
         f : ndarray
-            Symmetry breaking harmonics of B (T).
+            Symmetry breaking harmonics of B (T), dimensionless for `scale_invariant`.
 
         """
         constants = self._get_deprecated_constants(constants)
@@ -193,7 +219,12 @@
         B_mn = constants["matrix"] @ B_mn.T
         # output order = (rho, mn).flatten(), ie all the surfaces concatenated
         # one after the other
-        return B_mn[constants["idx"]].T.flatten()
+        if not self._scale_invariant:
+            return B_mn[constants["idx"]].T.flatten()
+        else:
+            # B_mn has shape (num modes, num rho), normalize each surface
+            f = B_mn[constants["idx"]] / jnp.linalg.norm(B_mn, axis=0)
+            return f.T.flatten()
 
     @property
     def helicity(self):
@@ -220,6 +251,14 @@
 class QuasisymmetryTwoTerm(_Objective):
     """Quasi-symmetry two-term error.
 
+    With B = ||𝐁||, ι the rotational transform, and G, I the Boozer currents:
+
+    f_C = [(M ι - N) (𝐁 × ∇ψ) - (M G + N I) 𝐁] ⋅ ∇B  (T³)
+
+    With ``scale_invariant`` this is divided by the local field strength cubed:
+
+    f̂_C = f_C / B³
+
     Parameters
     ----------
     eq : Equilibrium
@@ -229,6 +268,11 @@
         Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``.
     helicity : tuple, optional
         Type of quasi-symmetry (M, N).
+    scale_invariant : bool, optional
+        The scale_invariant version divides by the cube of the local field
+        strength, making the output dimensionless and invariant to the magnetic
+        field strength. Default is False, no normalization. See Basic Optimization
+        tutorial for details.
 
     """
 
@@ -239,6 +283,7 @@
     _coordinates = "rtz"
     _units = "(T^3)"
     _print_value_fmt = "Quasi-symmetry two-term error: "
+    _static_attrs = _Objective._static_attrs + ["_scale_invariant"]
 
     def __init__(
         self,
@@ -252,13 +297,18 @@
         deriv_mode="auto",
         grid=None,
         helicity=(1, 0),
+        scale_invariant=False,
         name="QS two-term",
         jac_chunk_size=None,
     ):
         if target is None and bounds is None:
             target = 0
         self._grid = grid
         self.helicity = helicity
+        self._scale_invariant = scale_invariant
+        if scale_invariant:
+            normalize = False
+            self._units = "(dimensionless)"
         super().__init__(
             things=eq,
             target=target,
@@ -308,6 +358,8 @@
 
         self._dim_f = grid.num_nodes
         self._data_keys = ["f_C"]
+        if self._scale_invariant:
+            self._data_keys += ["|B|"]
 
         timer = Timer()
         if verbose > 0:
@@ -346,7 +398,8 @@
         Returns
         -------
         f : ndarray
-            Quasi-symmetry flux function error at each node (T^3).
+            Quasi-symmetry flux function error at each node (T^3), dimensionless
+            for `scale_invariant`.
 
         """
         constants = self._get_deprecated_constants(constants)
@@ -358,7 +411,10 @@
             profiles=constants["profiles"],
             helicity=constants["helicity"],
         )
-        return data["f_C"]
+        if not self._scale_invariant:
+            return data["f_C"]
+        else:
+            return data["f_C"] / data["|B|"] ** 3
 
     @property
     def helicity(self):
@@ -384,13 +440,27 @@
 class QuasisymmetryTripleProduct(_Objective):
     """Quasi-symmetry triple product error.
 
+    With B = ||𝐁||:
+
+    f_T = ∇ψ × ∇B ⋅ ∇(𝐁 ⋅ ∇B)  (T⁴/m²)
+
+    With ``scale_invariant`` this is made dimensionless with the major radius and
+    the local field strength:
+
+    f̂_T = R² f_T / B⁴
+
     Parameters
     ----------
     eq : Equilibrium
         Equilibrium 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)``.
+    scale_invariant : bool, optional
+        The scale_invariant version multiplies by R² and divides by the local B⁴,
+        making the output dimensionless and invariant to the magnetic field
+        strength. Default is False, no normalization. See Basic Optimization
+        tutorial for details.
 
     """
 
@@ -401,6 +471,7 @@
     _coordinates = "rtz"
     _units = "(T^4/m^2)"
     _print_value_fmt = "Quasi-symmetry error: "
+    _static_attrs = _Objective._static_attrs + ["_scale_invariant"]
 
     def __init__(
         self,
@@ -413,12 +484,17 @@
         loss_function=None,
         deriv_mode="auto",
         grid=None,
+        scale_invariant=False,
         name="QS triple product",
         jac_chunk_size=None,
     ):
         if target is None and bounds is None:
             target = 0
         self._grid = grid
+        self._scale_invariant = scale_invariant
+        if scale_invariant:
+            normalize = False
+            self._units = "(dimensionless)"
         super().__init__(
             things=eq,
             target=target,
@@ -451,6 +527,8 @@
 
         self._dim_f = grid.num_nodes
         self._data_keys = ["f_T"]
+        if self._scale_invariant:
+            self._data_keys += ["R", "|B|"]
 
         timer = Timer()
         if verbose > 0:
@@ -488,7 +566,8 @@
         Returns
         -------
         f : ndarray
-            Quasi-symmetry flux function error at each node (T^4/m^2).
+            Quasi-symmetry flux function error at each node (T^4/m^2),
+            dimensionless for `scale_invariant`.
 
         """
         constants = self._get_deprecated_constants(constants)
@@ -499,7 +578,10 @@
             transforms=constants["transforms"],
             profiles=constants["profiles"],
         )
-        return data["f_T"]
+        if not self._scale_invariant:
+            return data["f_T"]
+        else:
+            return data["R"] ** 2 * data["f_T"] / data["|B|"] ** 4
 
 
 class Omnigenity(_Objective):

tests/test_objective_funs.py

@@ -495,6 +495,77 @@ def test(eq):
         test(Equilibrium(iota=PowerSeriesProfile(0)))
         test(Equilibrium(current=PowerSeriesProfile(0)))
 
+    @pytest.mark.unit
+    def test_qs_hat_modes(self):
+        """Test the dimensionless "hat" modes of the QS objectives."""
+        eq = get("HELIOTRON")
+        eq.change_resolution(M=6, M_grid=12, N=2, N_grid=4)
+        helicity = (1, eq.NFP)
+        rho = np.array([0.6, 1.0])
+        grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, rho=rho)
+
+        booz = {"helicity": helicity, "M_booz": eq.M, "N_booz": eq.N}
+        objs = {
+            "fb": QuasisymmetryBoozer(eq=eq, grid=grid, scale_invariant=False, **booz),
+            "fb_hat": QuasisymmetryBoozer(
+                eq=eq, grid=grid, scale_invariant=True, **booz
+            ),
+            "fc_hat": QuasisymmetryTwoTerm(
+                eq=eq, grid=grid, helicity=helicity, scale_invariant=True
+            ),
+            "ft_hat": QuasisymmetryTripleProduct(
+                eq=eq, grid=grid, scale_invariant=True
+            ),
+        }
+        f = {}
+        for mode, obj in objs.items():
+            obj.build()
+            if mode.endswith("_hat"):
+                assert obj._units == "(dimensionless)"
+                assert obj.normalization == 1
+            assert obj.compute_scaled_error(*obj.xs(eq)).size == obj.dim_f
+            f[mode] = obj.compute_unscaled(*obj.xs(eq))
+            assert np.all(np.isfinite(f[mode]))
+
+        # on each surface this is the ratio of the norm of the symmetry breaking
+        # harmonics to the norm of all the harmonics, so it is between 0 and 1
+        ratio = np.linalg.norm(f["fb_hat"].reshape((rho.size, -1)), axis=-1)
+        assert np.all(ratio > 0) and np.all(ratio < 1)
+
+    @pytest.mark.unit
+    def test_qs_hat_modes_field_strength_invariance(self):
+        """Test that the QS "hat" modes are invariant to the field strength."""
+        # has an iota profile, so |B| is proportional to Psi
+        eq1 = get("HELIOTRON")
+        eq1.change_resolution(M=6, M_grid=12, N=2, N_grid=4)
+        eq2 = eq1.copy()
+        eq2.Psi = 2 * eq1.Psi
+        helicity = (1, eq1.NFP)
+        grid = LinearGrid(M=eq1.M_grid, N=eq1.N_grid, NFP=eq1.NFP, rho=np.array([0.6]))
+
+        def test(obj, mode, ratio, **kwargs):
+            obj1 = obj(
+                eq=eq1, grid=grid, scale_invariant=mode, normalize=False, **kwargs
+            )
+            obj2 = obj(
+                eq=eq2, grid=grid, scale_invariant=mode, normalize=False, **kwargs
+            )
+            obj1.build()
+            obj2.build()
+            f1 = ratio * obj1.compute_scaled_error(*obj1.xs(eq1))
+            f2 = obj2.compute_scaled_error(*obj2.xs(eq2))
+            atol = 1e-10 * np.max(np.abs(f1))
+            np.testing.assert_allclose(f2, f1, rtol=1e-10, atol=atol)
+
+        # scale variant quantities scale as |B|^n, others are invariant
+        booz = {"helicity": helicity, "M_booz": eq1.M, "N_booz": eq1.N}
+        test(QuasisymmetryBoozer, False, 2, **booz)
+        test(QuasisymmetryBoozer, True, 1, **booz)
+        test(QuasisymmetryTwoTerm, False, 2**3, helicity=helicity)
+        test(QuasisymmetryTwoTerm, True, 1, helicity=helicity)
+        test(QuasisymmetryTripleProduct, False, 2**4)
+        test(QuasisymmetryTripleProduct, True, 1)
+
     @pytest.mark.unit
     def test_isodynamicity(self):
         """Test calculation of isodynamicity metric."""