Add fft grid and raz grid to test against master

desc/compute/_fast_ion.py

@@ -119,9 +119,9 @@ def _Gamma_c(params, transforms, profiles, data, **kwargs):
         num_pitch,
         pitch_batch_size,
         surf_batch_size,
-        quad,
         nufft_eps,
         spline,
+        quad,
         vander,
     ) = Bounce2D._defaults(-2, grid, **kwargs)
 
@@ -270,9 +270,9 @@ def _little_gamma_c_Nemov(params, transforms, profiles, data, **kwargs):
         num_pitch,
         _,
         _,
-        quad,
         nufft_eps,
         spline,
+        quad,
         vander,
     ) = Bounce2D._defaults(-2, grid, **kwargs)
 
@@ -382,9 +382,9 @@ def _Gamma_c_Velasco(params, transforms, profiles, data, **kwargs):
         num_pitch,
         pitch_batch_size,
         surf_batch_size,
-        quad,
         nufft_eps,
         spline,
+        quad,
         vander,
     ) = Bounce2D._defaults(-1, grid, **kwargs)
 

desc/compute/_neoclassical.py

@@ -20,7 +20,10 @@
         If the option ``spline`` is ``True``, the bounce points are found with
         8th order accuracy in this parameter. If the option ``spline`` is ``False``,
         then the bounce points are found with spectral accuracy in this parameter.
-        A reference value for the ``spline`` option is 100.
+        A reference value for the ``spline=True`` option is
+        ``grid.NFP*(grid.num_theta+grid.num_zeta)//2``.
+        A reference value for the ``spline=False`` option is
+        ``(grid.num_theta+grid.num_zeta)//2``.
 
         An error of ε in a bounce point manifests
         𝒪(ε¹ᐧ⁵) error in bounce integrals with (v_∥)¹ and
@@ -52,6 +55,8 @@
         A tighter upper bound than ``num_well=(Aι+C)*num_transit`` is preferable.
         The ``check_points`` or ``plot`` methods in ``desc.integrals.Bounce2D``
         are useful to select a reasonable value.
+
+        This is the most important parameter to specify for performance.
         """,
     "num_quad": """int :
         Resolution for quadrature of bounce integrals.
@@ -70,11 +75,6 @@
         If given ``None``, then ``surf_batch_size`` is ``grid.num_rho``.
         Default is ``1``. Only consider increasing if ``pitch_batch_size`` is ``None``.
         """,
-    "quad": """tuple[jnp.ndarray] :
-        Used to compute bounce integrals.
-        Quadrature points xₖ and weights wₖ for the
-        approximate evaluation of the integral ∫₋₁¹ f(x) dx ≈ ∑ₖ wₖ f(xₖ).
-        """,
     "nufft_eps": """float :
         Precision requested for interpolation with non-uniform fast Fourier
         transform (NUFFT). If less than ``1e-14`` then NUFFT will not be used.
@@ -83,6 +83,11 @@
         Whether to use cubic splines to compute bounce points instead of
         Chebyshev series. Default is ``True``.
         """,
+    "quad": """tuple[jnp.ndarray] :
+        Used to compute bounce integrals.
+        Quadrature points xₖ and weights wₖ for the
+        approximate evaluation of the integral ∫₋₁¹ f(x) dx ≈ ∑ₖ wₖ f(xₖ).
+        """,
     "_vander": """dict[str,jnp.ndarray] :
         Precomputed transform matrix "dct spline".
         This private parameter is intended to be used only by
@@ -199,9 +204,9 @@ def _epsilon_32(params, transforms, profiles, data, **kwargs):
         num_pitch,
         pitch_batch_size,
         surf_batch_size,
-        quad,
         nufft_eps,
         spline,
+        quad,
         vander,
     ) = Bounce2D._defaults(1, grid, **kwargs)
 

desc/equilibrium/equilibrium.py

