Non-uniform FFTs

.github/workflows/jax_tests.yml

@@ -12,10 +12,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        jax-version: [0.4.29, 0.4.30, 0.4.31, 0.4.33, 0.4.34, 0.4.35, 0.4.37,
-                      0.4.38, 0.5.0, 0.5.3, 0.6.0, 0.6.1, 0.6.2, 0.7.2]
-                      # 0.4.32 is not available on PyPI
-                      # 0.4.36 has a bug that causes tests to fail
+        jax-version: [0.5.0, 0.5.3, 0.6.0, 0.6.1, 0.6.2, 0.7.2]
                       # 0.5.1 and 0.5.2 installations are broken, see jax#26781
                       # 0.7.0 and 0.7.1 have performance issues, see diffrax#680
         group: [1, 2]
@@ -30,7 +27,7 @@ jobs:
           python -m pip install --upgrade pip
       - name: Install dependencies with given JAX version
         run: |
-          sed -i '/jax/d' ./requirements.txt
+          sed -i '1{/^jax/d}' requirements.txt
           sed -i '1i\jax[cpu] == ${{ matrix.jax-version }}' ./requirements.txt
           cat ./requirements.txt
           pip install -r ./devtools/dev-requirements.txt

CHANGELOG.md

@@ -26,6 +26,12 @@ Backend
 - ``desc.equilibrium.Equilibrium.set_initial_guess`` now sets lambda to zero for most use cases, and the docstring has been updated to be more explicit on what is done in each case.
 
 
