Merge branch 'master' into change_resolution_grid_attribute

.github/workflows/nbtests.yml

@@ -40,4 +40,4 @@ jobs:
           pwd
           lscpu
           export PYTHONPATH=$(pwd)
-          pytest -v --nbmake "./docs/notebooks" --nbmake-timeout=600 --ignore=./docs/notebooks/zernike_eval.ipynb --ignore=./docs/notebooks/DESC_Solve_from_pyQSC.ipynb
+          pytest -v --nbmake "./docs/notebooks" --nbmake-timeout=1000 --ignore=./docs/notebooks/zernike_eval.ipynb --ignore=./docs/notebooks/DESC_Solve_from_pyQSC.ipynb

desc/compute/utils.py

@@ -24,6 +24,8 @@
     "ne_l",
     "Ti_l",
     "Zeff_l",
+    "Ra_n",
+    "Za_n",
     "Rb_lmn",
     "Zb_lmn",
 )

desc/continuation.py

@@ -6,6 +6,7 @@
 import numpy as np
 from termcolor import colored
 
+from desc.compute import arg_order
 from desc.equilibrium import EquilibriaFamily, Equilibrium
 from desc.objectives import get_equilibrium_objective, get_fixed_boundary_constraints
 from desc.optimize import Optimizer
@@ -571,10 +572,10 @@ def solve_continuation_automatic(  # noqa: C901
         verbose,
         checkpoint_path,
     )
-
-    eq.R_lmn = eqfam[-1].R_lmn
-    eq.Z_lmn = eqfam[-1].Z_lmn
-    eq.L_lmn = eqfam[-1].L_lmn
+    for arg in arg_order:
+        val = np.asarray(getattr(eqfam[-1], arg))
+        if val.size:
+            setattr(eq, arg, val)
     eqfam[-1] = eq
     timer.stop("Total time")
     if verbose > 0:
@@ -726,9 +727,10 @@ def solve_continuation(  # noqa: C901
                 verbose=verbose,
                 copy=False,
             )
-            eqi.R_lmn = eqp.R_lmn
-            eqi.Z_lmn = eqp.Z_lmn
-            eqi.L_lmn = eqp.L_lmn
+            for arg in arg_order:
+                val = np.asarray(getattr(eqp, arg))
+                if val.size:
+                    setattr(eqi, arg, val)
             deltas = {}
             del eqp
 

desc/equilibrium/configuration.py

@@ -341,6 +341,7 @@ def __init__(  # noqa: C901 - FIXME: break this up into simpler pieces
                 )
         else:
             raise TypeError("Got unknown axis type {}".format(axis))
+        self._axis.change_resolution(self.N)
 
         # profiles
         self._pressure = None
@@ -576,6 +577,7 @@ def change_resolution(
         self.surface.change_resolution(
             self.L, self.M, self.N, NFP=self.NFP, sym=self.sym
         )
+        self.axis.change_resolution(self.N, NFP=self.NFP, sym=self.sym)
 
         self._R_lmn = copy_coeffs(self.R_lmn, old_modes_R, self.R_basis.modes)
         self._Z_lmn = copy_coeffs(self.Z_lmn, old_modes_Z, self.Z_basis.modes)
@@ -705,6 +707,40 @@ def get_profile(self, name, grid=None, **kwargs):
             x, grid.nodes[grid.unique_rho_idx, 0], grid=grid, name=name
         )
 
+    def get_axis(self):
+        """Return a representation for the magnetic axis.
+
+        Returns
+        -------
+        axis : FourierRZCurve
+            object representing the magnetic axis.
+        """
+        # value of Zernike polynomials at rho=0 for unique radial modes (+/-1)
+        sign_l = np.atleast_2d(((np.arange(0, self.L + 1, 2) / 2) % 2) * -2 + 1).T
+        # indices where m=0
+        idx0_R = np.where(self.R_basis.modes[:, 1] == 0)[0]
+        idx0_Z = np.where(self.Z_basis.modes[:, 1] == 0)[0]
+        # indices where l=0 & m=0
+        idx00_R = np.where((self.R_basis.modes[:, :2] == [0, 0]).all(axis=1))[0]
+        idx00_Z = np.where((self.Z_basis.modes[:, :2] == [0, 0]).all(axis=1))[0]
+        # this reshaping assumes the FourierZernike bases are sorted
+        R_n = np.sum(
+            sign_l * np.reshape(self.R_lmn[idx0_R], (-1, idx00_R.size), order="F"),
+            axis=0,
+        )
+        modes_R = self.R_basis.modes[idx00_R, 2]
+        if len(idx00_Z):
+            Z_n = np.sum(
+                sign_l * np.reshape(self.Z_lmn[idx0_Z], (-1, idx00_Z.size), order="F"),
+                axis=0,
+            )
+            modes_Z = self.Z_basis.modes[idx00_Z, 2]
+        else:  # catch cases such as axisymmetry with stellarator symmetry
+            Z_n = 0
+            modes_Z = 0
+        axis = FourierRZCurve(R_n, Z_n, modes_R, modes_Z, NFP=self.NFP, sym=self.sym)
+        return axis
+
     @property
     def surface(self):
         """Surface: Geometric surface defining boundary conditions."""
@@ -723,6 +759,24 @@ def surface(self, new):
                 f"surfaces should be of type Surface or a subclass, got {new}"
             )
 
