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,18 +75,18 @@
         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.
         """,
     "spline": """bool :
-        Whether to use cubic splines to compute bounce points instead of
-        Chebyshev series. Default is ``True``.
+        Whether to use cubic splines to compute initial guess for 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".
@@ -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/external/neo.py

@@ -26,8 +26,8 @@ def __init__(self, name, eq, ns=256, M_booz=None, N_booz=None):
 
         self.eq = eq
         self.ns = ns  # number of surfaces
-        self.M_booz = setdefault(M_booz, 5 * eq.M + 1)
-        self.N_booz = setdefault(N_booz, 5 * eq.N)
+        self.M_booz = setdefault(M_booz, 6 * eq.M + 1)
+        self.N_booz = setdefault(N_booz, 6 * eq.N)
 
     @staticmethod
     def read(name):
@@ -39,12 +39,36 @@ def read(name):
         neo_eps[~good] = np.interp(neo_rho[~good], neo_rho[good], neo_eps[good])
         return neo_rho, neo_eps
 
-    def write(self, **kwargs):
-        """Write neo input file."""
+    def write(self, *, theta_n=300, phi_n=300, num_pitch=150, nstep_per=500, **kwargs):
+        """Write neo input file.
+
+        Parameters
+        ----------
+        theta_n, phi_n : int
+            Spline resolution.
+        num_pitch : int
+            Resolution for quadrature over velocity coordinate.
+            Default is 150.
+        nstep_per : int
+            Runge-Kutta steps to perform within each field period. Default is 500.
+
+            NEO employs a constant step size (2π/NFP/``nstep_per``),
+            explicit Runge-Kutta scheme which has algebraic convergence of order 3/2
+            for the bounce integrals of interest. The accuracy of the bounce points is
+            also only found with first order accuracy in (2π/NFP/``nstep_per``) so the
+            algebraic convergence is capped at 3/2 independent of the integrator.
+
+        """
         print(f"Writing VMEC wout to {self.vmec_file}")
         VMECIO.save(self.eq, self.vmec_file, surfs=self.ns, verbose=0)
         self._write_booz()
-        self._write_neo(**kwargs)
+        self._write_neo(
+            theta_n=theta_n,
+            phi_n=phi_n,
+            num_pitch=num_pitch,
+            nstep_per=nstep_per,
+            **kwargs,
+        )
 
     def _write_booz(self):
         print(f"Writing boozer output file to {self.booz_file}")
@@ -59,17 +83,18 @@ def _write_booz(self):
 
     def _write_neo(
         self,
-        theta_n=200,
-        phi_n=200,
-        num_pitch=150,
+        theta_n,
+        phi_n,
+        num_pitch,
+        nstep_per,
         multra=1,
         acc_req=0.01,
         nbins=100,
-        nstep_per=75,
         nstep_min=500,
         nstep_max=2000,
         verbose=2,
     ):
+        """Write NEO file."""
         print(f"Writing NEO input file to {self.neo_in_file}")
         f = open(self.neo_in_file, "w")
 

desc/integrals/__init__.py

@@ -1,6 +1,5 @@
 """Classes for function integration."""
 
-from ._bounce_utils import fast_chebyshev, fast_cubic_spline
 from ._interp_utils import nufft1d2r, nufft2d2r
 from .bounce_integral import Bounce1D, Bounce2D
 from .singularities import (