Sparse pullback for big performance gain

.github/dependabot.yml

@@ -42,6 +42,7 @@ updates:
           - "orthax"
           - "quadax"
           - "jax-finufft"
+          - "adv-jax-math"
   - package-ecosystem: "github-actions"
     directory: "/"
     schedule:

CHANGELOG.md

@@ -4,9 +4,16 @@ Changelog
 Performance Improvements
 
 - Improves memory management to reduce the base memory used during optimization while using `lsq-exact`, `lsq-auglag` and `fmin-auglag` optimizers.
+- Sparse reverse-mode differentiation was introduced to yield significant performance improvements [#2170](https://github.com/PlasmaControl/DESC/pull/2170). Plumbing to use this method was added that will be progressively taken advantage of in the future.
 - Speeds up ``field_line_integrate`` and ``trace_particles`` for filamentary coils (``Coil``, ``CoilSet``, ``MixedCoilSet``) by precomputing the constant source information, so that the ODE right hand side only evaluates a single fused Biot-Savart kernel instead of recomputing the coil geometry at every solver step.
 - Improves the non-singular Biot-Savart kernel which should give a speed/memory improvement to objectives that compute magnetic field from coils such as ``QuadraticFlux``.
 
+Breaking Changes and Deprecations
+
+- The parameter ``num_transit`` in ``EffectiveRipple``, ``Gamma_c``, ``Bounce2D`` and related functions has been changed to ``field_period_transits``. This should make using a consistent resolution across different equilibria easier. The now-deprecated ``num_transit`` may still be used but note the equivalence ``field_period_transits = num_transit * grid.NFP``.
+- The parameter ``Y_B`` in ``EffectiveRipple``, ``Gamma_c``, ``Bounce2D`` is now the resolution over a single field period rather than a full toroidal transit. This should make using a consistent resolution across different equilibria easier.
+- Objectives using ``Bounce2D`` now do not support fwd mode differentiation for JAX versions <0.11.0.
+
 Bug Fixes
 
 - Fixes bug in ``auglag`` optimizers which prevented them from accepting solver hyperparameters.
@@ -16,6 +23,7 @@ Bug Fixes
   ``fmin-auglag`` with the default ``tr_method="exact"``, and in
   ``lsq-exact``/``lsq-auglag`` with ``tr_method="cho"``.
 
+
 v0.17.3
 -------
 
@@ -44,11 +52,11 @@ Bug Fixes
     - Updates ``"reactor_QA"`` in ``desc.examples`` to fix this. Note that if using ``"reactor_QA"`` example from ``v0.16.0`` until this fix, the current profile in that example has this issue.
 - Fixes bug in `CoilSet.from_symmetry` that ignored the passed in `check_intersection` value. This caused redundant checks in various other functions such as `plot_coils`.
 
-
 Breaking Changes
 
 - Name change in `_CoilObjective` replacing `coilset_mask` with `objective_mask`. Custom subclasses with `_broadcast_input="node"` that previously used `coilset_mask` should switch to `objective_mask`.
 
+
 v0.17.2
 -------
 
@@ -58,7 +66,6 @@ New Features
 - Sub-objectives of an `ObjectiveFunction` can now have different `use_jit` values than the `ObjectiveFunction`. These objectives have to be built before building the `ObjectiveFunction`.
 - Adds ``num_neighbors`` parameter to ``CoilSetMinDistance`` that limits the pairwise distance computation to the nearest neighbors per coil, reducing memory useage for large coilsets.
 - Method to plot frequency spectrum of inverse stream map in field line coordinates ``Bounce2D.plot_angle_spectrum``.
-- Method to compute bounce integrals in batches is now added to the public API ``Bounce2D.batch``.
 - Initiated deprecation of ``Bounce2D.compute_fieldline_length`` in favor of ``eq.compute("V_psi")``.
 - The quadrature resolution in ``Bounce2D.compute_fieldline_length`` now corresponds to the resolution over a single field period instead of the resolution over a toroidal transit.
 - Adds an optional attribute `ion_density` to the `Equilibrium` class, to allow the ion density profile to be set independently of the electron density and effective atomic number. Also adds compute functions for ``"ni_rr"`` and ``"Zeff_rr"``.
@@ -68,6 +75,7 @@ New Features
 Bug Fixes
 
 - Fixes SyntaxError thrown when loading hdf5 data from file-like objects.
+- Fixes ``pitch_batch_size`` argument getting ignored in compute functions.
 - Fixes a bug in `OmnigenousField.change_resolution` when changing `L_B`.
 - Scaling a `ScaledProfile` or taking power of a `PowerProfile` now only updates the `scale`/`power` attributes instead of nesting the `ScaledProfile`/`PowerProfile`s.
 - `jax.Array`s in `_static_attrs` will be automatically converted to `np.ndarray` to prevent stalling code. In general, jax arrays should be omitted in `_static_attrs`.
@@ -79,7 +87,6 @@ Performance Improvements
     - Now, `desc.compute._build_data_index` uses depth-first search algorithm to construct the dependency tree.
     - Some of the default value computations at import time are removed (i.e. `desc.integrals.bounce_integral.default_quad`)
 - [Significantly improves convergence of inverse stream maps](https://github.com/PlasmaControl/DESC/pull/1919).
-- Check-pointing to bounce integrals to improve speed and reduce memory of reverse mode differentiation.
 - Resolves a JAX memory regression in bounce integrals by avoiding materialization of a large tensor in memory. Previously, we had closed the issue by adding nuffts as a workaround. This update actually solves the issue for the case when a user specifies to not use nuffts as well.
 - ``ObjectiveFunction.print_value`` can now use the previously computed ``compute_scaled_error`` values to print. For bounded objectives, we fall back to computing ``compute_unscaled``. Additionally, ``compute_scaled_error`` and array splitting are used in other parts of the code to prevent recompilation for one-time tasks, which makes initialization faster.
 

desc/batching.py

@@ -2,6 +2,7 @@
 
 from functools import partial
 
+from adv_jax_math import batch_map  # noqa: F401
 from jax._src.api import (
     _check_input_dtype_jacfwd,
     _check_input_dtype_jacrev,
@@ -197,6 +198,21 @@ def vmap_chunked(
     - https://github.com/jax-ml/jax/issues/26689
     - https://github.com/jax-ml/jax/issues/27591
     - https://github.com/jax-ml/jax/issues/31919
+    - Due to an actively worked on issue in JAX,
+      https://docs.jax.dev/en/latest/jep/
+      2026-custom-derivatives.html#main-problem-descriptions,
+      this function can simply ignore custom derivative rules
+      of the function in wraps if ``chunk_size`` is not ``None``,
+      and therefore can damp the effeciency gains of ``sparse_pullback``.
+      Use ``batch_map`` instead to avoid this,
+      or try to make a hack with jax.custom_transforms to bypass this.
+    - Only out axes = 0 is supported.
+
+    See Also
+    --------
+    batch_map
+        If the function supports native vectorization, use ``batch_map`` instead
+        for the reasons discussed the docstring.
 
     Parameters
     ----------
@@ -230,49 +246,6 @@ def vmap_chunked(
     )
 
 
-def batch_map(
-    fun, fun_input, /, batch_size=None, *, reduction=None, chunk_reduction=identity
-):
-    """Compute ``chunk_reduction(fun(fun_input))`` in batches.
-
-    This utility is like ``vmap_chunked`` except that ``fun`` is assumed to be
-    vectorized natively. No JAX vectorization such as ``vmap`` is applied to the
-    supplied function. This makes compilation faster and avoids the weaknesses of
-    applying JAX vectorization, such as executing all branches of code conditioned on
-    dynamic values. For example, this function would be useful for GitHub issue #1303.
-
-    Parameters
-    ----------
-    fun : callable
-        Natively vectorized function.
-    fun_input : pytree
-        Data to split into batches to feed to ``fun``.
-    batch_size : int or None
-        Size of batches. If no batching should be done or the batch size is the
-        full input then supply ``None``.
-    reduction : callable or None
-        Binary reduction operation.
-        Should take two arguments and return one output, e.g. ``jnp.add``.
-    chunk_reduction : callable
-        Chunk-wise reduction operation.
-        Should typically apply ``reduction`` along the mapped axis,
-        e.g. ``jnp.add.reduce``.
-
-    Returns
-    -------
-    fun_output
-        Returns ``chunk_reduction(fun(fun_input))``.
-
-    """
-    return (
-        chunk_reduction(fun(fun_input))
-        if batch_size is None
-        else _evaluate_in_chunks(
-            fun, batch_size, (0,), reduction, chunk_reduction, fun_input
-        )
-    )
-
-
 def batched_vectorize(pyfunc, *, excluded=frozenset(), signature=None, chunk_size=None):
     """Define a vectorized function with broadcasting and batching.