Clean up API for free boundary with surface current

CHANGELOG.md

@@ -1,6 +1,40 @@
 Changelog
 =========
 
+
+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.
+
+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.
+
+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.
+
+
 v0.10.4
 -------
 

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 @@
             "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 @@
     @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, Zb_lmn):
         self.surface.Z_lmn = Zb_lmn
 
+    @optimizable_parameter
+    @property
+    def I(self):  # noqa: E743
+        """float: Net current linking the plasma surface toroidally."""
+        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 current linking the plasma surface poloidally."""
+        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 @@
     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 @@
         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 @@
     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 @@
         self,
         eq,
         ext_field,
-        sheet_current=None,
         target=None,
         bounds=None,
         weight=1,
@@ -419,11 +431,8 @@
         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 @@
         }
 
         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 @@
 
         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 @@
             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,