Add checkpointing

Author: unalmis

desc/compute/_fast_ion.py

@@ -125,7 +125,6 @@ def _Gamma_c(params, transforms, profiles, data, **kwargs):
         nufft_eps,
         spline,
         vander,
-        low_ram,
     ) = Bounce2D._default_kwargs("weak", grid.NFP, **kwargs)
 
     def Gamma_c(data):
@@ -152,7 +151,6 @@ def fun(pitch_inv):
                 ["|grad(psi)|*kappa_g", "|B|_r|v,p", "K"],
                 points,
                 nufft_eps=nufft_eps,
-                low_ram=low_ram,
                 is_fourier=True,
             )
             # This is γ_c π/2.
@@ -278,7 +276,6 @@ def _little_gamma_c_Nemov(params, transforms, profiles, data, **kwargs):
         nufft_eps,
         spline,
         vander,
-        low_ram,
     ) = Bounce2D._default_kwargs("weak", grid.NFP, **kwargs)
 
     def gamma_c0(data):
@@ -305,7 +302,6 @@ def fun(pitch_inv):
                 ["|grad(psi)|*kappa_g", "|B|_r|v,p", "K"],
                 points,
                 nufft_eps=nufft_eps,
-                low_ram=low_ram,
                 is_fourier=True,
             )
             return (2 / jnp.pi) * jnp.arctan(
@@ -402,7 +398,6 @@ def _Gamma_c_Velasco(params, transforms, profiles, data, **kwargs):
         nufft_eps,
         spline,
         vander,
-        low_ram,
     ) = Bounce2D._default_kwargs("weak", grid.NFP, **kwargs)
 
     def Gamma_c(data):
@@ -428,7 +423,6 @@ def fun(pitch_inv):
                 ["cvdrift0", "gbdrift (periodic)", "gbdrift (secular)/phi"],
                 num_well=num_well,
                 nufft_eps=nufft_eps,
-                low_ram=low_ram,
                 is_fourier=True,
             )
             # This is γ_c π/2.

desc/compute/_neoclassical.py

@@ -80,10 +80,6 @@
         This private parameter is intended to be used only by
         developers for objectives.
         """,
-    "low_ram": """bool :
-        If true, then will switch to a slower algorithm whose differentiation
-        consumes less memory. Default is false.
-        """,
     "theta": "",
 }
 
@@ -97,7 +93,6 @@
     "surf_batch_size",
     "nufft_eps",
     "spline",
-    "low_ram",
 )
 
 
@@ -201,7 +196,6 @@ def _epsilon_32(params, transforms, profiles, data, **kwargs):
         nufft_eps,
         spline,
         vander,
-        low_ram,
     ) = Bounce2D._default_kwargs("deriv", grid.NFP, **kwargs)
 
     def eps_32(data):
@@ -231,7 +225,6 @@ def fun(pitch_inv):
                 "|grad(rho)|*kappa_g",
                 num_well=num_well,
                 nufft_eps=nufft_eps,
-                low_ram=low_ram,
                 is_fourier=True,
             )
             return safediv(I_1**2, I_2).sum(-1).mean(-2)

desc/integrals/_bounce_utils.py

@@ -21,14 +21,14 @@
 from matplotlib import pyplot as plt
 from orthax.chebyshev import chebroots, chebvander
 
-from desc.backend import dct, flatnonzero, fori_loop, idct, ifft, jnp
+from desc.backend import dct, flatnonzero, fori_loop, idct, ifft, jax, jnp
 from desc.integrals._interp_utils import (
     _eps,
     _filter_distinct,
     _subtract_first,
+    cubic_val,
     nufft1d2r,
     polyroot_vec,
-    polyval_vec,
 )
 from desc.integrals.quad_utils import bijection_from_disc, bijection_to_disc
 from desc.io import IOAble
@@ -97,7 +97,7 @@ def _in_epigraph_and(is_intersect, df_dy, /):
     return is_intersect.at[idx[0]].set(edge_case)
 
 
-def bounce_points(pitch_inv, knots, B, dB_dz, num_well=None):
+def bounce_points(pitch_inv, knots, B, num_well=None):
     """Compute the bounce points given 1D spline of B and pitch λ.
 
     Parameters
