Merge pull request #894 from PlasmaControl/rc/surface_current

Clean up API for free boundary with surface current

CHANGELOG.md

@@ -1,7 +1,52 @@
 Changelog
 =========
 
-- Adds a bounding box to the `field_line_integrate` defined by `bounds_R` and `bounds_Z` keyword arguments, which form a hollow cylindrical bounding box. If the field line trajectory exits these bounds, the RHS will be multiplied by an exponentially decaying function of the distance to the box to stop the trajectory and prevent tracing the field line out to infinity, which is both costly and unnecessary when making a Poincare plot, the principle purpose of the function.
+
+New Features
+
+- Adds new objectives for free boundary equilibria: ``BoundaryError`` and
+``VacuumBoundaryError``, along with a new tutorial notebook demonstrating their usage.
+- Objectives ``Volume``, ``AspectRatio``, ``Elongation`` now work for
+``FourierRZToroidalSurface`` objects as well as ``Equilibrium``.
+- ``MagneticField`` objects now have a method ``save_mgrid`` for saving field data
+in the MAKEGRID format for use with other codes.
+- ``SplineMagneticField.from_mgrid`` now defaults to using ``extcur`` from the mgrid file.
+- When converting a near axis solution from QSC/QIC to a desc ``Equilibrium``, the
+least squares fit is now weighted inversely with the distance from the axis to improve
+the accuracy for low aspect ratio.
+- Adds a bounding box to the `field_line_integrate` defined by `bounds_R` and `bounds_Z`
+keyword arguments, which form a hollow cylindrical bounding box. If the field line 
+trajectory exits these bounds, the RHS will be multiplied by an exponentially decaying 
+function of the distance to the box to stop the trajectory and prevent tracing the field 
+line out to infinity, which is both costly and unnecessary when making a Poincare plot, 
+the principle purpose of the function.
+
+
+Speed Improvements
+
+- ``CoilSet`` is now more efficient when stellarator or field period symmetry is used.
+- Improves the efficiency of ``proximal`` optimizers by reducing the number of objective
+derivative evaluations. Optimization steps should now be 2-5x faster.
+- Improved performance of Zernike polynomial evaluation.
+- Adds a bounding box to the `field_line_integrate` defined by `bounds_R` and `bounds_Z`
+keyword arguments, which form a hollow cylindrical bounding box. If the field line
+trajectory exits these bounds, the RHS will be multiplied by an exponentially decaying
+function of the distance to the box to stop the trajectory and prevent tracing the
+field line out to infinity, which is both costly and unnecessary when making a Poincare
+plot, the principle purpose of the function.
+
+
+Bug Fixes
+
+- Fix bug causing NaN in ``ForceBalance`` objective when the grid contained nodes at
+the magnetic axis.
+- When saving VMEC output, ``buco`` and ``bvco`` are now correctly saved on the half
+mesh. Previously they were saved on the full mesh.
+- Fixed a bug where hdf5 files were not properly closed after reading.
+- Fixed bugs relating to `Curve` objects not being optimizable.
+- Fixed incorrect rotation matrix for `FourierPlanarCurve`.
+- Fixed bug where ``plot_boundaries`` with a single ``phi`` value would return an
+empty plot.
 - Renames the method for comparing equivalence between DESC objects from `eq` to `equiv`
 to avoid confusion with the common shorthand for `Equilibrium`.
 

desc/equilibrium/equilibrium.py

@@ -19,7 +19,6 @@
 from desc.geometry import (
     FourierRZCurve,
     FourierRZToroidalSurface,
-    Surface,
     ZernikeRZToroidalSection,
 )
 from desc.grid import LinearGrid, QuadratureGrid, _Grid
@@ -424,6 +423,9 @@ def _sort_args(self, args):
             "Za_n",
             "Rb_lmn",
             "Zb_lmn",
+            "I",
+            "G",
+            "Phi_mn",
         )
         assert sorted(args) == sorted(arg_order)
         return [arg for arg in arg_order if arg in args]
@@ -1132,14 +1134,12 @@ def surface(self):
     @surface.setter
     def surface(self, new):
         assert isinstance(
-            new, Surface
-        ), f"surfaces should be of type Surface or a subclass, got {new}"
+            new, FourierRZToroidalSurface
+        ), f"surface should be of type FourierRZToroidalSurface or subclass, got {new}"
         assert (
             self.sym == new.sym
         ), "Surface and Equilibrium must have the same symmetry"
-        assert self.NFP == getattr(
-            new, "NFP", self.NFP
-        ), "Surface and Equilibrium must have the same NFP"
+        assert self.NFP == new.NFP, "Surface and Equilibrium must have the same NFP"
         new.change_resolution(self.L, self.M, self.N)
         self._surface = new
 
@@ -1297,6 +1297,51 @@ def Zb_lmn(self):
     def Zb_lmn(self, Zb_lmn):
         self.surface.Z_lmn = Zb_lmn
 
+    @optimizable_parameter
+    @property
+    def I(self):  # noqa: E743
+        """float: Net toroidal current on the sheet current at the LCFS."""
+        return self.surface.I if hasattr(self.surface, "I") else np.empty(0)
+
+    @I.setter
+    def I(self, new):  # noqa: E743
+        errorif(
+            not hasattr(self.surface, "I"),
+            ValueError,
+            "Attempt to set I on an equilibrium without sheet current",
+        )
+        self.surface.I = new
+
+    @optimizable_parameter
+    @property
+    def G(self):
+        """float: Net poloidal current on the sheet current at the LCFS."""
+        return self.surface.G if hasattr(self.surface, "G") else np.empty(0)
+
+    @G.setter
+    def G(self, new):
+        errorif(
+            not hasattr(self.surface, "G"),
+            ValueError,
+            "Attempt to set G on an equilibrium without sheet current",
+        )
+        self.surface.G = new
+
+    @optimizable_parameter
+    @property
+    def Phi_mn(self):
+        """ndarray: coeffs of single-valued part of surface current potential."""
+        return self.surface.Phi_mn if hasattr(self.surface, "Phi_mn") else np.empty(0)
+
+    @Phi_mn.setter
+    def Phi_mn(self, new):
+        errorif(
+            not hasattr(self.surface, "Phi_mn"),
+            ValueError,
+            "Attempt to set Phi_mn on an equilibrium without sheet current",
+        )
+        self.surface.Phi_mn = new
+
     @optimizable_parameter
     @property
     def Ra_n(self):

