Dommaschk potentials implemented in JAX (#483)

Dommaschk potentials implemented in JAX

CHANGELOG.md

@@ -15,11 +15,14 @@ in the MAKEGRID format for use with other codes.
 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, 
+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.
+- Adds a new class ``DommaschkPotentialField`` which allows creation of magnetic fields based
+off of the vacuum potentials detailed in Representations for Vacuum Potentials in Stellarators
+https://doi.org/10.1016/0010-4655(86)90109-8.
 
 
 Speed Improvements

desc/compute/_surface.py

@@ -626,7 +626,9 @@ def _e_zeta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwa
     profiles=[],
     coordinates="tz",
     data=[],
-    parameterization="desc.magnetic_fields.FourierCurrentPotentialField",
+    parameterization=[
+        "desc.magnetic_fields._current_potential.FourierCurrentPotentialField"
+    ],
 )
 def _Phi_FourierCurrentPotentialField(params, transforms, profiles, data, **kwargs):
     data["Phi"] = (
@@ -649,7 +651,9 @@ def _Phi_FourierCurrentPotentialField(params, transforms, profiles, data, **kwar
     profiles=[],
     coordinates="tz",
     data=[],
-    parameterization="desc.magnetic_fields.FourierCurrentPotentialField",
+    parameterization=[
+        "desc.magnetic_fields._current_potential.FourierCurrentPotentialField"
+    ],
 )
 def _Phi_t_FourierCurrentPotentialField(params, transforms, profiles, data, **kwargs):
     data["Phi_t"] = (
@@ -670,7 +674,9 @@ def _Phi_t_FourierCurrentPotentialField(params, transforms, profiles, data, **kw
     profiles=[],
     coordinates="tz",
     data=[],
-    parameterization="desc.magnetic_fields.FourierCurrentPotentialField",
+    parameterization=[
+        "desc.magnetic_fields._current_potential.FourierCurrentPotentialField"
+    ],
 )
 def _Phi_z_FourierCurrentPotentialField(params, transforms, profiles, data, **kwargs):
     data["Phi_z"] = (
@@ -691,7 +697,7 @@ def _Phi_z_FourierCurrentPotentialField(params, transforms, profiles, data, **kw
     profiles=[],
     coordinates="tz",
     data=[],
-    parameterization="desc.magnetic_fields.CurrentPotentialField",
+    parameterization="desc.magnetic_fields._current_potential.CurrentPotentialField",
 )
 def _Phi_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
     data["Phi"] = transforms["potential"](
@@ -714,7 +720,7 @@ def _Phi_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="tz",
     data=[],
-    parameterization="desc.magnetic_fields.CurrentPotentialField",
+    parameterization="desc.magnetic_fields._current_potential.CurrentPotentialField",
 )
 def _Phi_t_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
     data["Phi_t"] = transforms["potential_dtheta"](
@@ -737,7 +743,7 @@ def _Phi_t_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="tz",
     data=[],
-    parameterization="desc.magnetic_fields.CurrentPotentialField",
+    parameterization="desc.magnetic_fields._current_potential.CurrentPotentialField",
 )
 def _Phi_z_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
     data["Phi_z"] = transforms["potential_dzeta"](
@@ -762,8 +768,8 @@ def _Phi_z_CurrentPotentialField(params, transforms, profiles, data, **kwargs):
     coordinates="tz",
     data=["Phi_t", "Phi_z", "e_theta", "e_zeta", "|e_theta x e_zeta|"],
     parameterization=[
-        "desc.magnetic_fields.CurrentPotentialField",
-        "desc.magnetic_fields.FourierCurrentPotentialField",
+        "desc.magnetic_fields._current_potential.CurrentPotentialField",
+        "desc.magnetic_fields._current_potential.FourierCurrentPotentialField",
     ],
 )
 def _K_CurrentPotentialField(params, transforms, profiles, data, **kwargs):

desc/compute/data_index.py

@@ -211,15 +211,15 @@ def _decorator(func):
         "desc.geometry.curve.FourierPlanarCurve",
         "desc.geometry.core.Curve",
     ],
-    "desc.magnetic_fields.CurrentPotentialField": [
+    "desc.magnetic_fields._current_potential.CurrentPotentialField": [
         "desc.geometry.surface.FourierRZToroidalSurface",
         "desc.geometry.core.Surface",
-        "desc.magnetic_fields.MagneticField",
+        "desc.magnetic_fields._core.MagneticField",
     ],
-    "desc.magnetic_fields.FourierCurrentPotentialField": [
+    "desc.magnetic_fields._current_potential.FourierCurrentPotentialField": [
         "desc.geometry.surface.FourierRZToroidalSurface",
         "desc.geometry.core.Surface",
-        "desc.magnetic_fields.MagneticField",
+        "desc.magnetic_fields._core.MagneticField",
     ],
     "desc.coils.SplineXYZCoil": [
         "desc.geometry.curve.SplineXYZCurve",

desc/magnetic_fields/__init__.py

@@ -0,0 +1,16 @@
+"""Classes for Magnetic Fields."""
+
+from ._core import (
+    PoloidalMagneticField,
+    ScalarPotentialField,
+    ScaledMagneticField,
+    SplineMagneticField,
+    SumMagneticField,
+    ToroidalMagneticField,
+    VerticalMagneticField,
+    _MagneticField,
+    field_line_integrate,
+    read_BNORM_file,
+)
+from ._current_potential import CurrentPotentialField, FourierCurrentPotentialField
+from ._dommaschk import DommaschkPotentialField, dommaschk_potential