@@ -109,15 +109,10 @@ def bounce_points(pitch_inv, knots, B, dB_dz, num_well=None):
         Shape (N, ).
         ζ coordinates of spline knots. Must be strictly increasing.
     B : jnp.ndarray
-        Shape (..., N - 1, B.shape[-1]).
+        Shape (..., N - 1, 4).
         Polynomial coefficients of the spline of B in local power basis.
         Last axis enumerates the coefficients of power series. Second to
         last axis enumerates the polynomials that compose a particular spline.
-    dB_dz : jnp.ndarray
-        Shape (..., N - 1, B.shape[-1] - 1).
-        Polynomial coefficients of the spline of (∂B/∂ζ)|(ρ,α) in local power basis.
-        Last axis enumerates the coefficients of power series. Second to
-        last axis enumerates the polynomials that compose a particular spline.
     num_well : int or None
         Specify to return the first ``num_well`` pairs of bounce points for each
         pitch and field line. Choosing ``-1`` will detect all wells, but due
@@ -145,8 +140,9 @@ def bounce_points(pitch_inv, knots, B, dB_dz, num_well=None):
         line and pitch, is padded with zero.
 
     """
+    B = B[..., None, :, :]
     intersect = polyroot_vec(
-        c=B[..., None, :, :],
+        c=B,
         k=jnp.atleast_1d(pitch_inv)[..., None],
         a_min=jnp.array([0.0]),
         a_max=jnp.diff(knots),
@@ -156,9 +152,7 @@ def bounce_points(pitch_inv, knots, B, dB_dz, num_well=None):
     )
     assert intersect.shape[-2:] == (knots.size - 1, B.shape[-1] - 1)
 
-    dB_dz = flatten_mat(
-        jnp.sign(polyval_vec(x=intersect, c=dB_dz[..., None, :, None, :]))
-    )
+    dB_dz = flatten_mat(jnp.sign(cubic_val(x=intersect, c=B[..., None, :], der=True)))
     # Only consider intersect if it is within knots that bound that polynomial.
     mask = flatten_mat(intersect >= 0)
     z1 = (dB_dz <= 0) & mask
@@ -532,7 +526,7 @@ def _plot_intersect(
         )
 
 
-def get_extrema(knots, g, dg_dz, sentinel=jnp.nan):
+def get_extrema(knots, g, sentinel=jnp.nan):
     """Return extrema (z*, g(z*)).
 
     Parameters
@@ -541,30 +535,28 @@ def get_extrema(knots, g, dg_dz, sentinel=jnp.nan):
         Shape (N, ).
         ζ coordinates of spline knots. Must be strictly increasing.
     g : jnp.ndarray
-        Shape (..., N - 1, g.shape[-1]).
+        Shape (..., N - 1, 4).
         Polynomial coefficients of the spline of g in local power basis.
         Last axis enumerates the coefficients of power series. Second to
         last axis enumerates the polynomials that compose a particular spline.
-    dg_dz : jnp.ndarray
-        Shape (..., N - 1, g.shape[-1] - 1).
-        Polynomial coefficients of the spline of ∂g/∂z in local power basis.
-        Last axis enumerates the coefficients of power series. Second to
-        last axis enumerates the polynomials that compose a particular spline.
     sentinel : float
         Value with which to pad array to return fixed shape.
 
     Returns
     -------
     ext, g_ext : jnp.ndarray