@@ -946,6 +946,8 @@ def compute(  # noqa: C901
                 inbasis=[inbasis[char] for char in grid.coordinates],
                 outbasis=("rho", "theta", "zeta"),
                 period=grid.period,
+                tol=kwargs.pop("tol", 1e-6),
+                maxiter=kwargs.pop("maxiter", 30),
             )
             grid = Grid(
                 nodes=rtz_nodes,

desc/integrals/_bounce_utils.py

@@ -189,7 +189,7 @@
     return z[..., 0, :, :], z[..., 1, :, :]
 
 
-@partial(jax.custom_jvp, nondiff_argnames=("num_well", "nufft_eps"))
+@partial(jax.custom_jvp, nondiff_argnums=(2, 3))
 def regular_points(o, pitch_inv, num_well, nufft_eps):
     """Bounce points then newton, with regularized jvp."""
     return _newton(
@@ -744,15 +744,12 @@
     (ϑ, NFP ζ) coordinates, then f(ϑ(α=α₀, ζ), ζ) will sample the approximation to
     F(α=α₀ ± ε, ζ) with ε → 0 as f converges to F.
 
-    This property was mentioned because parameterizing the stream map in (α, ζ) enables
-    partial summation. However, the small discontinuity due to discretization error
-    between branch cuts is undesirable as it can give significant error to the singular
-    integrals whose integration boundary is near a branch cut. If we were using splines
-    instead of pseudo-spectral methods to interpolate then we would have to account
-    for this.
-
     """
+    X = angle.shape[-2]
+    Y = truncate_rule(angle.shape[-1])
     num_alpha = alpha.size
+    domain = (0, 2 * jnp.pi / NFP)
+
     # peeling off field lines
     alpha = get_alphas(alpha, iota, num_transit, NFP)
     if angle.ndim == 2:
@@ -762,16 +759,20 @@
     # (since this avoids modding on more points later and keeps θ bounded).
     alpha %= 2 * jnp.pi
 
-    domain = (0, 2 * jnp.pi / NFP)
-    Y = truncate_rule(angle.shape[-1])
     delta = (
         FourierChebyshevSeries(angle, domain, truncate=Y)
         .compute_cheb(alpha)
         .swapaxes(0, -3)
     )
     alpha = alpha.swapaxes(0, -2)
-    delta = delta.at[..., 0].add(alpha)
+    delta = delta.at[..., 0].add(alpha)  # This is now θ = α + δ.
     assert delta.shape == (*angle.shape[:-2], num_alpha, num_transit * NFP, Y)
+
+    if X < 24:
+        # Short circuit convergence to enforce continuity. This is needed as
+        # we assume continuity of |B| along B when gathering bounce points.
+        delta = PiecewiseChebyshevSeries.stitch(delta)
+
     return PiecewiseChebyshevSeries(delta, domain)
 
 
@@ -1066,21 +1067,27 @@
     return (Y_B * grid.NFP) if spline else Y_B
 
 
-def num_well_rule(num_transit, NFP, Y_B=None):
+def num_well_rule(num_transit, NFP, extrema_per_transit=None):
     """Guess upper bound for number of wells based on spectrum.
 
     This should be loose enough that it is equivalent to ``num_well=None``,
     but more performant.
     """
     num_well = num_transit * (20 + NFP)
-    return num_well if Y_B is None else min(num_well, num_transit * Y_B)
+    return (
+        num_well
+        if extrema_per_transit is None
+        else min(num_well, num_transit * extrema_per_transit)
+    )
 
 
-def get_vander(grid, Y, Y_B, NFP):
+def get_vander(grid, Y, Y_B, NFP, spline):
     """Builds Vandermonde matrices for objectives."""
-    Y_trunc = truncate_rule(Y)
-    Y_B, num_z = round_up_rule(Y_B, NFP, grid.num_zeta == 1)
-    x = jnp.linspace(
-        -1, 1, (Y_B // NFP) if (grid.num_zeta == 1) else num_z, endpoint=False
-    )
-    return {"dct spline": chebvander(x, Y_trunc - 1)}
+    if spline:
+        Y_trunc = truncate_rule(Y)
+        Y_B, num_z = round_up_rule(Y_B, NFP, grid.num_zeta == 1)
+        x = jnp.linspace(
+            -1, 1, (Y_B // NFP) if (grid.num_zeta == 1) else num_z, endpoint=False
+        )
+        return {"dct spline": chebvander(x, Y_trunc - 1)}
+    return {}

desc/integrals/_interp_utils.py

@@ -160,8 +160,6 @@
         s = s[..., jnp.newaxis, :] if vec else s
         f = jnp.fft.ifftshift(f, rfft_axis)
 
-    opts = options.Opts(modeord=1)
-
     # TODO: Delete this if block after
     # https://github.com/flatironinstitute/jax-finufft/pull/216 is merged
     # and then bump min version requirement.
@@ -172,6 +170,7 @@
         f = jnp.fft.fftshift(f, (-2, -1))
         return (nufft2(f, x0, x1, iflag=1, eps=eps, opts=opts) * s).real
 
+    opts = options.Opts(modeord=1)
     return (nufft2(f, x0, x1, points_mask=mask, iflag=1, eps=eps, opts=opts) * s).real
 
 
@@ -261,7 +260,7 @@
 _eps = max(jnp.finfo(jnp.array(1.0).dtype).eps, 2.5e-12)
 
 
-@partial(jax.custom_jvp, nondiff_argnames=("sort", "sentinel", "eps", "distinct"))
+@partial(jax.custom_jvp, nondiff_argnums=(4, 5, 6, 7))
 def polyroot_vec(
     c,
     k=0.0,