desc/objectives/_free_boundary.py

@@ -321,7 +321,8 @@ class BoundaryError(_Objective):
     current potential on the LCFS. All residuals are weighted by the local area
     element ||𝐞_θ × 𝐞_ζ|| Δθ Δζ
 
-    The third equation is only included if a sheet current is supplied, otherwise it
+    The third equation is only included if a sheet current is supplied by making
+    the ``equilibrium.surface`` object a FourierCurrentPotentialField, otherwise it
     is trivially satisfied. If it is known that the external field accurately reproduces
     the target equilibrium with low normal field error and pressure at the edge is zero,
     then the sheet current will generally be negligible and can be omitted to save
@@ -337,10 +338,6 @@ class BoundaryError(_Objective):
         Equilibrium that will be optimized to satisfy the Objective.
     ext_field : MagneticField
         External field produced by coils.
-    sheet_current : FourierCurrentPotentialField
-        Sheet current on the last closed flux surface. Not required for vacuum fields,
-        but generally needed to correctly solve at finite beta/current. Geometry will
-        be automatically constrained to be the same as the equilibrium LCFS.
     target : float, ndarray, optional
         Target value(s) of the objective. Only used if bounds is None.
         len(target) must be equal to Objective.dim_f
@@ -381,6 +378,22 @@ class BoundaryError(_Objective):
     name : str
         Name of the objective function.
 
+
+    Examples
+    --------
+    Assigning a surface current to the equilibrium:
+
+    .. code-block:: python
+
+        from desc.magnetic_fields import FourierCurrentPotentialField
+        # turn the regular FourierRZToroidalSurface into a current potential on the
+        # last closed flux surface
+        eq.surface = FourierCurrentPotentialField.from_suface(eq.surface,
+                                                              M_Phi=eq.M,
+                                                              N_Phi=eq.N,
+                                                              )
+        objective = BoundaryError(eq, ext_field)
+
     """
 
     _scalar = False
@@ -394,7 +407,6 @@ def __init__(
         self,
         eq,
         ext_field,
-        sheet_current=None,
         target=None,
         bounds=None,
         weight=1,
@@ -419,11 +431,8 @@ def __init__(
         self._ext_field = ext_field
         self._field_grid = field_grid
         self._loop = loop
-        self._sheet_current = None
+        self._sheet_current = hasattr(eq.surface, "Phi_mn")
         things = [eq]
-        if sheet_current:
-            things += [sheet_current]
-            self._sheet_current = True
 
         super().__init__(
             things=things,
@@ -550,19 +559,12 @@ def build(self, use_jit=True, verbose=1):
         }
 
         if self._sheet_current:
-            sheet_current = self.things[1]
-            assert (
-                (sheet_current.M == eq.surface.M)
-                and (sheet_current.N == eq.surface.N)
-                and (sheet_current.NFP == eq.surface.NFP)
-                and (sheet_current.sym == eq.surface.sym)
-            ), "sheet current must have same resolution as eq.surface"
             self._sheet_data_keys = ["K"]
             sheet_eval_transforms = get_transforms(
-                self._sheet_data_keys, obj=sheet_current, grid=eval_grid
+                self._sheet_data_keys, obj=eq.surface, grid=eval_grid
             )
             sheet_source_transforms = get_transforms(
-                self._sheet_data_keys, obj=sheet_current, grid=source_grid
+                self._sheet_data_keys, obj=eq.surface, grid=source_grid
             )
             self._constants["sheet_eval_transforms"] = sheet_eval_transforms
             self._constants["sheet_source_transforms"] = sheet_source_transforms
@@ -590,15 +592,13 @@ def build(self, use_jit=True, verbose=1):
 
         super().build(use_jit=use_jit, verbose=verbose)
 
-    def compute(self, eq_params, sheet_params=None, constants=None):
+    def compute(self, eq_params, constants=None):
         """Compute boundary force error.
 
         Parameters
         ----------
         eq_params : dict
             Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        sheet_params : dict
-            Dictionary of sheet current degrees of freedom.
         constants : dict
             Dictionary of constant data, eg transforms, profiles etc. Defaults to
             self.constants
@@ -628,10 +628,13 @@ def compute(self, eq_params, sheet_params=None, constants=None):
             profiles=constants["eval_profiles"],
         )
         if self._sheet_current:
-            assert sheet_params is not None
-            # enforce that they're the same surface
-            sheet_params["R_lmn"] = eq_params["Rb_lmn"]
-            sheet_params["Z_lmn"] = eq_params["Zb_lmn"]
+            sheet_params = {
+                "R_lmn": eq_params["Rb_lmn"],
+                "Z_lmn": eq_params["Zb_lmn"],
+                "I": eq_params["I"],
+                "G": eq_params["G"],
+                "Phi_mn": eq_params["Phi_mn"],
+            }
             sheet_source_data = compute_fun(
                 "desc.magnetic_fields.FourierCurrentPotentialField",
                 self._sheet_data_keys,