-        Shape (..., (N - 1) * (g.shape[-1] - 2)).
+        Shape (..., (N - 1) * 2).
         First array enumerates z*. Second array enumerates g(z*)
         Sorting order of extrema is arbitrary.
 
     """
     ext = polyroot_vec(
-        c=dg_dz, a_min=jnp.array([0.0]), a_max=jnp.diff(knots), sentinel=sentinel
+        c=g[..., :-1] * jnp.arange(g.shape[-1] - 1, 0, -1),
+        a_min=jnp.array([0.0]),
+        a_max=jnp.diff(knots),
+        sentinel=sentinel,
     )
-    g_ext = flatten_mat(polyval_vec(x=ext, c=g[..., None, :]))
+    g_ext = flatten_mat(cubic_val(x=ext, c=g[..., None, :]))
     # Transform out of local power basis expansion.
     ext = flatten_mat(ext + knots[:-1, None])
     assert ext.shape == g_ext.shape
@@ -991,23 +983,6 @@ def round_up_rule(Y, NFP, axisymmetric=False):
     return num_z * NFP, num_z
 
 
-def fieldline_quad_rule(Y):
-    """Ensure field line quadrature has reasonable resolution.
-
-    Parameters
-    ----------
-    Y : int
-        Resolution of Chebyshev spectrum of angle over one field period.
-
-    Returns
-    -------
-    Y : int
-        Resolution for Gauss-Legendre quadrature over one field period.
-
-    """
-    return max(Y, 8)
-
-
 def Y_B_rule(Y, NFP, spline=True):
     """Guess Y_B from resolution of Chebyshev spectrum of angle."""
     return (2 * Y * int(np.sqrt(NFP))) if spline else Y
@@ -1037,6 +1012,45 @@ def get_vander(grid, Y, Y_B, NFP):
     return {"dct spline": chebvander(x, Y_trunc - 1)}
 
 
+@jax.checkpoint
+def _gather_reduce(y, cheb, x_idx):
+    """Gather then reduce with checkpointing.
+
+    Checkpointing this makes it faster and reduces memory.
+    On many architectures, the cost of allocating contiguous blocks of memory
+    exceeds the cost of flops. By checkpointing, we avoid that in favor of
+    recomputing derivatives in the backward pass.
+
+    Lies:
+    https://docs.jax.dev/en/latest/notebooks/autodiff_remat.html#practical-notes
+    """
+    return idct_mmt(y, jnp.take_along_axis(cheb, x_idx[..., None], axis=-2))
+
+
+@jax.checkpoint
+def _loop(y, cheb, x_idx):
+    """Memory efficient Clenshaw recursion.
+
+    Checkpointing on a CPU observed a minor reduction in memory usage while not
+    affecting speed. JAX/XLA has poor performance with iterative algorithms compared
+    to languages like Julia. On JAX version 0.7.2, this is slower than the product sum
+    reduction above. Without checkpointing either, this uses signficantly less memory
+    than above.
+    """
+
+    def body(i, val):
+        c0, c1 = val
+        return jnp.take_along_axis(cheb[-i], x_idx, axis=-1) - c1, c0 + c1 * y2
+
+    num_coef = cheb.shape[-1]
+    cheb = jnp.moveaxis(cheb, -1, 0)  # to minimize cache misses
+    y2 = 2 * y
+    c0 = jnp.take_along_axis(cheb[-2], x_idx, axis=-1)
+    c1 = jnp.take_along_axis(cheb[-1], x_idx, axis=-1)
+    c0, c1 = fori_loop(3, num_coef + 1, body, (c0, c1))
+    return c0 + c1 * y
+
+
 class PiecewiseChebyshevSeries(IOAble):
     """Chebyshev series.
 
@@ -1167,8 +1181,10 @@ def eval1d(self, z, cheb=None, loop=False):
             Shape (..., X, Y).
             Chebyshev coefficients to use. If not given, uses ``self.cheb``.
         loop : bool
-            Whether to use Clenshaw recursion.
-            This is slower on CPU, but it reduces memory of the Jacobian.
+            If ``True``, then uses Clenshaw recursion which is memory efficient.
+            If ``False``, then gathers a large block of memory and computes
+            a product sum reduction while checkpointing the derivative
+            to reduce memory consumption of the Jacobian.
 
         Returns
         -------
