Merge branch 'master' of https://github.com/PlasmaControl/DESC into rg/Mercier_rational

Author: rahulgaur104

CHANGELOG.md

@@ -1,10 +1,21 @@
 Changelog
 =========
 
+New Features
+- Adds a new utility function ``desc.compat.contract_equilibrium`` which takes in an ``Equilibrium`` object and an argument ``inner_rho``, and returns a new ``Equilibrium`` with original ``Equilibrium``'s ``inner_rho`` flux surface as its boundary.
+Optionally can also contract the profiles of the original ``Equilibrium`` so that the new ``Equilibrium``'s profiles match the original's in real space.
+- Adds second-order NAE constraints, accessible by passing ``order=2`` to ``desc.objectives.get_NAE_constraints``.
+- Adds automatically generated header file showing date the input file was created with `desc.vmec.VMECIO.write_vmec_input`
+
+
+v0.14.1
+-------
+
 Bug Fixes
 - Fixes bug in ``desc.vmec.VMECIO.write_vmec_input`` for current-constrained equilibria, where DESC was incorrectly writing the ``s**0`` mode, where VMEC actually assumes it is zero and starts at the  ``s**1`` (which is different than the usual convention VMEC uses for its current profile when it uses the current derivative, where it starts with the ``s**0`` mode).
 - Fixes error that occurs when using the default grid for ``SplineXYZCoil`` in an optimization.
 - Fixes bug in ``desc.coils.CoilSet.save_in_makegrid_format`` for ``CoilSet`` objects with ``sym=True`` or ``NFP>1``
+- Adds missing `&END` to the input file created from `desc.vmec.VMECIO.write_vmec_input`
 
 v0.14.0
 -------
@@ -37,6 +48,7 @@ for compatibility with other codes which expect such files from the Booz_Xform c
 - Adds a new objective ``desc.objectives.ExternalObjective`` for wrapping external codes with finite differences.
 - DESC/JAX version and device info is no longer printed by default, but can be accessed with the function `desc.backend.print_backend_info()`.
 
+
 Performance Improvements
 
 - A number of minor improvements to basis function evaluation and spectral transforms to improve speed. These will also enable future improvements for larger gains.

desc/compat.py

@@ -5,7 +5,8 @@
 import numpy as np
 
 from desc.grid import Grid, LinearGrid, QuadratureGrid
-from desc.utils import errorif
+from desc.profiles import FourierZernikeProfile, PowerSeriesProfile, SplineProfile
+from desc.utils import errorif, setdefault, warnif
 
 
 def ensure_positive_jacobian(eq):
@@ -330,3 +331,135 @@ def _get_new_coeffs(fun):
     eq_rotated.axis = eq_rotated.get_axis()
 
     return eq_rotated
