Merge branch 'master' into dp/coil-paper-figs

CHANGELOG.md

@@ -1,7 +1,15 @@
 Changelog
 =========
 
+New Feature
+
+- Adds a new profile class ``PowerProfile`` for raising profiles to a power.
+
+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)
@@ -13,14 +21,19 @@ New Features
     * use of both this and the ``QuadraticFlux`` objective allows for REGCOIL solutions to be obtained through the optimization framework, and combined with other objectives as well.
 - Changes local area weighting of Bn in QuadraticFlux objective to be the square root of the local area element (Note that any existing optimizations using this objective may need different weights to achieve the same result now.)
 - Adds a new tutorial showing how to use``REGCOIL`` features.
-
+- Adds an ``NFP`` attribute to ``ScalarPotentialField``, ``VectorPotentialField`` and ``DommaschkPotentialField``, to allow ``SplineMagneticField.from_field`` and ``MagneticField.save_mgrid`` to efficiently take advantage of the discrete toroidal symmetry of these fields, if present.
+- Adds ``SurfaceQuadraticFlux`` objective which minimizes the quadratic magnetic flux through a ``FourierRZToroidalSurface`` object, allowing for optimizing for Quadratic flux minimizing (QFM) surfaces.
+- Allows ``ToroidalFlux`` objective to accept ``FourierRZToroidalSurface`` so it can be used to specify the toroidal flux through a QFM surface.
+- Adds ``eq_fixed`` flag to ``ToroidalFlux`` to allow for the equilibrium/QFM surface to vary during optimization, useful for single-stage optimizations.
+- Adds tutorial notebook showcasing QFM surface capability.
+- Adds ``rotate_zeta`` function to ``desc.compat`` to rotate an ``Equilibrium`` around Z axis.
 
 Bug Fixes
 
-- Fixes bug that occurs when taking the gradient of ``root`` and ``root_scalar`` with newer versions of JAX (>=0.4.34) and unpins the JAX version
+- Fixes bug that occurs when taking the gradient of ``root`` and ``root_scalar`` with newer versions of JAX (>=0.4.34) and unpins the JAX version.
 - Changes ``FixLambdaGauge`` constraint to now enforce zero flux surface average for lambda, instead of enforcing lambda(rho,0,0)=0 as it was incorrectly doing before.
 - Fixes bug in ``softmin/softmax`` implementation.
-
+- Fixes bug that occured when using ``ProximalProjection`` with a scalar optimization algorithm.
 
 v0.12.3
 -------

desc/compat.py

