Implement FourierRZCurveWindingSurface Class

desc/coils.py

@@ -22,6 +22,7 @@
 from desc.geometry import (
     FourierPlanarCurve,
     FourierRZCurve,
+    FourierRZWindingSurfaceCurve,
     FourierXYZCurve,
     SplineXYZCurve,
 )
@@ -699,6 +700,64 @@ def from_values(cls, current, coords, N=10, NFP=1, basis="rpz", sym=False, name=
         )
 
 
+class FourierRZWindingSurfaceCoil(_Coil, FourierRZWindingSurfaceCurve):
+    """Coil parameterized by Fourier series for theta,zeta in terms of parameter s.
+
+    This curve will lie on the given winding surface, parameterized by a
+    Fourier series given by Rb_mn and Zb_mn.
+
+    Parameters
+    ----------
+    current : float
+        current through coil, in Amperes
+    surface : FourierRZToroidalSurface
+        Winding surface that the coil will lie on.
+    theta_n, zeta_n: array-like
+        Fourier coefficients for theta, zeta in terms of curve parameter s.
+    modes_theta : array-like, optional
+        Mode numbers associated with theta_n. If not given defaults to [-n:n].
+    modes_zeta : array-like, optional
+        Mode numbers associated with zeta_n, If not given defaults to [-n:n]].
+    sym_theta : {"cos", "sin", False}, optional
+        Whether to enforce symmetry for the theta(t) Fourier series. Defaults to "sin"
+    sym_zeta : {"cos", "sin", False}, optional
+        Whether to enforce symmetry for the zeta(t) Fourier series. Defaults to "sin"
+    name : str
+        Name for this coil.
+
+    """
+
+    _io_attrs_ = _Coil._io_attrs_ + FourierRZWindingSurfaceCurve._io_attrs_
+
+    def __init__(
+        self,
+        current=1,
+        surface=None,
+        theta_n=1,
+        zeta_n=1,
+        secular_theta=1.0,
+        secular_zeta=1.0,
+        modes_theta=None,
+        modes_zeta=None,
+        sym_theta="sin",
+        sym_zeta="sin",
+        name="",
+    ):
+        super().__init__(
+            current,
+            surface,
+            theta_n,
+            zeta_n,
+            secular_theta,
+            secular_zeta,
+            modes_theta,
+            modes_zeta,
+            sym_theta,
+            sym_zeta,
+            name,
+        )
+
+
 class FourierXYZCoil(_Coil, FourierXYZCurve):
     """Coil parameterized by fourier series for X,Y,Z in terms of arbitrary angle s.
 
@@ -1184,6 +1243,8 @@ def _check_type(coil0, coil):
         FourierXYZCoil: ["X_basis", "Y_basis", "Z_basis"],
         FourierPlanarCoil: ["r_basis"],
         SplineXYZCoil: ["method", "N", "knots"],
+        # TODO: probably add a surf_R_basis or R_basis and Z_basis to this
+        FourierRZWindingSurfaceCoil: ["N", "theta_basis", "zeta_basis", "surface"],
     }
 
     for attr in attrs[coil0.__class__]:

desc/compute/_core.py

@@ -2758,6 +2758,10 @@ def _phi_r(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["omega_rr"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+    ],
 )
 def _phi_rr(params, transforms, profiles, data, **kwargs):
     data["phi_rr"] = data["omega_rr"]
@@ -2777,6 +2781,10 @@ def _phi_rr(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["omega_rrz"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+    ],
 )
 def _phi_rrz(params, transforms, profiles, data, **kwargs):
     data["phi_rrz"] = data["omega_rrz"]
@@ -2796,6 +2804,10 @@ def _phi_rrz(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["omega_rt"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+    ],
 )
 def _phi_rt(params, transforms, profiles, data, **kwargs):
     data["phi_rt"] = data["omega_rt"]
@@ -2815,6 +2827,10 @@ def _phi_rt(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["omega_rtz"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+    ],
 )
 def _phi_rtz(params, transforms, profiles, data, **kwargs):
     data["phi_rtz"] = data["omega_rtz"]
@@ -2834,6 +2850,10 @@ def _phi_rtz(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["omega_rz"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+    ],
 )
 def _phi_rz(params, transforms, profiles, data, **kwargs):
     data["phi_rz"] = data["omega_rz"]
@@ -2853,6 +2873,10 @@ def _phi_rz(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["omega_rzz"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+    ],
 )
 def _phi_rzz(params, transforms, profiles, data, **kwargs):
     data["phi_rzz"] = data["omega_rzz"]
@@ -2894,12 +2918,39 @@ def _phi_t(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["omega_tt"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+    ],
 )
 def _phi_tt(params, transforms, profiles, data, **kwargs):
     data["phi_tt"] = data["omega_tt"]
     return data
 
 
+@register_compute_fun(
+    name="phi_ttt",
+    label="\\partial_{\\theta \\theta \\theta} \\phi",
+    units="rad",
+    units_long="radians",
+    description="Toroidal angle in lab frame, third derivative wrt poloidal "
+    "coordinate",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["omega_ttt"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+    ],
+)
+def _phi_ttt(params, transforms, profiles, data, **kwargs):
+    data["phi_ttt"] = data["omega_ttt"]
+    return data
+
+
 @register_compute_fun(
     name="phi_ttz",
     label="\\partial_{\\theta \\theta \\zeta} \\phi",
@@ -2913,6 +2964,10 @@ def _phi_tt(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["omega_ttz"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+    ],
 )
 def _phi_ttz(params, transforms, profiles, data, **kwargs):
     data["phi_ttz"] = data["omega_ttz"]
@@ -2932,6 +2987,10 @@ def _phi_ttz(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["omega_tz"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+    ],
 )
 def _phi_tz(params, transforms, profiles, data, **kwargs):
     data["phi_tz"] = data["omega_tz"]
@@ -2951,6 +3010,10 @@ def _phi_tz(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["omega_tzz"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+    ],
 )
 def _phi_tzz(params, transforms, profiles, data, **kwargs):
     data["phi_tzz"] = data["omega_tzz"]
@@ -2992,6 +3055,10 @@ def _phi_z(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["omega_zz"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+    ],
 )
 def _phi_zz(params, transforms, profiles, data, **kwargs):
     data["phi_zz"] = data["omega_zz"]
@@ -3011,6 +3078,10 @@ def _phi_zz(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["omega_zzz"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+    ],
 )
 def _phi_zzz(params, transforms, profiles, data, **kwargs):
     data["phi_zzz"] = data["omega_zzz"]