+
+
+def contract_equilibrium(
+    eq, inner_rho, contract_profiles=True, profile_num_points=None, copy=True
+):
+    """Contract an equilibrium so that an inner flux surface is the new boundary.
+
+    Parameters
+    ----------
+    eq : Equilibrium
+        Equilibrium to contract.
+    inner_rho: float
+        rho value (<1) to contract the Equilibrium to.
+    contract_profiles :  bool
+        Whether or not to contract the profiles.
+        If True, the new profile's value at ``rho=1.0`` will be the same as the old
+        profile's value at ``rho=inner_rho``, i.e. in physical space, the new
+        profile is the same as the old profile. If the profile is a
+        ``PowerSeriesProfile`` or ``SplineProfile``, the same profile type will be
+        returned. If not one of these two classes, a ``SplineProfile`` will be returned
+        as other profile classes cannot be safely contracted.
+        If False, the new profile will have the same functional form
+        as the old profile, with no rescaling performed. This means the new equilibrium
+        has a physically different profile than the original equilibrium.
+    profile_num_points : int
+        Number of points to use when fitting or re-scaling the profiles. Defaults to
+        ``eq.L_grid``
+    copy : bool
+        Whether or not to return a copy or to modify the original equilibrium.
+
+    Returns
+    -------
+    eq_inner: New Equilibrium object, contracted from the old one such that
+        eq.pressure(rho=inner_rho) = eq_inner.pressure(rho=1), and
+        eq_inner LCFS = eq's rho=inner_rho surface.
+        Note that this will not be in force balance, and so must be re-solved.
+
+    """
+    errorif(
+        not (inner_rho < 1 and inner_rho > 0),
+        ValueError,
+        f"Surface must be in the range 0 < inner_rho < 1, instead got {inner_rho}.",
+    )
+    profile_num_points = setdefault(profile_num_points, eq.L_grid)
+
+    def scale_profile(profile, rho):
+        errorif(
+            isinstance(profile, FourierZernikeProfile),
+            ValueError,
+            "contract_equilibrium does not support FourierZernikeProfile",
+        )
+        if profile is None:
+            return profile
+        is_power_series = isinstance(profile, PowerSeriesProfile)
+        if contract_profiles and is_power_series:
+            # only PowerSeriesProfile both
+            # a) has a from_values
+            # b) can safely use that from_values to represent a
+            #    subset of itself.
+            x = np.linspace(0, 1, profile_num_points)
+            grid = LinearGrid(rho=x / rho)
+            y = profile.compute(grid)
+            return profile.from_values(x=x, y=y)
+        elif contract_profiles:
+            warnif(
+                not isinstance(profile, SplineProfile),
+                UserWarning,
+                f"{profile} is not a PowerSeriesProfile or SplineProfile,"
+                " so cannot safely contract using the same profile type."
+                "Falling back to fitting the values with a SplineProfile",
+            )
+            x = np.linspace(0, 1, profile_num_points)
+            grid = LinearGrid(rho=x / rho)
+            y = profile.compute(grid)
+            return SplineProfile(knots=x, values=y)
+        else:  # don't do any scaling of the profile
+            return profile
+
+    # create new profiles for contracted equilibrium
+    pressure = scale_profile(eq.pressure, inner_rho)
+    iota = scale_profile(eq.iota, inner_rho)
+    current = scale_profile(eq.current, inner_rho)
+    electron_density = scale_profile(eq.electron_density, inner_rho)
+    electron_temperature = scale_profile(eq.electron_temperature, inner_rho)
+    ion_temperature = scale_profile(eq.ion_temperature, inner_rho)
+    atomic_number = scale_profile(eq.atomic_number, inner_rho)
+    anisotropy = scale_profile(eq.anisotropy, inner_rho)
+
+    surf_inner = eq.get_surface_at(rho=inner_rho)
+    surf_inner.rho = 1.0
+    from .equilibrium import Equilibrium
+
+    eq_inner = Equilibrium(
+        surface=surf_inner,
+        pressure=pressure,
+        iota=iota,
+        current=current,
+        electron_density=electron_density,
+        electron_temperature=electron_temperature,
+        ion_temperature=ion_temperature,
+        atomic_number=atomic_number,
+        anisotropy=anisotropy,
+        Psi=(
+            eq.Psi * inner_rho**2
+        ),  # flux (in Webers) within the new last closed flux surface
+        NFP=eq.NFP,
+        L=eq.L,
+        M=eq.M,
+        N=eq.N,
+        L_grid=eq.L_grid,
+        M_grid=eq.M_grid,
+        N_grid=eq.N_grid,
+        sym=eq.sym,
+        ensure_nested=False,  # we fit the surfaces later so don't check now
+    )
+    inner_grid = LinearGrid(
+        rho=np.linspace(0, inner_rho, eq.L_grid * 2),
+        M=eq.M_grid,
+        N=eq.N_grid,
+        NFP=eq.NFP,
+        axis=True,
+    )
+    inner_data = eq.compute(["R", "Z", "lambda"], grid=inner_grid)
+    nodes = inner_grid.nodes
+    nodes[:, 0] = nodes[:, 0] / inner_rho
+    eq_inner.set_initial_guess(
+        nodes, inner_data["R"], inner_data["Z"], inner_data["lambda"]
+    )
+    if not copy:  # overwrite the original eq
+        eq.__dict__.update(eq_inner.__dict__)
+        return eq
+    return eq_inner

desc/equilibrium/equilibrium.py

