Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add mirror ratio objective #1366

Merged
merged 12 commits into from
Dec 13, 2024
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -7,6 +7,7 @@ New Feature
- Add ``desc.objectives.LinkingCurrentConsistency`` for ensuring that coils in a stage 2 or single stage optimization provide the required linking current for a given equilibrium.
- 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

12 changes: 12 additions & 0 deletions desc/compute/_field.py
Original file line number Diff line number Diff line change
@@ -3611,6 +3611,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|"])
@@ -3630,6 +3634,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|"])
@@ -3648,6 +3656,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|"]) / (
1 change: 1 addition & 0 deletions desc/objectives/__init__.py
Original file line number Diff line number Diff line change
@@ -32,6 +32,7 @@
Elongation,
GoodCoordinates,
MeanCurvature,
MirrorRatio,
PlasmaVesselDistance,
PrincipalCurvature,
Volume,
143 changes: 143 additions & 0 deletions desc/objectives/_geometry.py
Original file line number Diff line number Diff line change
@@ -1355,3 +1355,146 @@
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(

Check warning on line 1444 in desc/objectives/_geometry.py

Codecov / codecov/patch

desc/objectives/_geometry.py#L1444

Added line #L1444 was not covered by tests
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")

Check warning on line 1469 in desc/objectives/_geometry.py

Codecov / codecov/patch

desc/objectives/_geometry.py#L1469

Added line #L1469 was not covered by tests

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"])
1 change: 1 addition & 0 deletions docs/api.rst
Original file line number Diff line number Diff line change
@@ -212,6 +212,7 @@ Objective Functions
desc.objectives.MagneticWell
desc.objectives.MeanCurvature
desc.objectives.MercierStability
desc.objectives.MirrorRatio
desc.objectives.ObjectiveFromUser
desc.objectives.ObjectiveFunction
desc.objectives.Omnigenity
1 change: 1 addition & 0 deletions docs/api_objectives.rst
Original file line number Diff line number Diff line change
@@ -49,6 +49,7 @@ Geometry
desc.objectives.PrincipalCurvature
desc.objectives.PlasmaVesselDistance
desc.objectives.BScaleLength
desc.objectives.MirrorRatio
desc.objectives.GoodCoordinates


Binary file modified tests/inputs/master_compute_data_rpz.pkl
Binary file not shown.
48 changes: 48 additions & 0 deletions tests/test_objective_funs.py
Original file line number Diff line number Diff line change
@@ -65,6 +65,7 @@
MagneticWell,
MeanCurvature,
MercierStability,
MirrorRatio,
ObjectiveFromUser,
ObjectiveFunction,
Omnigenity,
@@ -1426,6 +1427,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_linking_current(self):
"""Test calculation of signed linking current from coils to plasma."""
Loading