Merge branch 'master' into dp/order_rho_sq_NAE

CHANGELOG.md

@@ -4,18 +4,21 @@ Changelog
 New Feature
 
 - Adds a new profile class ``PowerProfile`` for raising profiles to a power.
+- 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``.
+- Adds the output quantities ``wb`` and ``wp`` to ``VMECIO.save``.
 
 Bug Fixes
 
-- Small bug fix to use the correct normalization length ``a`` in the BallooningStability objective
+- Small bug fix to use the correct normalization length ``a`` in the BallooningStability objective.
+- Fixed I/O bug when saving/loading ``_Profile`` classes that do not have a ``_params`` attribute.
 
 v0.13.0
 -------
 
 New Features
-
 - Adds ``from_input_file`` method to ``Equilibrium`` class to generate an ``Equilibrium`` object with boundary, profiles, resolution and flux specified in a given DESC or VMEC input file
 - Adds function ``solve_regularized_surface_current`` to ``desc.magnetic_fields`` module that implements the REGCOIL algorithm (Landreman, (2017)) for surface current normal field optimization
     * Can specify the tuple ``current_helicity=(M_coil, N_coil)`` to determine if resulting contours correspond to helical topology (both ``(M_coil, N_coil)`` not equal to 0), modular (``N_coil`` equal to 0 and ``M_coil`` nonzero) or windowpane/saddle (``M_coil`` and ``N_coil`` both zero)
@@ -298,7 +301,7 @@ optimization. Set to False by default.
 non-singular, non-degenerate) coordinate mappings for initial guesses. This is applied
 automatically when creating a new `Equilibrium` if the default initial guess of scaling
 the boundary surface produces self-intersecting surfaces. This can be disabled by
-passing `ensure_nested=False` when constructing the `Equilibrum`.
+passing `ensure_nested=False` when constructing the `Equilibrium`.
 - Adds `loss_function` argument to all `Objective`s for applying one of min/max/mean
 to objective function values (for targeting the average value of a profile, etc).
 - `Equilibrium.get_profile` now allows user to choose a profile type (power series, spline, etc)
@@ -450,7 +453,7 @@ Breaking changes
 - Renames ``theta_sfl`` to ``theta_PEST`` in compute functions to avoid confusion with
 other straight field line coordinate systems.
 - Makes plotting kwargs a bit more uniform. ``zeta``, ``nzeta``, ``nphi`` have all been
-superceded by ``phi`` which can be an integer for equally spaced angles or a float or
+superseded by ``phi`` which can be an integer for equally spaced angles or a float or
 array of float to specify angles manually.
 
 Bug fixes
@@ -520,7 +523,7 @@ the future all quantities should evaluate correctly at the magnetic axis. Note t
 evaluating quantities at the axis generally requires higher order derivatives and so
 can be much more expensive than evaluating at nonsingular points, so during optimization
 it is not recommended to include a grid point at the axis. Generally a small finite value
-such as ``rho = 1e-6`` will avoid the singuarlity with a negligible loss in accuracy for
+such as ``rho = 1e-6`` will avoid the singularity with a negligible loss in accuracy for
 analytic quantities.
 - Adds new optimizers ``fmin-auglag`` and ``lsq-auglag`` for performing constrained
 optimization using the augmented Lagrangian method. These generally perform much better
@@ -536,7 +539,7 @@ the existing methods ``compute_theta_coordinates`` and ``compute_flux_coordinate
 but allows mapping between arbitrary coordinates.
 - Adds calculation of $\nabla \mathbf{B}$ tensor and corresponding $L_{\nabla B}$ metric
 - Adds objective ``BScaleLength`` for penalizing strong magnetic field curvature.
-- Adds objective ``ObjectiveFromUser`` for wrapping an arbitary user defined function.
+- Adds objective ``ObjectiveFromUser`` for wrapping an arbitrary user defined function.
 - Adds utilities ``desc.grid.find_least_rational_surfaces`` and
 ``desc.grid.find_most_rational_surfaces`` for finding the least/most rational surfaces
 for a given rotational transform profile.
@@ -1072,7 +1075,7 @@ New Features:
 -   Updates `Equilibrium` to make creating them more straightforward.
     -   Instead of a dictionary of arrays and values, init method now
         takes individual arguments. These can either be objects of the