@@ -2047,16 +2047,15 @@ def from_near_axis(
         A_Zw = A_Z * W[:, None]
         A_Lw = A_L * W[:, None]
 
+        rho = grid.nodes[grid.unique_rho_idx, 0]
         R_1D = np.zeros((grid.num_nodes,))
         Z_1D = np.zeros((grid.num_nodes,))
         L_1D = np.zeros((grid.num_nodes,))
+        phi_cyl_ax = np.linspace(0, 2 * np.pi / na_eq.nfp, na_eq.nphi, endpoint=False)
+        nu_B_ax = na_eq.nu_spline(phi_cyl_ax)
+        phi_B = phi_cyl_ax + nu_B_ax
         for rho_i in rho:
             R_2D, Z_2D, phi0_2D = na_eq.Frenet_to_cylindrical(r * rho_i, ntheta)
-            phi_cyl_ax = np.linspace(
-                0, 2 * np.pi / na_eq.nfp, na_eq.nphi, endpoint=False
-            )
-            nu_B_ax = na_eq.nu_spline(phi_cyl_ax)
-            phi_B = phi_cyl_ax + nu_B_ax
             nu_B = phi_B - phi0_2D
             idx = np.nonzero(grid.nodes[:, 0] == rho_i)[0]
             R_1D[idx] = R_2D.flatten(order="F")

desc/objectives/getters.py

@@ -1,6 +1,6 @@
 """Utilities for getting standard groups of objectives and constraints."""
 
-from desc.utils import flatten_list, get_all_instances, isposint, unique_list
+from desc.utils import errorif, flatten_list, get_all_instances, isposint, unique_list
 
 from ._equilibrium import Energy, ForceBalance, HelicalForceBalance, RadialForceBalance
 from .linear_objectives import (
@@ -29,7 +29,11 @@
     FixPsi,
     FixSheetCurrent,
 )
-from .nae_utils import calc_zeroth_order_lambda, make_RZ_cons_1st_order
+from .nae_utils import (
+    calc_zeroth_order_lambda,
+    make_RZ_cons_1st_order,
+    make_RZ_cons_2nd_order,
+)
 from .objective_funs import ObjectiveFunction
 
 _PROFILE_CONSTRAINTS = {
@@ -189,6 +193,15 @@ def get_NAE_constraints(
         A list of the linear constraints used in fixed-axis problems.
 
     """
+    if qsc_eq is not None:
+        errorif(
+            qsc_eq.lasym and fix_lambda is not False,
+            NotImplementedError,
+            "NAE Constrained equilibria with lambda constrained "
+            " do not yet work correctly with asymmetric equilibria, "
+            " as the NAE-prescribed lambda may not have the correct "
+            " gauge that DESC enforces (zero flux-surface average).",
+        )
     kwargs = {"eq": desc_eq, "normalize": normalize, "normalize_target": normalize}
     if not isinstance(fix_lambda, bool):
         fix_lambda = int(fix_lambda)
@@ -266,6 +279,7 @@ def _get_NAE_constraints(
         Whether to constrain lambda to match that of the NAE near-axis
         if an `int`, fixes lambda up to that order in rho {0,1}
         if `True`, fixes lambda up to the specified order given by `order`
+        (maximum of `order=1`)
     normalize : bool
         Whether to apply constraints in normalized units.
 
@@ -298,10 +312,19 @@ def _get_NAE_constraints(
 
     if order >= 1:  # first order constraints
         constraints += make_RZ_cons_1st_order(
-            qsc=qsc_eq, desc_eq=desc_eq, N=N, fix_lambda=fix_lambda and fix_lambda > 0
+            qsc=qsc_eq,
+            desc_eq=desc_eq,
+            N=N,
+            fix_lambda=fix_lambda and fix_lambda > 0,
+        )
+    if order == 2:  # 2nd order constraints
+        constraints += make_RZ_cons_2nd_order(
+            qsc=qsc_eq,
+            desc_eq=desc_eq,
+            N=N,
         )
-    if order >= 2:  # 2nd order constraints
-        raise NotImplementedError("NAE constraints only implemented up to O(rho) ")
+    if order > 2:
+        raise NotImplementedError("NAE constraints only implemented up to O(rho^2) ")
 
     return constraints