Add JAX implementation of REGCOIL algorithm and Add Ability to Disc…

…retize Current Potentials into Coilsets (#579)

- 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)
* ``M_coil`` is the number of poloidal transits a coil makes before
returning to itself, while ``N_coil`` is the number of toroidal transits
a coil makes before returning to itself (this is sort of like the QS
``helicity``)
* Can also specify `regularization_type` which can be `simple` or
`regcoil`, the difference explained
[here](https://www.overleaf.com/2924393292wskpdrwbtkhh#2ddf44), tldr is
simple is not guaranteed to monotonically decrease $\chi^2_K$ as
regularization increases, but gives the same qualitative behavior and is
less expensive to run.
- Adds method ``To_CoilSet`` to ``FourierCurrentPotentialField`` which
implements a coil cutting algorithm to discretize the surface current
into coils
    * works for both modular and helical coils
* uses `skimage.measure.find_contours` instead of `matplotlib`'s
deprecated contour finding algorithm, so this PR adds a new dependency
in `skimage`.
- Adds a new objective ``SurfaceCurrentRegularization`` (which minimizes
``w*|K|``, the regularization term from surface current in the REGCOIL
algorithm, with `w` being the objective weight which act as the sqrt of
the regularization parameter)
* 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 $\sqrt{|\mathbf{e}_{\theta} \times \mathbf{e}_{\zeta}}|$
- Adds a new tutorial showing how to use``REGCOIL`` features.

Resolves #578 
Resolves #1351 
Resolves #1172

CHANGELOG.md

@@ -2,7 +2,17 @@ Changelog
 =========
 
 New Features
-- Add ``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 ``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)
+    * ``M_coil`` is the number of poloidal transits a coil makes before returning to itself, while ``N_coil`` is the number of toroidal transits a coil makes before returning to itself (this is sort of like the QS ``helicity``)
+    * if multiple values of the regularization parameter are input, will return a family of surface current fields (as a list) corresponding to the solution at each regularization value
+- Adds method ``to_CoilSet`` to ``FourierCurrentPotentialField`` which implements a coil cutting algorithm to discretize the surface current into coils
+    * works for both modular and helical coils
+- Adds a new objective ``SurfaceCurrentRegularization`` (which minimizes ``w*|K|``, the regularization term from surface current in the REGCOIL algorithm, with `w` being the objective weight which act as the regularization parameter)
+    * 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.
 
 
 Bug Fixes

desc/compute/_surface.py

