Adding a separate factor for NFP to calculate umbilic configurations

desc/backend.py

@@ -68,6 +68,9 @@
 if use_jax:  # noqa: C901 - FIXME: simplify this, define globally and then assign?
     from jax import custom_jvp, jit, vmap
 
+    # --no-verifyfrom jax import custom_jvp, vmap
+    # --no-verifyjit = lambda func, *args, **kwargs: func
+
     imap = jax.lax.map
     from jax.experimental.ode import odeint
     from jax.lax import cond, fori_loop, scan, switch, while_loop

desc/basis.py

@@ -30,6 +30,7 @@ class _Basis(IOAble, ABC):
         "_M",
         "_N",
         "_NFP",
+        # --no-verify "_NFP_umbilic_factor",
         "_modes",
         "_sym",
         "_spectral_indexing",
@@ -186,6 +187,11 @@ def NFP(self):
         """int: Number of field periods."""
         return self.__dict__.setdefault("_NFP", 1)
 
+    @property
+    def NFP_umbilic_factor(self):
+        """int: Umbilic factor field periods."""
+        return self.__dict__.setdefault("_NFP_umbilic_factor", 1)
+
     @property
     def sym(self):
         """str: Type of symmetry."""
@@ -213,8 +219,14 @@ def __repr__(self):
             type(self).__name__
             + " at "
             + str(hex(id(self)))
-            + " (L={}, M={}, N={}, NFP={}, sym={}, spectral_indexing={})".format(
-                self.L, self.M, self.N, self.NFP, self.sym, self.spectral_indexing
+            + " (L={}, M={}, N={}, NFP={}, sym={},\
+                spectral_indexing={})".format(
+                self.L,
+                self.M,
+                self.N,
+                self.NFP,
+                self.sym,
+                self.spectral_indexing,
             )
         )
 
@@ -239,6 +251,7 @@ def __init__(self, L, sym="even"):
         self._M = 0
         self._N = 0
         self._NFP = 1
+        self._sym = sym
         self._sym = bool(sym) if not sym else str(sym)
         self._spectral_indexing = "linear"
 
@@ -340,18 +353,24 @@ class FourierSeries(_Basis):
         Maximum toroidal resolution.
     NFP : int
         number of field periods
+    NFP_umbilic_factor : int
+        Prefactor of the form 1/NFP_umbilic_fac.
+        This is needed for the umbilic torus design.
     sym : {``'cos'``, ``'sin'``, False}
         * ``'cos'`` for cos(m*t-n*z) symmetry
         * ``'sin'`` for sin(m*t-n*z) symmetry
         * ``False`` for no symmetry (Default)
 
     """
 
-    def __init__(self, N, NFP=1, sym=False):
+    def __init__(self, N, NFP=1, NFP_umbilic_factor=1, sym=False):
         self._L = 0
         self._M = 0
         self._N = check_nonnegint(N, "N", False)
         self._NFP = check_posint(NFP, "NFP", False)
+        self._NFP_umbilic_factor = check_posint(
+            NFP_umbilic_factor, "NFP_umbilic_factor", False
+        )
         self._sym = bool(sym) if not sym else str(sym)
         self._spectral_indexing = "linear"
 
@@ -422,13 +441,20 @@ def evaluate(
             z = z[zidx]
             n = n[nidx]
 
-        toroidal = fourier(z[:, np.newaxis], n, self.NFP, derivatives[2])
+        toroidal = fourier(
+            z[:, np.newaxis],
+            n,
+            self.NFP,
+            self.NFP_umbilic_factor,
+            derivatives[2],
+        )
+
         if unique:
             toroidal = toroidal[zoutidx][:, noutidx]
 
         return toroidal
 
-    def change_resolution(self, N, NFP=None, sym=None):
+    def change_resolution(self, N, NFP=None, NFP_umbilic_factor=None, sym=None):
         """Change resolution of the basis to the given resolutions.
 
         Parameters
