diff --git a/CHANGELOG.md b/CHANGELOG.md index 139f8ba85d..d2259883e7 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/desc/compute/_omnigenity.py b/desc/compute/_omnigenity.py index 57f6ae85c5..03ae0e6419 100644 --- a/desc/compute/_omnigenity.py +++ b/desc/compute/_omnigenity.py @@ -754,6 +754,26 @@ def _f_C(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="f_C_normalized", + label="\\frac{[(M \\iota - N) (\\mathbf{B} \\times \\nabla \\psi)" + + " - (M G + N I) \\mathbf{B}] \\cdot \\nabla B}{B^3}", + units="~", + units_long="None", + description="Two-term quasisymmetry metric, normalized by the cube of the " + "local field strength", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["f_C", "|B|"], +) +def _f_C_normalized(params, transforms, profiles, data, **kwargs): + data["f_C_normalized"] = data["f_C"] / data["|B|"] ** 3 + return data + + @register_compute_fun( name="f_T", label="\\nabla \\psi \\times \\nabla B \\cdot \\nabla " @@ -776,6 +796,84 @@ def _f_T(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="f_T_normalized", + label="\\frac{R^2 \\nabla \\psi \\times \\nabla B \\cdot \\nabla " + + "(\\mathbf{B} \\cdot \\nabla B)}{B^4}", + units="~", + units_long="None", + description="Triple product quasisymmetry metric, normalized by the major " + "radius and the local field strength", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["f_T", "|B|", "R"], +) +def _f_T_normalized(params, transforms, profiles, data, **kwargs): + data["f_T_normalized"] = data["R"] ** 2 * data["f_T"] / data["|B|"] ** 4 + return data + + +@register_compute_fun( + name="f_B", + label="\\{B_{mn}^{\\mathrm{Boozer}}(\\rho) \\vert m/n \\neq M/N\\}", + units="T", + units_long="Tesla", + description="Symmetry breaking Boozer harmonics of magnetic field, " + "shape (num rho, num modes)", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="rtz", + data=["|B|_mn_B"], + resolution_requirement="tz", + grid_requirement={"is_meshgrid": True, "sym": False}, + matrix="ndarray: Transform matrix from the double-Fourier coefficients to the " + "double-angle coefficients, as returned by ``ptolemy_linear_transform``.", + idx="ndarray: Indices of the symmetry breaking modes, as returned by " + "``ptolemy_linear_transform``.", +) +def _f_B(params, transforms, profiles, data, **kwargs): + B_mn = data["|B|_mn_B"].reshape((transforms["grid"].num_rho, -1)) + # B_mn has shape (num modes, num rho) + B_mn = kwargs["matrix"] @ B_mn.T + data["f_B"] = B_mn[kwargs["idx"]].T + return data + + +@register_compute_fun( + name="f_B_normalized", + label="\\{B_{mn}^{\\mathrm{Boozer}}(\\rho) \\vert m/n \\neq M/N\\} / " + "(\\sum_{mn} B_{mn}^{\\mathrm{Boozer}}(\\rho)^2)^{1/2}", + units="~", + units_long="None", + description="Symmetry breaking Boozer harmonics of magnetic field, normalized " + "by the norm of all the harmonics on that surface, shape (num rho, num modes)", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="rtz", + data=["|B|_mn_B"], + resolution_requirement="tz", + grid_requirement={"is_meshgrid": True, "sym": False}, + matrix="ndarray: Transform matrix from the double-Fourier coefficients to the " + "double-angle coefficients, as returned by ``ptolemy_linear_transform``.", + idx="ndarray: Indices of the symmetry breaking modes, as returned by " + "``ptolemy_linear_transform``.", +) +def _f_B_normalized(params, transforms, profiles, data, **kwargs): + B_mn = data["|B|_mn_B"].reshape((transforms["grid"].num_rho, -1)) + # B_mn has shape (num modes, num rho), normalize each surface + B_mn = kwargs["matrix"] @ B_mn.T + norm = jnp.linalg.norm(B_mn, axis=0) + data["f_B_normalized"] = (B_mn[kwargs["idx"]] / norm).T + return data + + @register_compute_fun( name="eta", label="\\eta", diff --git a/desc/objectives/_omnigenity.py b/desc/objectives/_omnigenity.py index 88f59ebab6..da075e2df0 100644 --- a/desc/objectives/_omnigenity.py +++ b/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 @@ class QuasisymmetryBoozer(_Objective): 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 @@ class QuasisymmetryBoozer(_Objective): _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 @@ def __init__( 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 @@ def __init__( 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, @@ -123,7 +149,7 @@ def build(self, use_jit=True, verbose=1): "resolution for surface averages", ) - self._data_keys = ["|B|_mn_B"] + self._data_keys = ["f_B_normalized"] if self._scale_invariant else ["f_B"] timer = Timer() if verbose > 0: @@ -177,7 +203,7 @@ def compute(self, params, constants=None): 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) @@ -187,13 +213,13 @@ def compute(self, params, constants=None): params=params, transforms=constants["transforms"], profiles=constants["profiles"], + matrix=constants["matrix"], + idx=constants["idx"], surf_batch_size=self._surf_batch_size, ) - B_mn = data["|B|_mn_B"].reshape((constants["transforms"]["grid"].num_rho, -1)) - 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() + return data[self._data_keys[0]].flatten() @property def helicity(self): @@ -220,6 +246,14 @@ def helicity(self, helicity): 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 +263,11 @@ class QuasisymmetryTwoTerm(_Objective): 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 +278,7 @@ class QuasisymmetryTwoTerm(_Objective): _coordinates = "rtz" _units = "(T^3)" _print_value_fmt = "Quasi-symmetry two-term error: " + _static_attrs = _Objective._static_attrs + ["_scale_invariant"] def __init__( self, @@ -252,6 +292,7 @@ def __init__( deriv_mode="auto", grid=None, helicity=(1, 0), + scale_invariant=False, name="QS two-term", jac_chunk_size=None, ): @@ -259,6 +300,10 @@ def __init__( 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, @@ -307,7 +352,7 @@ def build(self, use_jit=True, verbose=1): ) self._dim_f = grid.num_nodes - self._data_keys = ["f_C"] + self._data_keys = ["f_C_normalized"] if self._scale_invariant else ["f_C"] timer = Timer() if verbose > 0: @@ -346,7 +391,8 @@ def compute(self, params, constants=None): 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 +404,7 @@ def compute(self, params, constants=None): profiles=constants["profiles"], helicity=constants["helicity"], ) - return data["f_C"] + return data[self._data_keys[0]] @property def helicity(self): @@ -384,6 +430,15 @@ def helicity(self, helicity): 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 @@ -391,6 +446,11 @@ class QuasisymmetryTripleProduct(_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 +461,7 @@ class QuasisymmetryTripleProduct(_Objective): _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 +474,17 @@ def __init__( 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, @@ -450,7 +516,7 @@ def build(self, use_jit=True, verbose=1): grid = self._grid self._dim_f = grid.num_nodes - self._data_keys = ["f_T"] + self._data_keys = ["f_T_normalized"] if self._scale_invariant else ["f_T"] timer = Timer() if verbose > 0: @@ -488,7 +554,8 @@ def compute(self, params, constants=None): 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 +566,7 @@ def compute(self, params, constants=None): transforms=constants["transforms"], profiles=constants["profiles"], ) - return data["f_T"] + return data[self._data_keys[0]] class Omnigenity(_Objective): diff --git a/desc/plotting.py b/desc/plotting.py index ffcc6e0998..7d549fa638 100644 --- a/desc/plotting.py +++ b/desc/plotting.py @@ -3440,16 +3440,18 @@ def plot_qs_error( # noqa: 16 fxn too complex grid = LinearGrid(M=2 * eq.M_grid, N=2 * eq.N_grid, NFP=eq.NFP, rho=rho) names = [] + booz = {} if fB: - names += ["|B|_mn_B"] + names += ["f_B_normalized"] transforms = get_transforms( - "|B|_mn_B", obj=eq, grid=grid, M_booz=M_booz, N_booz=N_booz + "f_B_normalized", obj=eq, grid=grid, M_booz=M_booz, N_booz=N_booz ) - matrix, modes, idx = ptolemy_linear_transform( + matrix, _, idx = ptolemy_linear_transform( transforms["B"].basis.modes, helicity=helicity, NFP=transforms["B"].basis.NFP, ) + booz = {"matrix": matrix, "idx": idx} if fC or fT: names += ["sqrt(g)"] if fC: @@ -3460,15 +3462,12 @@ def plot_qs_error( # noqa: 16 fxn too complex with warnings.catch_warnings(): warnings.simplefilter("ignore") data = eq.compute( - names, grid=grid, M_booz=M_booz, N_booz=N_booz, helicity=helicity + names, grid=grid, M_booz=M_booz, N_booz=N_booz, helicity=helicity, **booz ) if fB: - B_mn = data["|B|_mn_B"].reshape((len(rho), -1)) - B_mn = (matrix @ B_mn.T).T - f_B = np.sqrt(np.sum(B_mn[:, idx] ** 2, axis=-1)) / np.sqrt( - np.sum(B_mn**2, axis=-1) - ) + # norm of the symmetry breaking modes relative to all the modes + f_B = np.linalg.norm(data["f_B_normalized"], axis=-1) plot_data["f_B"] = f_B if fC: sqrtg = grid.meshgrid_reshape(data["sqrt(g)"], "rtz") diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index ab211c6599..da9c82f6be 100644 --- a/tests/test_objective_funs.py +++ b/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."""