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=["periodic(grad(alpha))", "secular(grad(alpha))"],
 )
 def _grad_alpha(params, transforms, profiles, data, **kwargs):
-    data["grad(alpha)"] = (
-        data["alpha_r"] * data["e^rho"].T
+    data["grad(alpha)"] = data["periodic(grad(alpha))"] + data["secular(grad(alpha))"]
+    return data
+
+
+@register_compute_fun(
+    name="periodic(grad(alpha))",
+    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", "periodic(alpha_r)", "alpha_t", "alpha_z"],
+)
+def _periodic_grad_alpha(params, transforms, profiles, data, **kwargs):
+    data["periodic(grad(alpha))"] = (
+        data["periodic(alpha_r)"] * 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="secular(grad(alpha))",
+    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", "secular(alpha_r)"],
+)
+def _secular_grad_alpha(params, transforms, profiles, data, **kwargs):
+    data["secular(grad(alpha))"] = (
+        data["secular(alpha_r)"][:, jnp.newaxis] * data["e^rho"]
+    )
+    return data
+
+
 @register_compute_fun(
     name="grad(psi)",
     label="\\nabla\\psi",

desc/compute/_core.py

@@ -1518,14 +1518,48 @@
     transforms={},
     profiles=[],
     coordinates="rtz",
-    data=["theta_PEST_r", "phi", "phi_r", "iota", "iota_r"],
+    data=["periodic(alpha_r)", "secular(alpha_r)"],
 )
 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["periodic(alpha_r)"] + data["secular(alpha_r)"]
+    return data
+
+
+@register_compute_fun(
+    name="periodic(alpha_r)",
+    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["periodic(alpha_r)"] = data["theta_PEST_r"] - data["iota"] * data["phi_r"]
+    return data
+
+
+@register_compute_fun(
+    name="secular(alpha_r)",
+    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["secular(alpha_r)"] = -data["iota_r"] * data["phi"]
     return data
 
 

desc/compute/_metric.py

@@ -1941,24 +1941,85 @@
     # 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=["periodic(gbdrift)", "secular(gbdrift)"],
 )
 def _gbdrift(params, transforms, profiles, data, **kwargs):
-    data["gbdrift"] = (
-        1
+    data["gbdrift"] = data["periodic(gbdrift)"] + data["secular(gbdrift)"]
+    return data
+
+
+@register_compute_fun(
+    name="periodic(gbdrift)",
+    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", "periodic(grad(alpha))", "grad(|B|)"],
+)
+def _periodic_gbdrift(params, transforms, profiles, data, **kwargs):
+    data["periodic(gbdrift)"] = (
+        dot(data["b"], cross(data["grad(|B|)"], data["periodic(grad(alpha))"]))
+        / data["|B|^2"]
+    )
+    return data
+
+
+@register_compute_fun(
+    name="secular(gbdrift)",
+    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", "secular(grad(alpha))", "grad(|B|)"],
+)
+def _secular_gbdrift(params, transforms, profiles, data, **kwargs):
+    data["secular(gbdrift)"] = (
+        dot(data["b"], cross(data["grad(|B|)"], data["secular(grad(alpha))"]))
+        / data["|B|^2"]
+    )
+    return data
+
+
+@register_compute_fun(
+    name="secular(gbdrift)/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["secular(gbdrift)/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 @@
     # 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=["periodic(cvdrift)", "secular(gbdrift)"],
 )
 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["periodic(cvdrift)"] + data["secular(gbdrift)"]
+    return data
+
+
+@register_compute_fun(
+    name="periodic(cvdrift)",
+    label="\\mathrm{periodic(cvdrift)}",
+    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", "periodic(gbdrift)"],
+)
+def _periodic_cvdrift(params, transforms, profiles, data, **kwargs):
+    data["periodic(cvdrift)"] = (
+        mu_0 * data["p_r"] / data["psi_r"] / data["|B|^2"] + data["periodic(gbdrift)"]
+    )
     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 @@
         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/grid.py

@@ -43,6 +43,7 @@
         "_inverse_poloidal_idx",
         "_inverse_zeta_idx",
         "_is_meshgrid",
+        "_can_fft2",
     ]
 
     @abstractmethod
@@ -231,6 +232,16 @@
         """
         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 @@
         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 @@
         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 @@
         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 @@
             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

desc/integrals/__init__.py

@@ -1,6 +1,6 @@
 """Classes for function integration."""
 
-from .bounce_integral import Bounce1D
+from .bounce_integral import Bounce1D, Bounce2D
 from .singularities import (
     DFTInterpolator,
     FFTInterpolator,