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
@@ -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/compute/_basis_vectors.py

@@ -3223,17 +3223,61 @@ def _e_sub_zeta_zz(params, transforms, profiles, data, **kwargs):
     transforms={},
     profiles=[],
     coordinates="rtz",
-    data=["e^rho", "e^theta", "e^zeta", "alpha_r", "alpha_t", "alpha_z"],
+    data=["grad(alpha) (periodic)", "grad(alpha) (secular)"],
 )
 def _grad_alpha(params, transforms, profiles, data, **kwargs):
-    data["grad(alpha)"] = (
-        data["alpha_r"] * data["e^rho"].T
+    data["grad(alpha)"] = data["grad(alpha) (periodic)"] + data["grad(alpha) (secular)"]
+    return data
+
+
+@register_compute_fun(
+    name="grad(alpha) (periodic)",
+    label="\\mathrm{periodic}(\\nabla \\alpha)",
+    units="m^{-1}",
+    units_long="Inverse meters",
+    description=(
+        "Gradient of field line label, which is perpendicular to the field line, "
+        "periodic component"
+    ),
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["e^rho", "e^theta", "e^zeta", "alpha_r (periodic)", "alpha_t", "alpha_z"],
+)
+def _periodic_grad_alpha(params, transforms, profiles, data, **kwargs):
+    data["grad(alpha) (periodic)"] = (
+        data["alpha_r (periodic)"] * data["e^rho"].T
         + data["alpha_t"] * data["e^theta"].T
         + data["alpha_z"] * data["e^zeta"].T
     ).T
     return data
 
 
+@register_compute_fun(
+    name="grad(alpha) (secular)",
+    label="\\mathrm{secular}(\\nabla \\alpha)",
+    units="m^{-1}",
+    units_long="Inverse meters",
+    description=(
+        "Gradient of field line label, which is perpendicular to the field line, "
+        "secular component"
+    ),
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["e^rho", "alpha_r (secular)"],
+)
+def _secular_grad_alpha(params, transforms, profiles, data, **kwargs):
+    data["grad(alpha) (secular)"] = (
+        data["alpha_r (secular)"][:, jnp.newaxis] * data["e^rho"]
+    )
+    return data
+
+
 @register_compute_fun(
     name="grad(psi)",
     label="\\nabla\\psi",

desc/compute/_core.py

@@ -1518,14 +1518,48 @@ def _alpha(params, transforms, profiles, data, **kwargs):
     transforms={},
     profiles=[],
     coordinates="rtz",
-    data=["theta_PEST_r", "phi", "phi_r", "iota", "iota_r"],
+    data=["alpha_r (periodic)", "alpha_r (secular)"],
 )
 def _alpha_r(params, transforms, profiles, data, **kwargs):
-    data["alpha_r"] = (
-        data["theta_PEST_r"]
-        - data["iota_r"] * data["phi"]
-        - data["iota"] * data["phi_r"]
-    )
+    data["alpha_r"] = data["alpha_r (periodic)"] + data["alpha_r (secular)"]
+    return data
+
+
+@register_compute_fun(
+    name="alpha_r (periodic)",
+    label="\\mathrm{periodic}(\\partial_\\rho \\alpha)",
+    units="~",
+    units_long="None",
+    description="Field line label, derivative wrt radial coordinate, "
+    "periodic component",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["theta_PEST_r", "iota", "phi_r"],
+)
+def _periodic_alpha_r(params, transforms, profiles, data, **kwargs):
+    data["alpha_r (periodic)"] = data["theta_PEST_r"] - data["iota"] * data["phi_r"]
+    return data
+
+
+@register_compute_fun(
+    name="alpha_r (secular)",
+    label="\\mathrm{secular}(\\partial_\\rho \\alpha)",
+    units="~",
+    units_long="None",
+    description="Field line label, derivative wrt radial coordinate, "
+    "secular component",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["iota_r", "phi"],
+)
+def _secular_alpha_r(params, transforms, profiles, data, **kwargs):
+    data["alpha_r (secular)"] = -data["iota_r"] * data["phi"]
     return data
 
 

desc/compute/_metric.py

@@ -1941,24 +1941,85 @@ def _g_sup_ra(params, transforms, profiles, data, **kwargs):
     # Exact definition of the magnetic drifts taken from
     # eqn. 48 of Introduction to Quasisymmetry by Landreman
     # https://tinyurl.com/54udvaa4
-    label="\\mathrm{gbdrift} = 1/B^{2} (\\mathbf{b}\\times\\nabla B) \\cdot"
-    + "\\nabla \\alpha",
-    units="1/(T-m^{2})",
-    units_long="inverse Tesla meters^2",
-    description="Binormal component of the geometric part of the gradB drift"
-    + " used for local stability analyses, Gamma_c, epsilon_eff etc.",
+    label="(\\nabla \\vert B \\vert)_{\\mathrm{drift}} = "
+    "(\\mathbf{b} \\times \\nabla B) \\cdot \\nabla \\alpha / \\vert B \\vert^{2}",
+    units="1 / Wb",
+    units_long="Inverse webers",
+    description="Binormal, geometric part of the gradB drift. "
+    "Used for local stability analyses, gyrokinetics, and Gamma_c.",
     dim=1,
     params=[],
     transforms={},
     profiles=[],
     coordinates="rtz",
-    data=["|B|^2", "b", "grad(alpha)", "grad(|B|)"],
+    data=["gbdrift (periodic)", "gbdrift (secular)"],
 )
 def _gbdrift(params, transforms, profiles, data, **kwargs):