-        correct type (ie `Surface` objects for boundary condiitons,
+        correct type (ie `Surface` objects for boundary conditions,
         `Profile` for pressure and iota etc,) or ndarrays which will get
         parsed into objects of the correct type (for backwards
         compatibility)
@@ -1356,7 +1359,7 @@ Major changes:
     indexing, where only M+1 points were used instead of the correct
     2\*M+1
 -   Rotated concentric grids by 2pi/3M to avoid symmetry plane at
-    theta=0,pi. Previously, for stellarator symmetic cases, the nodes at
+    theta=0,pi. Previously, for stellarator symmetric cases, the nodes at
     theta=0 did not contribute to helical force balance.
 -   Added [L\_grid]{.title-ref} parameter to specify radial resolution
     of grid nodes directly and making the API more consistent.
@@ -1522,7 +1525,7 @@ saved, and objectives getting compiled more often than necessary
 
 Major Changes:
 
--   Changes to Equilibium/EquilibriaFamily:
+-   Changes to Equilibrium/EquilibriaFamily:
     -   general switching to using properties rather than direct
         attributes when referencing things (ie, `eq.foo`, not
         `eq._foo`). This allows getter methods to have safeguards if

desc/coils.py

@@ -1300,6 +1300,15 @@ def current(self, new):
         for coil, cur in zip(self.coils, new):
             coil.current = cur
 
+    def _all_currents(self, currents=None):
+        """Return an array of all the currents (including those in virtual coils)."""
+        if currents is None:
+            currents = self.current
+        currents = jnp.asarray(currents)
+        if self.sym:
+            currents = jnp.concatenate([currents, -1 * currents[::-1]])
+        return jnp.tile(currents, self.NFP)
+
     def _make_arraylike(self, x):
         if isinstance(x, dict):
             x = [x] * len(self)
@@ -2337,6 +2346,22 @@ def num_coils(self):
         """int: Number of coils."""
         return sum([c.num_coils for c in self])
 
+    def _all_currents(self, currents=None):
+        """Return an array of all the currents (including those in virtual coils)."""
+        if currents is None:
+            currents = jnp.array(flatten_list(self.current))
+        all_currents = []
+        i = 0
+        for coil in self.coils:
+            if isinstance(coil, CoilSet):
+                curr = currents[i : i + len(coil)]
+                all_currents += [coil._all_currents(curr)]
+                i += len(coil)
+            else:
+                all_currents += [jnp.atleast_1d(currents[i])]
+                i += 1
+        return jnp.concatenate(all_currents)
+
     def compute(
         self,
         names,

desc/compute/_field.py

@@ -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|"]) / (

desc/objectives/__init__.py

@@ -9,6 +9,7 @@
     CoilSetLinkingNumber,
     CoilSetMinDistance,
     CoilTorsion,
+    LinkingCurrentConsistency,
     PlasmaCoilSetMinDistance,
     QuadraticFlux,
     SurfaceCurrentRegularization,
@@ -31,6 +32,7 @@
     Elongation,
     GoodCoordinates,
     MeanCurvature,
+    MirrorRatio,
     PlasmaVesselDistance,
     PrincipalCurvature,
     Volume,

desc/objectives/_coils.py

@@ -2,6 +2,7 @@
 import warnings
 
 import numpy as np
+from scipy.constants import mu_0
 
 from desc.backend import (
     fori_loop,
@@ -15,7 +16,7 @@
 from desc.compute.utils import _compute as compute_fun
 from desc.grid import LinearGrid, _Grid
 from desc.integrals import compute_B_plasma
-from desc.utils import Timer, errorif, safenorm, warnif
+from desc.utils import Timer, broadcast_tree, errorif, safenorm, warnif
 
 from .normalization import compute_scaling_factors
 from .objective_funs import _Objective, collect_docs
@@ -1762,6 +1763,202 @@ def compute(self, params_1, params_2=None, constants=None):
         return Psi
 
 
+class LinkingCurrentConsistency(_Objective):
+    """Target the self-consistent poloidal linking current between the plasma and coils.
+
+    A self-consistent coil + plasma configuration must have the sum of the signed
+    currents in the coils that poloidally link the plasma equal to the total poloidal
+    current required to be linked by the plasma according to the loop integral of its
+    toroidal magnetic field, given by `G(rho=1)`. This objective computes the difference
+    between these two quantities, such that a value of zero means the coils create the
+    correct net poloidal current.
+
+    Assumes the coil topology does not change (ie the linking number with the plasma
+    is fixed).
+
+    Parameters
+    ----------
+    eq : Equilibrium
+        Equilibrium that will be optimized to satisfy the Objective.
+    coil : CoilSet
+        Coil(s) that are to be optimized.
+    grid : Grid, optional
+        Collocation grid containing the nodes to evaluate plasma current at.
+        Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym)``.
+    eq_fixed : bool
+        Whether the equilibrium is assumed fixed (should be true for stage 2, false
+        for single stage).
+    """
+
+    __doc__ = __doc__.rstrip() + collect_docs(
+        target_default="``target=0``.",
+        bounds_default="``target=0``.",
+    )
+
+    _scalar = True
+    _units = "(A)"
+    _print_value_fmt = "Linking current error: "
+
+    def __init__(
+        self,
+        eq,
+        coil,
+        *,
+        grid=None,
+        eq_fixed=False,
+        target=None,
+        bounds=None,
+        weight=1,
+        normalize=True,
+        normalize_target=True,
+        loss_function=None,
+        deriv_mode="auto",
+        jac_chunk_size=None,
+        name="linking current",
+    ):
+        if target is None and bounds is None:
+            target = 0
+        self._grid = grid
+        self._eq_fixed = eq_fixed
+        self._linear = eq_fixed
+        self._eq = eq
+        self._coil = coil
+
+        super().__init__(
+            things=[coil] if eq_fixed else [coil, eq],
+            target=target,
+            bounds=bounds,
+            weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
+            loss_function=loss_function,
+            deriv_mode=deriv_mode,
+            jac_chunk_size=jac_chunk_size,
+            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.
+
+        """
+        eq = self._eq
+        coil = self._coil
+        grid = self._grid or LinearGrid(
+            M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym
+        )
+        warnif(
+            not np.allclose(grid.nodes[:, 0], 1),
+            UserWarning,
+            "grid includes interior points, should be rho=1.",
+        )
+
+        self._dim_f = 1
+        self._data_keys = ["G"]
+
+        all_params = tree_map(lambda dim: np.arange(dim), coil.dimensions)
+        current_params = tree_map(lambda idx: {"current": idx}, True)
+        # indices of coil currents
+        self._indices = tree_leaves(broadcast_tree(current_params, all_params))
+        self._num_coils = coil.num_coils
+
+        profiles = get_profiles(self._data_keys, obj=eq, grid=grid)
+        transforms = get_transforms(self._data_keys, obj=eq, grid=grid)
+
+        # compute linking number of coils with plasma. To do this we add a fake "coil"
+        # along the magnetic axis and compute the linking number of that coilset
+        from desc.coils import FourierRZCoil, MixedCoilSet
+
+        axis_coil = FourierRZCoil(
+            1.0,
+            eq.axis.R_n,
+            eq.axis.Z_n,
+            eq.axis.R_basis.modes[:, 2],
+            eq.axis.Z_basis.modes[:, 2],
+            eq.axis.NFP,
+        )
+        dummy_coilset = MixedCoilSet(axis_coil, coil, check_intersection=False)
+        # linking number for coils with axis
+        link = np.round(dummy_coilset._compute_linking_number())[0, 1:]
+
+        self._constants = {
+            "quad_weights": 1.0,
+            "link": link,
+        }
+
+        if self._eq_fixed:
+            data = compute_fun(
+                "desc.equilibrium.equilibrium.Equilibrium",
+                self._data_keys,
+                params=eq.params_dict,
+                transforms=transforms,
+                profiles=profiles,
+            )
+            eq_linking_current = 2 * jnp.pi * data["G"][0] / mu_0
+            self._constants["eq_linking_current"] = eq_linking_current
+        else:
+            self._constants["profiles"] = profiles
+            self._constants["transforms"] = transforms
+
+        if self._normalize:
+            params = tree_leaves(
+                coil.params_dict, is_leaf=lambda x: isinstance(x, dict)
+            )
+            self._normalization = np.sum([np.abs(param["current"]) for param in params])
+
+        super().build(use_jit=use_jit, verbose=verbose)
+
+    def compute(self, coil_params, eq_params=None, constants=None):
+        """Compute linking current error.
+
+        Parameters
+        ----------
+        coil_params : dict
+            Dictionary of coilset degrees of freedom, eg ``CoilSet.params_dict``
+        eq_params : dict
+            Dictionary of equilibrium degrees of freedom, eg ``Equilibrium.params_dict``
+            Only required if eq_fixed=False.
+        constants : dict
+            Dictionary of constant data, eg transforms, profiles etc.
+            Defaults to self._constants.
+
+        Returns
+        -------
+        f : array of floats
+            Linking current error.
+
+        """
+        if constants is None:
+            constants = self.constants
+        if self._eq_fixed:
+            eq_linking_current = constants["eq_linking_current"]
+        else:
+            data = compute_fun(
+                "desc.equilibrium.equilibrium.Equilibrium",
+                self._data_keys,
+                params=eq_params,
+                transforms=constants["transforms"],
+                profiles=constants["profiles"],
+            )
+            eq_linking_current = 2 * jnp.pi * data["G"][0] / mu_0
+
+        coil_currents = jnp.concatenate(
+            [
+                jnp.atleast_1d(param[idx])
+                for param, idx in zip(tree_leaves(coil_params), self._indices)
+            ]
+        )
+        coil_currents = self.things[0]._all_currents(coil_currents)
+        coil_linking_current = jnp.sum(constants["link"] * coil_currents)
+        return eq_linking_current - coil_linking_current
+
+
 class CoilSetLinkingNumber(_Objective):
     """Prevents coils from becoming interlinked.