+Performance Improvements
+
+- [Partial summation in coordinate mapping](https://github.com/PlasmaControl/DESC/pull/1826).
+- [NUFFTS](https://github.com/PlasmaControl/DESC/pull/1834) are now used by default for computing bounce integrals.
+
+
 v0.15.0
 -------
 

desc/batching.py

@@ -30,7 +30,7 @@
 )
 
 from desc.backend import jax, jnp, scan, vmap
-from desc.utils import errorif
+from desc.utils import errorif, identity
 
 try:
     from jax.extend import linear_util as lu
@@ -70,11 +70,6 @@ def _batch_and_remainder(x, batch_size: int):
         return scan_tree, remainder_tree
 
 
-def _identity(y):
-    """Returns the input."""
-    return y
-
-
 _unchunk = partial(tree_map, lambda y: y.reshape(-1, *y.shape[2:]))
 _concat = partial(tree_map, lambda y1, y2: jnp.concatenate((y1, y2)))
 _get_first_chunk = partial(tree_map, lambda x: x[0])
@@ -106,7 +101,7 @@ def body(carry, x):
     return result
 
 
-def _scanmap(fun, argnums=0, reduction=None, chunk_reduction=_identity):
+def _scanmap(fun, argnums=0, reduction=None, chunk_reduction=identity):
     """A helper function to wrap f with a scan_fun.
 
     Refrences
@@ -141,7 +136,7 @@ def _evaluate_in_chunks(
     chunk_size,
     argnums,
     reduction=None,
-    chunk_reduction=_identity,
+    chunk_reduction=identity,
     *args,
     **kwargs,
 ):
@@ -192,7 +187,7 @@ def vmap_chunked(
     *,
     chunk_size=None,
     reduction=None,
-    chunk_reduction=_identity,
+    chunk_reduction=identity,
 ):
     """Behaves like ``vmap`` but uses scan to chunk the computations in smaller chunks.
 
@@ -232,7 +227,7 @@ def vmap_chunked(
 
 
 def batch_map(
-    fun, fun_input, /, batch_size=None, *, reduction=None, chunk_reduction=_identity
+    fun, fun_input, /, batch_size=None, *, reduction=None, chunk_reduction=identity
 ):
     """Compute ``chunk_reduction(fun(fun_input))`` in batches.
 

desc/compute/__init__.py

@@ -33,6 +33,7 @@
     _core,
     _curve,
     _equil,
+    _fast_ion,
     _field,
     _geometry,
     _metric,

desc/compute/_fast_ion.py

@@ -100,6 +100,7 @@ def _drift2(data, B, pitch):
         "num_pitch",
         "pitch_batch_size",
         "surf_batch_size",
+        "nufft_eps",
         "spline",
     ],
 )
@@ -135,7 +136,6 @@ def _Gamma_c(params, transforms, profiles, data, **kwargs):
     assert (
         surf_batch_size == 1 or pitch_batch_size is None
     ), f"Expected pitch_batch_size to be None, got {pitch_batch_size}."
-    spline = kwargs.get("spline", True)
     fl_quad = (
         kwargs["fieldline_quad"] if "fieldline_quad" in kwargs else leggauss(Y_B // 2)
     )
@@ -147,6 +147,9 @@ def _Gamma_c(params, transforms, profiles, data, **kwargs):
             (automorphism_sin, grad_automorphism_sin),
         )
     )
+    nufft_eps = kwargs.get("nufft_eps", 1e-7)
+    spline = kwargs.get("spline", True)
+    vander = kwargs.get("_vander", None)
 
     def Gamma_c(data):
         bounce = Bounce2D(
@@ -157,8 +160,10 @@ def Gamma_c(data):
             alpha,
             num_transit,
             quad,
+            nufft_eps=nufft_eps,
             is_fourier=True,
             spline=spline,
+            vander=vander,
         )
 
         def fun(pitch_inv):
@@ -169,6 +174,7 @@ def fun(pitch_inv):
                 data,
                 ["|grad(psi)|*kappa_g", "|B|_r|v,p", "K"],
                 points,
+                nufft_eps=nufft_eps,
                 is_fourier=True,
             )
             # This is γ_c π/2.
@@ -177,7 +183,10 @@ def fun(pitch_inv):
                     drift1,
                     drift2
                     * bounce.interp_to_argmin(
-                        data["|grad(rho)|*|e_alpha|r,p|"], points, is_fourier=True
+                        data["|grad(rho)|*|e_alpha|r,p|"],
+                        points,
+                        nufft_eps=nufft_eps,
+                        is_fourier=True,
                     ),
                 )
             )
@@ -188,14 +197,14 @@ def fun(pitch_inv):
             * data["pitch_inv weight"]
             / data["pitch_inv"] ** 2,
             axis=-1,
-        ) / (bounce.compute_fieldline_length(fl_quad) * 2**1.5 * jnp.pi)
+        ) / (bounce.compute_fieldline_length(fl_quad, vander) * 2**1.5 * jnp.pi)
 
     # It is assumed the grid is sufficiently dense to reconstruct |B|,
     # so anything smoother than |B| may be captured accurately as a single
-    # Fourier series rather than transforming each component.
-    # Last term in K behaves as ∂log(|B|²/B^ϕ)/∂ρ |B| if one ignores the issue
-    # of a log argument with units. Smoothness determined by positive lower bound
-    # of log argument, and hence behaves as ∂log(|B|)/∂ρ |B| = ∂|B|/∂ρ.
+    # Fourier series rather than transforming each component. Last term in K
+    # behaves as ∂log(|B|²/(R₀B₀B^ϕ))/∂ρ |B| where R₀B₀ is a constant with
+    # units Tesla meters. Smoothness is determined by positive lower bound of
+    # log argument, and hence behaves as ∂log(|B|/B₀)/∂ρ |B| = ∂|B|/∂ρ.
     fun_data = {
         "|grad(psi)|*kappa_g": data["|grad(psi)|"] * data["kappa_g"],
         "|grad(rho)|*|e_alpha|r,p|": data["|grad(rho)|"] * data["|e_alpha|r,p|"],
@@ -266,6 +275,7 @@ def _poloidal_drift(data, B, pitch):
         "num_pitch",
         "pitch_batch_size",
         "surf_batch_size",
+        "nufft_eps",
         "spline",
     ],
 )
@@ -298,7 +308,9 @@ def _little_gamma_c_Nemov(params, transforms, profiles, data, **kwargs):
             (automorphism_sin, grad_automorphism_sin),
         )
     )
+    nufft_eps = kwargs.get("nufft_eps", 1e-7)
     spline = kwargs.get("spline", True)
+    vander = kwargs.get("_vander", None)
 
     def gamma_c0(data):
         bounce = Bounce2D(
@@ -309,8 +321,10 @@ def gamma_c0(data):
             alpha,
             num_transit,
             quad,
+            nufft_eps=nufft_eps,
             is_fourier=True,
             spline=spline,
+            vander=vander,
         )
 
         def fun(pitch_inv):
@@ -321,6 +335,7 @@ def fun(pitch_inv):
                 data,
                 ["|grad(psi)|*kappa_g", "|B|_r|v,p", "K"],
                 points,
+                nufft_eps=nufft_eps,
                 is_fourier=True,
             )
             return (2 / jnp.pi) * jnp.arctan(
@@ -330,6 +345,7 @@ def fun(pitch_inv):
                     * bounce.interp_to_argmin(
                         data["|grad(rho)|*|e_alpha|r,p|"],
                         points,
+                        nufft_eps=nufft_eps,
                         is_fourier=True,
                     ),
                 )
@@ -397,6 +413,7 @@ def fun(pitch_inv):
         "num_pitch",
         "pitch_batch_size",
         "surf_batch_size",
+        "nufft_eps",
         "spline",
     ],
 )
@@ -427,7 +444,6 @@ def _Gamma_c_Velasco(params, transforms, profiles, data, **kwargs):
     assert (
         surf_batch_size == 1 or pitch_batch_size is None
     ), f"Expected pitch_batch_size to be None, got {pitch_batch_size}."
-    spline = kwargs.get("spline", True)
     fl_quad = (
         kwargs["fieldline_quad"] if "fieldline_quad" in kwargs else leggauss(Y_B // 2)
     )
@@ -439,6 +455,9 @@ def _Gamma_c_Velasco(params, transforms, profiles, data, **kwargs):
             (automorphism_sin, grad_automorphism_sin),
         )
     )
+    nufft_eps = kwargs.get("nufft_eps", 1e-7)
+    spline = kwargs.get("spline", True)
+    vander = kwargs.get("_vander", None)
 
     def Gamma_c(data):
         bounce = Bounce2D(
@@ -449,8 +468,10 @@ def Gamma_c(data):
             alpha,
             num_transit,
             quad,
+            nufft_eps=nufft_eps,
             is_fourier=True,
             spline=spline,
+            vander=vander,
         )
 
         def fun(pitch_inv):
@@ -459,7 +480,8 @@ def fun(pitch_inv):
                 pitch_inv,
                 data,
                 ["cvdrift0", "gbdrift (periodic)", "gbdrift (secular)/phi"],
-                bounce.points(pitch_inv, num_well),
+                num_well=num_well,
+                nufft_eps=nufft_eps,
                 is_fourier=True,
             )
             # This is γ_c π/2.
@@ -471,7 +493,7 @@ def fun(pitch_inv):
             * data["pitch_inv weight"]
             / data["pitch_inv"] ** 2,
             axis=-1,
-        ) / (bounce.compute_fieldline_length(fl_quad) * 2**1.5 * jnp.pi)
+        ) / (bounce.compute_fieldline_length(fl_quad, vander) * 2**1.5 * jnp.pi)
 
     grid = transforms["grid"]
 

desc/compute/_metric.py

@@ -2626,7 +2626,7 @@ def _B_dot_grad_grad_rho(params, transforms, profiles, data, **kwargs):
 @register_compute_fun(
     name="finite-n instability drive",
     label="(\\mathbf{J} \\times (\\nabla \\rho))/{(g^{\\rho \\rho})}^2"
-    + " \\mathbf{B} \\cdot \\cdot \\mathbf{\\nabla} (\\mathbf{\\nabla} \\rho)",
+    + " \\mathbf{B} \\cdot \\mathbf{\\nabla} (\\mathbf{\\nabla} \\rho)",
     units="T A \\cdot m^{-1}",
     units_long="Tesla Amperes / meter",
     description="finite-n instability drive term",