@@ -382,8 +382,8 @@ def _Phi_z_FourierCurrentPotentialField(params, transforms, profiles, data, **kw
     units_long="Amperes",
     description="Surface current potential",
     dim=1,
-    params=["params"],
-    transforms={"grid": [], "potential": []},
+    params=[],
+    transforms={"grid": [], "potential": [], "params": []},
     profiles=[],
     coordinates="tz",
     data=[],
@@ -393,7 +393,7 @@ def _Phi_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
     data["Phi"] = transforms["potential"](
         transforms["grid"].nodes[:, 1],
         transforms["grid"].nodes[:, 2],
-        **params["params"]
+        **transforms["params"]
     )
     return data
 
@@ -405,8 +405,8 @@ def _Phi_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
     units_long="Amperes",
     description="Surface current potential, poloidal derivative",
     dim=1,
-    params=["params"],
-    transforms={"grid": [], "potential_dtheta": []},
+    params=[],
+    transforms={"grid": [], "potential_dtheta": [], "params": []},
     profiles=[],
     coordinates="tz",
     data=[],
@@ -416,7 +416,7 @@ def _Phi_t_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
     data["Phi_t"] = transforms["potential_dtheta"](
         transforms["grid"].nodes[:, 1],
         transforms["grid"].nodes[:, 2],
-        **params["params"]
+        **transforms["params"]
     )
     return data
 
@@ -428,8 +428,8 @@ def _Phi_t_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
     units_long="Amperes",
     description="Surface current potential, toroidal derivative",
     dim=1,
-    params=["params"],
-    transforms={"grid": [], "potential_dzeta": []},
+    params=[],
+    transforms={"grid": [], "potential_dzeta": [], "params": []},
     profiles=[],
     coordinates="tz",
     data=[],
@@ -439,11 +439,55 @@ def _Phi_z_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
     data["Phi_z"] = transforms["potential_dzeta"](
         transforms["grid"].nodes[:, 1],
         transforms["grid"].nodes[:, 2],
-        **params["params"]
+        **transforms["params"]
     )
     return data
 
 
+@register_compute_fun(
+    name="K^theta",
+    label="K^{\\theta}",
+    units="A/m^2",
+    units_long="Amperes per square meter",
+    description="Contravariant poloidal component of surface current density",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="tz",
+    data=["Phi_z", "|e_theta x e_zeta|"],
+    parameterization=[
+        "desc.magnetic_fields._current_potential.CurrentPotentialField",
+        "desc.magnetic_fields._current_potential.FourierCurrentPotentialField",
+    ],
+)
+def _K_sup_theta_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
+    data["K^theta"] = -data["Phi_z"] * (1 / data["|e_theta x e_zeta|"])
+    return data
+
+
+@register_compute_fun(
+    name="K^zeta",
+    label="K^{\\zeta}",
+    units="A/m^2",
+    units_long="Amperes per square meter",
+    description="Contravariant toroidal component of surface current density",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="tz",
+    data=["Phi_t", "|e_theta x e_zeta|"],
+    parameterization=[
+        "desc.magnetic_fields._current_potential.CurrentPotentialField",
+        "desc.magnetic_fields._current_potential.FourierCurrentPotentialField",
+    ],
+)
+def _K_sup_zeta_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
+    data["K^zeta"] = data["Phi_t"] * (1 / data["|e_theta x e_zeta|"])
+    return data
+
+
 @register_compute_fun(
     name="K",
     label="\\mathbf{K}",
@@ -456,16 +500,16 @@ def _Phi_z_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
     transforms={},
     profiles=[],
     coordinates="tz",
-    data=["Phi_t", "Phi_z", "e_theta", "e_zeta", "|e_theta x e_zeta|"],
+    data=["K^theta", "K^zeta", "e_theta", "e_zeta"],
     parameterization=[
         "desc.magnetic_fields._current_potential.CurrentPotentialField",
         "desc.magnetic_fields._current_potential.FourierCurrentPotentialField",
     ],
 )
 def _K_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
-    data["K"] = (
-        data["Phi_t"] * (1 / data["|e_theta x e_zeta|"]) * data["e_zeta"].T
-    ).T - (data["Phi_z"] * (1 / data["|e_theta x e_zeta|"]) * data["e_theta"].T).T
+    data["K"] = (data["K^zeta"] * data["e_zeta"].T).T + (
+        data["K^theta"] * data["e_theta"].T
+    ).T
     return data
 
 

desc/magnetic_fields/__init__.py

@@ -15,5 +15,9 @@
     field_line_integrate,
     read_BNORM_file,
 )
-from ._current_potential import CurrentPotentialField, FourierCurrentPotentialField
+from ._current_potential import (
+    CurrentPotentialField,
+    FourierCurrentPotentialField,
+    solve_regularized_surface_current,
+)
 from ._dommaschk import DommaschkPotentialField, dommaschk_potential

desc/magnetic_fields/_core.py

@@ -319,9 +319,11 @@ def compute_Bnormal(
         Returns
         -------
         Bnorm : ndarray
-            The normal magnetic field to the surface given, of size grid.num_nodes.
+            The normal magnetic field to the surface given, as an array of
+            size ``grid.num_nodes``.
         coords: ndarray
-            the locations (in specified basis) at which the Bnormal was calculated
+            the locations (in specified basis) at which the Bnormal was calculated,
+            given as a ``(grid.num_nodes , 3)`` shaped array.
 
         """
         calc_Bplasma = False