+    @property
+    def axis(self):
+        """Curve: object representing the magnetic axis."""
+        return self._axis
+
+    @axis.setter
+    def axis(self, new):
+        if isinstance(new, FourierRZCurve):
+            assert (
+                self.sym == new.sym
+            ), "Surface and Equilibrium must have the same symmetry"
+            new.change_resolution(self.N)
+            self._axis = new
+        else:
+            raise TypeError(
+                f"axis should be of type FourierRZCurve or a subclass, got {new}"
+            )
+
     @property
     def spectral_indexing(self):
         """str: Type of indexing used for the spectral basis."""
@@ -845,41 +899,18 @@ def Ra_n(self):
         """ndarray: R coefficients for axis Fourier series."""
         return self.axis.R_n
 
+    @Ra_n.setter
+    def Ra_n(self, Ra_n):
+        self.axis.R_n = Ra_n
+
     @property
     def Za_n(self):
         """ndarray: Z coefficients for axis Fourier series."""
         return self.axis.Z_n
 
-    @property
-    def axis(self):
-        """Curve: object representing the magnetic axis."""
-        # value of Zernike polynomials at rho=0 for unique radial modes (+/-1)
-        sign_l = np.atleast_2d(((np.arange(0, self.L + 1, 2) / 2) % 2) * -2 + 1).T
-        # indices where m=0
-        idx0_R = np.where(self.R_basis.modes[:, 1] == 0)[0]
-        idx0_Z = np.where(self.Z_basis.modes[:, 1] == 0)[0]
-        # indices where l=0 & m=0
-        idx00_R = np.where((self.R_basis.modes[:, :2] == [0, 0]).all(axis=1))[0]
-        idx00_Z = np.where((self.Z_basis.modes[:, :2] == [0, 0]).all(axis=1))[0]
-        # this reshaping assumes the FourierZernike bases are sorted
-        R_n = np.sum(
-            sign_l * np.reshape(self.R_lmn[idx0_R], (-1, idx00_R.size), order="F"),
-            axis=0,
-        )
-        modes_R = self.R_basis.modes[idx00_R, 2]
-        if len(idx00_Z):
-            Z_n = np.sum(
-                sign_l * np.reshape(self.Z_lmn[idx0_Z], (-1, idx00_Z.size), order="F"),
-                axis=0,
-            )
-            modes_Z = self.Z_basis.modes[idx00_Z, 2]
-        else:  # catch cases such as axisymmetry with stellarator symmetry
-            Z_n = 0
-            modes_Z = 0
-        self._axis = FourierRZCurve(
-            R_n, Z_n, modes_R, modes_Z, NFP=self.NFP, sym=self.sym
-        )
-        return self._axis
+    @Za_n.setter
+    def Za_n(self, Za_n):
+        self.axis.Z_n = Za_n
 
     @property
     def pressure(self):

desc/equilibrium/equilibrium.py

@@ -17,6 +17,7 @@
     ForceBalance,
     ObjectiveFunction,
     get_equilibrium_objective,
+    get_fixed_axis_constraints,
     get_fixed_boundary_constraints,
 )
 from desc.optimize import Optimizer
@@ -863,11 +864,18 @@ def perturb(
         if objective is None:
             objective = get_equilibrium_objective(eq=self)
         if constraints is None:
-            constraints = get_fixed_boundary_constraints(
-                eq=self,
-                iota=self.iota is not None,
-                kinetic=self.electron_temperature is not None,
-            )
+            if "Ra_n" in deltas or "Za_n" in deltas:
+                constraints = get_fixed_axis_constraints(
+                    eq=self,
+                    iota=self.iota is not None,
+                    kinetic=self.electron_temperature is not None,
+                )
+            else:
+                constraints = get_fixed_boundary_constraints(
+                    eq=self,
+                    iota=self.iota is not None,
+                    kinetic=self.electron_temperature is not None,
+                )
 
         eq = perturb(
             self,

desc/equilibrium/initial_guess.py

@@ -149,6 +149,13 @@ def set_initial_guess(eq, *args):  # noqa: C901 - FIXME: simplify this
             eq.R_lmn = copy_coeffs(eq1.R_lmn, eq1.R_basis.modes, eq.R_basis.modes)
             eq.Z_lmn = copy_coeffs(eq1.Z_lmn, eq1.Z_basis.modes, eq.Z_basis.modes)
             eq.L_lmn = copy_coeffs(eq1.L_lmn, eq1.L_basis.modes, eq.L_basis.modes)
+            eq.Ra_n = copy_coeffs(
+                eq1.Ra_n, eq1.axis.R_basis.modes, eq.axis.R_basis.modes
+            )
+            eq.Za_n = copy_coeffs(
+                eq1.Za_n, eq1.axis.Z_basis.modes, eq.axis.Z_basis.modes
+            )
+
         elif isinstance(args[0], (str, os.PathLike)):
             # from file
             path = args[0]
@@ -184,6 +191,12 @@ def set_initial_guess(eq, *args):  # noqa: C901 - FIXME: simplify this
             eq.R_lmn = copy_coeffs(eq1.R_lmn, eq1.R_basis.modes, eq.R_basis.modes)
             eq.Z_lmn = copy_coeffs(eq1.Z_lmn, eq1.Z_basis.modes, eq.Z_basis.modes)
             eq.L_lmn = copy_coeffs(eq1.L_lmn, eq1.L_basis.modes, eq.L_basis.modes)
+            eq.Ra_n = copy_coeffs(
+                eq1.Ra_n, eq1.axis.R_basis.modes, eq.axis.R_basis.modes
+            )
+            eq.Za_n = copy_coeffs(
+                eq1.Za_n, eq1.axis.Z_basis.modes, eq.axis.Z_basis.modes
+            )
 
         elif nargs > 2:  # assume we got nodes and ndarray of points
             grid = args[0]