-    data["gbdrift"] = (
-        1
+    data["gbdrift"] = data["gbdrift (periodic)"] + data["gbdrift (secular)"]
+    return data
+
+
+@register_compute_fun(
+    name="gbdrift (periodic)",
+    label="\\mathrm{periodic}(\\nabla \\vert B \\vert)_{\\mathrm{drift}}",
+    units="1 / Wb",
+    units_long="Inverse webers",
+    description="Periodic, binormal, geometric part of the gradB drift.",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["|B|^2", "b", "grad(alpha) (periodic)", "grad(|B|)"],
+)
+def _periodic_gbdrift(params, transforms, profiles, data, **kwargs):
+    data["gbdrift (periodic)"] = (
+        dot(data["b"], cross(data["grad(|B|)"], data["grad(alpha) (periodic)"]))
+        / data["|B|^2"]
+    )
+    return data
+
+
+@register_compute_fun(
+    name="gbdrift (secular)",
+    label="\\mathrm{secular}(\\nabla \\vert B \\vert)_{\\mathrm{drift}}",
+    units="1 / Wb",
+    units_long="Inverse webers",
+    description="Secular, binormal, geometric part of the gradB drift.",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["|B|^2", "b", "grad(alpha) (secular)", "grad(|B|)"],
+)
+def _secular_gbdrift(params, transforms, profiles, data, **kwargs):
+    data["gbdrift (secular)"] = (
+        dot(data["b"], cross(data["grad(|B|)"], data["grad(alpha) (secular)"]))
+        / data["|B|^2"]
+    )
+    return data
+
+
+@register_compute_fun(
+    name="gbdrift (secular)/phi",
+    label="\\mathrm{secular}(\\nabla \\vert B \\vert)_{\\mathrm{drift}} / \\phi",
+    units="1 / Wb",
+    units_long="Inverse webers",
+    description="Secular, binormal, geometric part of the gradB drift divided "
+    "by the toroidal angle. This quantity is periodic.",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["|B|^2", "b", "e^rho", "grad(|B|)", "iota_r"],
+)
+def _secular_gbdrift_over_phi(params, transforms, profiles, data, **kwargs):
+    data["gbdrift (secular)/phi"] = (
+        dot(data["b"], cross(data["e^rho"], data["grad(|B|)"]))
+        * data["iota_r"]
         / data["|B|^2"]
-        * dot(data["b"], cross(data["grad(|B|)"], data["grad(alpha)"]))
     )
     return data
 
@@ -1970,34 +2031,53 @@ def _gbdrift(params, transforms, profiles, data, **kwargs):
     # https://tinyurl.com/54udvaa4
     label="\\mathrm{cvdrift} = 1/B^{3} (\\mathbf{b}\\times\\nabla(p + B^2/2))"
     + "\\cdot \\nabla \\alpha",
-    units="1/(T-m^{2})",
-    units_long="inverse Tesla meters^2",
-    description="Binormal component of the geometric part of the curvature drift"
-    + " used for local stability analyses, Gamma_c, epsilon_eff etc.",
+    units="1 / Wb",
+    units_long="Inverse webers",
+    description="Binormal, geometric part of the curvature drift. "
+    "Used for local stability analyses and gyrokinetics.",
     dim=1,
     params=[],
     transforms={},
     profiles=[],
     coordinates="rtz",
-    data=["p_r", "psi_r", "|B|^2", "gbdrift"],
+    data=["cvdrift (periodic)", "gbdrift (secular)"],
 )
 def _cvdrift(params, transforms, profiles, data, **kwargs):
