Umbilic fourier bounce

desc/backend.py

@@ -73,7 +73,7 @@
     from jax.lax import cond, fori_loop, scan, switch, while_loop
     from jax.nn import softmax as softargmax
     from jax.numpy import bincount, flatnonzero, repeat, take
-    from jax.numpy.fft import irfft, rfft, rfft2
+    from jax.numpy.fft import ifft, irfft, irfft2, rfft, rfft2
     from jax.scipy.fft import dct, idct
     from jax.scipy.linalg import block_diag, cho_factor, cho_solve, qr, solve_triangular
     from jax.scipy.special import gammaln, logsumexp
@@ -434,7 +434,7 @@ def tangent_solve(g, y):
     jit = lambda func, *args, **kwargs: func
     execute_on_cpu = lambda func: func
     import scipy.optimize
-    from numpy.fft import irfft, rfft, rfft2  # noqa: F401
+    from numpy.fft import ifft, irfft, irfft2, rfft, rfft2  # noqa: F401
     from scipy.fft import dct, idct  # noqa: F401
     from scipy.integrate import odeint  # noqa: F401
     from scipy.linalg import (  # noqa: F401

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,7 @@ 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, sym=False, name=""):
         """Convert Coil to FourierRZCoil representation.
 
         Note that some types of coils may not be representable in this basis.
@@ -564,7 +567,13 @@ def to_FourierRZ(self, N=10, grid=None, NFP=None, sym=False, name="", **kwargs):
             grid = LinearGrid(N=2 * N + 1)
         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,
+            basis="xyz",
+            sym=sym,
+            name=name,
         )
 
     def to_FourierPlanar(self, N=10, grid=None, basis="xyz", name="", **kwargs):
@@ -669,7 +678,16 @@ def __init__(
         super().__init__(current, R_n, Z_n, modes_R, modes_Z, NFP, 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,
+        basis="rpz",
+        name="",
+        sym=False,
+    ):
         """Fit coordinates to FourierRZCoil representation.
 
         Parameters
@@ -699,7 +717,12 @@ 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,
+            basis=basis,
+            sym=sym,
+            name=name,
         )
         return FourierRZCoil(
             current=current,
@@ -2036,7 +2059,13 @@ 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,
+        sym=False,
+        name="",
+        check_intersection=True,
     ):
         """Convert all coils to FourierRZCoil representaion.
 

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,13 @@
     _field,
     _geometry,
     _metric,
+    _neoclassical,
+    _neoclassical_1D,
     _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