@@ -437,12 +463,20 @@ def change_resolution(self, N, NFP=None, sym=None):
             Maximum toroidal resolution.
         NFP : int
             Number of field periods.
+        NFP_umbilic_factor : float
+            Prefactor of the form 1/NFP_umbilic_factor.
+            This is needed for the umbilic torus design.
         sym : bool
             Whether to enforce stellarator symmetry.
 
         """
         NFP = check_posint(NFP, "NFP")
         self._NFP = NFP if NFP is not None else self.NFP
+        self._NFP_umbilic_factor = (
+            NFP_umbilic_factor
+            if NFP_umbilic_factor is not None
+            else self.NFP_umbilic_factor
+        )
         if N != self.N:
             self._N = check_nonnegint(N, "N", False)
             self._sym = sym if sym is not None else self.sym
@@ -561,8 +595,16 @@ def evaluate(
             m = m[midx]
             n = n[nidx]
 
-        poloidal = fourier(t[:, np.newaxis], m, 1, derivatives[1])
-        toroidal = fourier(z[:, np.newaxis], n, self.NFP, derivatives[2])
+        poloidal = fourier(
+            t[:, np.newaxis], m, NFP=1, NFP_umbilic_factor=1, dt=derivatives[1]
+        )
+        toroidal = fourier(
+            z[:, np.newaxis],
+            n,
+            NFP=self.NFP,
+            NFP_umbilic_factor=1,
+            dt=derivatives[2],
+        )
         if unique:
             poloidal = poloidal[toutidx][:, moutidx]
             toroidal = toroidal[zoutidx][:, noutidx]
@@ -770,7 +812,7 @@ def evaluate(
             m = m[midx]
 
         radial = zernike_radial(r[:, np.newaxis], lm[:, 0], lm[:, 1], dr=derivatives[0])
-        poloidal = fourier(t[:, np.newaxis], m, 1, derivatives[1])
+        poloidal = fourier(t[:, np.newaxis], m, 1, 1, derivatives[1])
 
         if unique:
             radial = radial[routidx][:, lmoutidx]
@@ -902,8 +944,8 @@ def evaluate(
         l, m, n = modes.T
 
         radial = chebyshev(r[:, np.newaxis], l, dr=derivatives[0])
-        poloidal = fourier(t[:, np.newaxis], m, 1, derivatives[1])
-        toroidal = fourier(z[:, np.newaxis], n, self.NFP, derivatives[2])
+        poloidal = fourier(t[:, np.newaxis], m, 1, 1, derivatives[1])
+        toroidal = fourier(z[:, np.newaxis], n, self.NFP, 1, derivatives[2])
 
         return radial * poloidal * toroidal
 
@@ -1132,8 +1174,14 @@ def evaluate(
             n = n[nidx]
 
         radial = zernike_radial(r[:, np.newaxis], lm[:, 0], lm[:, 1], dr=derivatives[0])
-        poloidal = fourier(t[:, np.newaxis], m, dt=derivatives[1])
-        toroidal = fourier(z[:, np.newaxis], n, NFP=self.NFP, dt=derivatives[2])
+        poloidal = fourier(t[:, np.newaxis], m, 1, 1, dt=derivatives[1])
+        toroidal = fourier(
+            z[:, np.newaxis],
+            n,
+            NFP=self.NFP,
+            NFP_umbilic_factor=1,
+            dt=derivatives[2],
+        )
         if unique:
             radial = radial[routidx][:, lmoutidx]
             poloidal = poloidal[toutidx][:, moutidx]
@@ -1652,7 +1700,7 @@ def chebyshev(r, l, dr=0):
 
 
 @jit
-def fourier(theta, m, NFP=1, dt=0):
+def fourier(theta, m, NFP=1, NFP_umbilic_factor=1, dt=0):
     """Fourier series.
 
     Parameters
@@ -1663,6 +1711,9 @@ def fourier(theta, m, NFP=1, dt=0):
         poloidal/toroidal mode number(s)
     NFP : int
         number of field periods (Default = 1)