-    dp_dpsi = mu_0 * data["p_r"] / data["psi_r"]
-    data["cvdrift"] = 1 / data["|B|^2"] * dp_dpsi + data["gbdrift"]
+    data["cvdrift"] = data["cvdrift (periodic)"] + data["gbdrift (secular)"]
+    return data
+
+
+@register_compute_fun(
+    name="cvdrift (periodic)",
+    label="\\mathrm{cvdrift (periodic)}",
+    units="1 / Wb",
+    units_long="Inverse webers",
+    description="Periodic, binormal, geometric part of the curvature drift.",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["p_r", "psi_r", "|B|^2", "gbdrift (periodic)"],
+)
+def _periodic_cvdrift(params, transforms, profiles, data, **kwargs):
+    data["cvdrift (periodic)"] = (
+        mu_0 * data["p_r"] / data["psi_r"] / data["|B|^2"] + data["gbdrift (periodic)"]
+    )
     return data
 
 
 @register_compute_fun(
     name="cvdrift0",
     # Exact definition of the magnetic drifts taken from
     # eqn. 48 of Introduction to Quasisymmetry by Landreman
-    # https://tinyurl.com/54udvaa4
+    # https://tinyurl.com/54udvaa4 up to dimensionless factors.
     label="\\mathrm{cvdrift0} = 1/B^{2} (\\mathbf{b}\\times\\nabla B)"
-    + "\\cdot \\nabla \\rho",
-    units="1/(T-m^{2})",
-    units_long="inverse Tesla meters^2",
-    description="Radial component of the geometric part of the curvature drift"
-    + " used for local stability analyses, gyrokinetics, Gamma_c.",
+    + "\\cdot (2 \\rho \\nabla \\rho)",
+    units="1 / Wb",
+    units_long="Inverse webers",
+    description="Radial, geometric part of the curvature drift."
+    + " Used for local stability analyses, gyrokinetics, and Gamma_c.",
     dim=1,
     params=[],
     transforms={},
@@ -2010,6 +2090,6 @@ def _cvdrift0(params, transforms, profiles, data, **kwargs):
         2
         * data["rho"]
         / data["|B|^2"]
-        * (dot(data["b"], cross(data["grad(|B|)"], data["e^rho"])))
+        * dot(data["b"], cross(data["grad(|B|)"], data["e^rho"]))
     )
     return data

desc/equilibrium/equilibrium.py

@@ -1061,8 +1061,10 @@ def need_src(name):
                 sym=self.sym
                 and all(
                     data_index[p][dep]["grid_requirement"].get("sym", True)
-                    # TODO: GitHub issue #1206.
-                    and not data_index[p][dep]["grid_requirement"].get("can_fft", False)
+                    # TODO (#1206)
+                    and not data_index[p][dep]["grid_requirement"].get(
+                        "can_fft2", False
+                    )
                     for dep in dep1dr
                 ),
             )
@@ -1109,8 +1111,10 @@ def need_src(name):
                 sym=self.sym
                 and all(
                     data_index[p][dep]["grid_requirement"].get("sym", True)
-                    # TODO: GitHub issue #1206.
-                    and not data_index[p][dep]["grid_requirement"].get("can_fft", False)
+                    # TODO (#1206)
+                    and not data_index[p][dep]["grid_requirement"].get(
+                        "can_fft2", False
+                    )
                     for dep in dep1dz
                 ),
             )

desc/grid.py

@@ -43,6 +43,7 @@ class _Grid(IOAble, ABC):
         "_inverse_poloidal_idx",
         "_inverse_zeta_idx",
         "_is_meshgrid",
+        "_can_fft2",
     ]
 
     @abstractmethod
@@ -231,6 +232,16 @@ def is_meshgrid(self):
         """
         return self.__dict__.setdefault("_is_meshgrid", False)
 
+    @property
+    def can_fft2(self):
+        """bool: Whether this grid is compatible with 2D FFT.
+
+        Tensor product grid with uniformly spaced points on
+        (θ, ζ) ∈ [0, 2π) × [0, 2π/NFP).
+        """
+        # TODO: GitHub issue 1243?
+        return self.__dict__.setdefault("_can_fft2", self.is_meshgrid and not self.sym)
+
     @property
     def coordinates(self):
         """Coordinates specified by the nodes.
@@ -682,6 +693,7 @@ class Grid(_Grid):
         Use np.inf to denote no periodicity.
     NFP : int
         Number of field periods (Default = 1).
+        Change this only if your nodes are placed within one field period.
     source_grid : Grid
         Grid from which coordinates were mapped from.
     sort : bool
@@ -805,7 +817,8 @@ def create_meshgrid(
         NFP : int
             Number of field periods (Default = 1).
             Only makes sense to change from 1 if last coordinate is periodic
-            with some constant divided by ``NFP``.
+            with some constant divided by ``NFP`` and the nodes are placed
+            within one field period.
 
         Returns
         -------
@@ -927,6 +940,8 @@ class LinearGrid(_Grid):
         Toroidal grid resolution.
     NFP : int
         Number of field periods (Default = 1).
+        Change this only if your nodes are placed within one field period
+        or should be interpreted as spanning one field period.
     sym : bool
         True for stellarator symmetry, False otherwise (Default = False).
     axis : bool
@@ -1022,6 +1037,8 @@ def _create_nodes(  # noqa: C901
             Toroidal grid resolution.
         NFP : int
             Number of field periods (Default = 1).
+            Only change this if your nodes are placed within one field period
+            or should be interpreted as spanning one field period.
         axis : bool
             True to include a point at rho=0 (default), False for rho[0] = rho[1]/2.
         endpoint : bool