@@ -1178,26 +1194,8 @@ def eval1d(self, z, cheb=None, loop=False):
         """
         cheb = setdefault(cheb, self.cheb)
         x_idx, y = self._isomorphism_to_C2(z)
-
         y = bijection_to_disc(y, self.domain[0], self.domain[-1])
-
-        # Recall that the Chebyshev coefficients αₙ for f(z) = ∑ₙ₌₀ᴺ⁻¹ αₙ(x[z]) Tₙ(y[z])
-        # are in cheb array whose shape is (..., num cheb series, spectral resolution).
-
-        if not loop or self.Y < 3:
-            cheb = jnp.take_along_axis(cheb, x_idx[..., None], axis=-2)
-            return idct_mmt(y, cheb)
-
-        def body(i, val):
-            c0, c1 = val
-            return jnp.take_along_axis(cheb[-i], x_idx, axis=-1) - c1, c0 + c1 * y2
-
-        cheb = jnp.moveaxis(cheb, -1, 0)  # to leverage cache
-        y2 = 2 * y
-        c0 = jnp.take_along_axis(cheb[-2], x_idx, axis=-1)
-        c1 = jnp.take_along_axis(cheb[-1], x_idx, axis=-1)
-        c0, c1 = fori_loop(3, self.Y + 1, body, (c0, c1))
-        return c0 + c1 * y
+        return (_loop if (loop and self.Y >= 3) else _gather_reduce)(y, cheb, x_idx)
 
     def intersect2d(self, k=0.0, *, eps=_eps):
         """Coordinates yᵢ such that f(x, yᵢ) = k(x).

desc/integrals/_interp_utils.py

@@ -161,32 +161,8 @@ def interp1d_Hermite_vec(xq, x, f, fx, /):
     return interp1d(xq, x, f, method="cubic", fx=fx)
 
 
-# TODO (#1388): Move to interpax.
-
-
-def polyder_vec(c):
-    """Coefficients for the derivatives of the given set of polynomials.
-
-    Parameters
-    ----------
-    c : jnp.ndarray
-        Last axis should store coefficients of a polynomial. For a polynomial given by
-        ∑ᵢⁿ cᵢ xⁱ, where n is ``c.shape[-1]-1``, coefficient cᵢ should be stored at
-        ``c[...,n-i]``.
-
-    Returns
-    -------
-    poly : jnp.ndarray
-        Coefficients of polynomial derivative, ignoring the arbitrary constant. That is,
-        ``poly[...,i]`` stores the coefficient of the monomial xⁿ⁻ⁱ⁻¹,  where n is
-        ``c.shape[-1]-1``.
-
-    """
-    return c[..., :-1] * jnp.arange(c.shape[-1] - 1, 0, -1)
-
-
-def polyval_vec(*, x, c):
-    """Evaluate the set of polynomials ``c`` at the points ``x``.
+def cubic_val(*, x, c, der=False):
+    """Evaluate the derivative of cubic polynomial ``c`` at the points ``x``.
 
     Parameters
     ----------
@@ -196,6 +172,8 @@ def polyval_vec(*, x, c):
         Last axis should store coefficients of a polynomial. For a polynomial given by
         ∑ᵢⁿ cᵢ xⁱ, where n is ``c.shape[-1]-1``, coefficient cᵢ should be stored at
         ``c[...,n-i]``.
+    der : bool
+        Whether to evaluate the derivative instead.
 
     Returns
     -------
@@ -207,17 +185,16 @@ def polyval_vec(*, x, c):
     .. code-block:: python
 
         np.testing.assert_allclose(
-            polyval_vec(x=x, c=c),
+            cubic_val(x=x, c=c),
             np.sum(polyvander(x, c.shape[-1] - 1) * c[..., ::-1], axis=-1),
         )
 
     """
-    # Better than Horner's method as we expect to evaluate low order polynomials.
-    # No need to use fast multipoint evaluation techniques for the same reason.
-    return jnp.sum(
-        c * x[..., jnp.newaxis] ** jnp.arange(c.shape[-1] - 1, -1, -1),
-        axis=-1,
-    )
+    assert c.shape[-1] == 4
+    if der:
+        return (3 * c[..., 0] * x + 2 * c[..., 1]) * x + c[..., 2]
+    else:
+        return ((c[..., 0] * x + c[..., 1]) * x + c[..., 2]) * x + c[..., 3]
 
 
 def _subtract_first(c, k):
@@ -265,6 +242,8 @@ def _filter_distinct(r, sentinel, eps):
 )
 _eps = max(jnp.finfo(jnp.array(1.0).dtype).eps, 2.5e-12)
 
+# TODO (#1388): Move to interpax.
+
 
 def polyroot_vec(
     c,