+    NFP_umbilic_factor : int
+        NFP prefactor of the form 1/NFP_umbilic_factor.
+        This is needed for the umbilic torus design.
     dt : int
         order of derivative (Default = 0)
 
@@ -1672,9 +1723,11 @@ def fourier(theta, m, NFP=1, dt=0):
         basis function(s) evaluated at specified points
 
     """
-    theta, m, NFP, dt = map(jnp.asarray, (theta, m, NFP, dt))
+    theta, m, NFP, NFP_umbilic_factor, dt = map(
+        jnp.asarray, (theta, m, NFP, NFP_umbilic_factor, dt)
+    )
     m_pos = (m >= 0).astype(int)
-    m_abs = jnp.abs(m) * NFP
+    m_abs = jnp.abs(m) * NFP / NFP_umbilic_factor
     shift = m_pos * jnp.pi / 2 + dt * jnp.pi / 2
     return m_abs**dt * jnp.sin(m_abs * theta + shift)
 

desc/coils.py

@@ -279,6 +279,9 @@ def _compute_position(self, params=None, grid=None, dx1=False, **kwargs):
         basis = kwargs.get("basis", "xyz")
         if basis.lower() == "rpz":
             x = x.at[:, :, 1].set(jnp.mod(x[:, :, 1], 2 * jnp.pi))
+
+            # --no-verify x = x.at[:, :, 1].set(x[:, :, 1])
+
         if dx1:
             return x, x_s
         return x
@@ -533,7 +536,9 @@ def to_SplineXYZ(self, knots=None, grid=None, method="cubic", name="", **kwargs)
             self.current, coords, knots=knots, method=method, name=name, basis="xyz"
         )
 
-    def to_FourierRZ(self, N=10, grid=None, NFP=None, sym=False, name="", **kwargs):
+    def to_FourierRZ(
+        self, N=10, grid=None, NFP=None, NFP_umbilic_factor=None, sym=False, name=""
+    ):
         """Convert Coil to FourierRZCoil representation.
 
         Note that some types of coils may not be representable in this basis.
@@ -560,11 +565,19 @@ def to_FourierRZ(self, N=10, grid=None, NFP=None, sym=False, name="", **kwargs):
 
         """
         NFP = 1 or NFP
+        NFP_umbilic_factor = 1 or NFP_umbilic_factor
         if grid is None:
-            grid = LinearGrid(N=2 * N + 1)
+            grid = LinearGrid(N=2 * N + 1, NFP_umbilic_factor=NFP_umbilic_factor)
         coords = self.compute("x", grid=grid, basis="xyz")["x"]
         return FourierRZCoil.from_values(
-            self.current, coords, N=N, NFP=NFP, basis="xyz", sym=sym, name=name
+            self.current,
+            coords,
+            N=N,
+            NFP=NFP,
+            NFP_umbilic_factor=NFP_umbilic_factor,
+            basis="xyz",
+            sym=sym,
+            name=name,
         )
 
     def to_FourierPlanar(self, N=10, grid=None, basis="xyz", name="", **kwargs):
@@ -616,6 +629,8 @@ class FourierRZCoil(_Coil, FourierRZCurve):
         mode numbers associated with Z_n, defaults to modes_R
     NFP : int
         number of field periods
+    NFP_umbilic_factor : int
+        umbilic factor
     sym : bool
         whether to enforce stellarator symmetry
     name : str
@@ -663,13 +678,26 @@ def __init__(
         modes_R=None,
         modes_Z=None,
         NFP=1,
+        NFP_umbilic_factor=1,
         sym="auto",
         name="",
     ):
-        super().__init__(current, R_n, Z_n, modes_R, modes_Z, NFP, sym, name)
+        super().__init__(
+            current, R_n, Z_n, modes_R, modes_Z, NFP, NFP_umbilic_factor, sym, name
+        )
 
     @classmethod
-    def from_values(cls, current, coords, N=10, NFP=1, basis="rpz", sym=False, name=""):
+    def from_values(
+        cls,
+        current,
+        coords,
+        N=10,
+        NFP=1,
+        NFP_umbilic_factor=1,
+        basis="rpz",
+        name="",
+        sym=False,
+    ):
         """Fit coordinates to FourierRZCoil representation.
 
         Parameters