@@ -261,3 +261,72 @@ def rescale(
     eq.surface = eq.get_surface_at(rho=1)
 
     return eq
+
+
+def rotate_zeta(eq, angle, copy=False):
+    """Rotate the equilibrium about the toroidal direction.
+
+    Parameters
+    ----------
+    eq : Equilibrium
+        Equilibrium to rotate.
+    angle : float
+        Angle to rotate the equilibrium in radians. The actual physical rotation
+        is by angle radians. Any rotation that is not a multiple of pi/NFP will
+        break the symmetry of a stellarator symmetric equilibrium.
+    copy : bool, optional
+        Whether to update the existing equilibrium or make a copy (Default).
+
+    Returns
+    -------
+    eq_rotated : Equilibrium
+        Equilibrium rotated about the toroidal direction
+    """
+    eq_rotated = eq.copy() if copy else eq
+    # We will apply the rotation in NFP domain
+    angle = angle * eq.NFP
+    # Check if the angle is a multiple of pi/NFP
+    kpi = np.isclose(angle % np.pi, 0, 1e-8, 1e-8) or np.isclose(
+        angle % np.pi, np.pi, 1e-8, 1e-8
+    )
+    if eq.sym and not kpi and eq.N != 0:
+        warnings.warn(
+            "Rotating a stellarator symmetric equilibrium by an angle "
+            "that is not a multiple of pi/NFP will break the symmetry. "
+            "Changing the symmetry to False to rotate the equilibrium."
+        )
+        eq_rotated.change_resolution(sym=0)
+
+    def _get_new_coeffs(fun):
+        if fun == "R":
+            f_lmn = np.array(eq_rotated.R_lmn)
+            basis = eq_rotated.R_basis
+        elif fun == "Z":
+            f_lmn = np.array(eq_rotated.Z_lmn)
+            basis = eq_rotated.Z_basis
+        elif fun == "L":
+            f_lmn = np.array(eq_rotated.L_lmn)
+            basis = eq_rotated.L_basis
+        else:
+            raise ValueError("fun must be 'R', 'Z' or 'L'")
+
+        new_coeffs = f_lmn.copy()
+        for i, (l, m, n) in enumerate(basis.modes):
+            id_sin = basis.get_idx(L=l, M=m, N=-n, error=False)
+            v_sin = np.sin(np.abs(n) * angle)
+            v_cos = np.cos(np.abs(n) * angle)
+            c_sin = f_lmn[id_sin] if isinstance(id_sin, int) else 0
+            if n >= 0:
+                new_coeffs[i] = f_lmn[i] * v_cos + c_sin * v_sin
+            elif n < 0:
+                new_coeffs[i] = f_lmn[i] * v_cos - c_sin * v_sin
+        return new_coeffs
+
+    eq_rotated.R_lmn = _get_new_coeffs(fun="R")
+    eq_rotated.Z_lmn = _get_new_coeffs(fun="Z")
+    eq_rotated.L_lmn = _get_new_coeffs(fun="L")
+
+    eq_rotated.surface = eq_rotated.get_surface_at(rho=1.0)
+    eq_rotated.axis = eq_rotated.get_axis()
+
+    return eq_rotated

desc/magnetic_fields/_core.py

@@ -1885,7 +1885,7 @@ def from_mgrid(cls, mgrid_file, extcur=None, method="cubic", extrap=False):
 
     @classmethod
     def from_field(
-        cls, field, R, phi, Z, params=None, method="cubic", extrap=False, NFP=1
+        cls, field, R, phi, Z, params=None, method="cubic", extrap=False, NFP=None
     ):
         """Create a splined magnetic field from another field for faster evaluation.
 
@@ -1903,14 +1903,16 @@ def from_field(
         extrap : bool
             whether to extrapolate splines beyond specified R,phi,Z
         NFP : int, optional
-            Number of toroidal field periods.
+            Number of toroidal field periods.  If not provided, will default to 1 or
+        the provided field's NFP, if it has that attribute.
 
         """
         R, phi, Z = map(np.asarray, (R, phi, Z))
         rr, pp, zz = np.meshgrid(R, phi, Z, indexing="ij")
         shp = rr.shape
         coords = np.array([rr.flatten(), pp.flatten(), zz.flatten()]).T
         BR, BP, BZ = field.compute_magnetic_field(coords, params, basis="rpz").T
+        NFP = getattr(field, "_NFP", 1)
         try:
             AR, AP, AZ = field.compute_magnetic_vector_potential(
                 coords, params, basis="rpz"
@@ -1948,12 +1950,22 @@ class ScalarPotentialField(_MagneticField):
         R,phi,Z are arrays of cylindrical coordinates.
     params : dict, optional
         default parameters to pass to potential function
+    NFP : int, optional
+        Whether the field has a discrete periodicity. This is only used when making
+        a ``SplineMagneticField`` from this field using its ``from_field`` method,
+        or when saving this field as an mgrid file using the ``save_mgrid`` method.
 
     """
 
-    def __init__(self, potential, params=None):
+    def __init__(self, potential, params=None, NFP=1):
         self._potential = potential
         self._params = params
+        self._NFP = NFP
+
+    @property
+    def NFP(self):
+        """int: Number of (toroidal) field periods."""
+        return self._NFP
 
     def compute_magnetic_field(
         self, coords, params=None, basis="rpz", source_grid=None, transforms=None
@@ -2042,12 +2054,22 @@ class VectorPotentialField(_MagneticField):
         R,phi,Z are arrays of cylindrical coordinates.
     params : dict, optional
         default parameters to pass to potential function
+    NFP : int, optional
+        Whether the field has a discrete periodicity. This is only used when making
+        a ``SplineMagneticField`` from this field using its ``from_field`` method,
+        or when saving this field as an mgrid file using the ``save_mgrid`` method.
 
     """
 
-    def __init__(self, potential, params=None):
+    def __init__(self, potential, params=None, NFP=1):
         self._potential = potential
         self._params = params
+        self._NFP = NFP
+
+    @property
+    def NFP(self):
+        """int: Number of (toroidal) field periods."""
+        return self._NFP
 
     def _compute_A_or_B(
         self,

desc/magnetic_fields/_dommaschk.py

@@ -38,6 +38,10 @@ class DommaschkPotentialField(ScalarPotentialField):
         d_m_l coefficients of V_m_l terms, which multiply the sin(m*phi)*N_m_l-1 terms
     B0: float
         scale strength of the magnetic field's 1/R portion
+    NFP : int, optional
+        Whether the field has a discrete periodicity. This is only used when making
+        a ``SplineMagneticField`` from this field using its ``from_field`` method,
+        or when saving this field as an mgrid file using the ``save_mgrid`` method.
 
     """
 
@@ -50,6 +54,7 @@ def __init__(
         c_arr=jnp.array([0.0]),
         d_arr=jnp.array([0.0]),
         B0=1.0,
+        NFP=1,
     ):
         ms = jnp.atleast_1d(jnp.asarray(ms))
         ls = jnp.atleast_1d(jnp.asarray(ls))
@@ -68,6 +73,11 @@ def __init__(
             jnp.isscalar(B0) or jnp.atleast_1d(B0).size == 1
         ), "B0 should be a scalar value!"
 
+        ms_over_NFP = ms / NFP
+        assert jnp.allclose(
+            ms_over_NFP, ms_over_NFP.astype(int)
+        ), "To enforce desired NFP, `ms` should be all integer multiples of NFP"
+
         params = {}
         params["ms"] = ms
         params["ls"] = ls
@@ -77,7 +87,7 @@ def __init__(
         params["d_arr"] = d_arr
         params["B0"] = B0
 
-        super().__init__(dommaschk_potential, params)
+        super().__init__(dommaschk_potential, params, NFP)
 
     @classmethod
     def fit_magnetic_field(  # noqa: C901

desc/objectives/__init__.py

@@ -12,6 +12,7 @@
     PlasmaCoilSetMinDistance,
     QuadraticFlux,
     SurfaceCurrentRegularization,
+    SurfaceQuadraticFlux,
     ToroidalFlux,
 )
 from ._equilibrium import (