@@ -684,6 +712,9 @@ def from_values(cls, current, coords, N=10, NFP=1, basis="rpz", sym=False, name=
         NFP : int
             Number of field periods, the curve will have a discrete toroidal symmetry
             according to NFP.
+        NFP_umbilic_factor : int
+            Umbilic factor to fit curves that go around multiple times toroidally before
+            closing on themselves.
         basis : {"rpz", "xyz"}
             basis for input coordinates. Defaults to "rpz"
         sym : bool
@@ -699,7 +730,13 @@ def from_values(cls, current, coords, N=10, NFP=1, basis="rpz", sym=False, name=
 
         """
         curve = super().from_values(
-            coords=coords, N=N, NFP=NFP, basis=basis, sym=sym, name=name
+            coords,
+            N=N,
+            NFP=NFP,
+            NFP_umbilic_factor=NFP_umbilic_factor,
+            basis=basis,
+            sym=sym,
+            name=name,
         )
         return FourierRZCoil(
             current=current,
@@ -708,6 +745,7 @@ def from_values(cls, current, coords, N=10, NFP=1, basis="rpz", sym=False, name=
             modes_R=curve.R_basis.modes[:, 2],
             modes_Z=curve.Z_basis.modes[:, 2],
             NFP=NFP,
+            NFP_umbilic_factor=NFP_umbilic_factor,
             sym=curve.sym,
             name=name,
         )
@@ -2036,7 +2074,14 @@ def to_FourierPlanar(
         )
 
     def to_FourierRZ(
-        self, N=10, grid=None, NFP=None, sym=False, name="", check_intersection=True
+        self,
+        N=10,
+        grid=None,
+        NFP=None,
+        NFP_umbilic_factor=None,
+        sym=False,
+        name="",
+        check_intersection=True,
     ):
         """Convert all coils to FourierRZCoil representaion.
 
@@ -2065,7 +2110,12 @@ def to_FourierRZ(
             New representation of the coilset parameterized by a Fourier series for R,Z.
 
         """
-        coils = [coil.to_FourierRZ(N=N, grid=grid, NFP=NFP, sym=sym) for coil in self]
+        coils = [
+            coil.to_FourierRZ(
+                N=N, grid=grid, NFP=NFP, NFP_umbilic_factor=NFP_umbilic_factor, sym=sym
+            )
+            for coil in self
+        ]
         return self.__class__(
             *coils,
             NFP=self.NFP,

desc/compat.py

@@ -210,7 +210,12 @@ def rescale(
         eq = eq.copy()
 
     # size scaling
-    grid_L = QuadratureGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP)
+    grid_L = QuadratureGrid(
+        L=eq.L_grid,
+        M=eq.M_grid,
+        N=eq.N_grid,
+        NFP=eq.NFP,
+    )
     data_L = eq.compute(L_keys, grid=grid_L)
     L_old = data_L[L_key]
     L_new = L_new or L_old
@@ -227,7 +232,12 @@ def rescale(
         data_B = eq.compute("<|B|>_vol", grid=grid_B)
         B_old = data_B["<|B|>_vol"]
     elif B_key == "B_max":
-        grid_B = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, rho=1)
+        grid_B = LinearGrid(
+            M=eq.M_grid,
+            N=eq.N_grid,
+            NFP=eq.NFP,
+            rho=1,
+        )
         data_B = eq.compute("|B|", grid=grid_B)
         B_old = np.max(data_B["|B|"])
     B_new = B_new or B_old

desc/compute/__init__.py

@@ -35,10 +35,12 @@
     _field,
     _geometry,
     _metric,
+    _neoclassical,
     _omnigenity,
     _profiles,
     _stability,
     _surface,
+    _umbiliccurve,
 )
 from .data_index import all_kwargs, allowed_kwargs, data_index
 from .geom_utils import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec