Improve import time

.gitignore

@@ -30,3 +30,5 @@ uv.lock
 # Environments
 .env
 .venv
+
+.vscode/

CHANGELOG.md

@@ -4,22 +4,38 @@ Changelog
 New Features
 
 - 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.
 
 Bug Fixes
 
-- Fixes incorrect units in the documentation of some curvature variables.
 - Fixes SyntaxError thrown when loading hdf5 data from file-like objects.
 
+Performance Improvements
+
+- Reduces import time of `desc` modules.
+    - 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.
+
+
+v0.17.1
+-------
+
+Bug Fixes
+
+- Fixes incorrect units in the documentation of some curvature variables.
+
+
 v0.17.0
 -------
 
 New Features
 
-- [Significantly improves convergence of inverse stream maps](https://github.com/PlasmaControl/DESC/pull/1919).
-    - 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 particle tracing capabilities in ``desc.particles`` module.
     - Particle tracing is done via ``desc.particles.trace_particles`` function.
     - Particles can be initialized in couple different ways:
@@ -48,16 +64,13 @@ Bug Fixes
 
 - No longer uses the full Hessian to compute the scale when ``x_scale="auto"`` and using a scipy optimizer that approximates the hessian (e.g. if using ``"scipy-bfgs"``, no longer attempts the Hessian computation to get the x_scale).
 - ``SplineMagneticField.from_field()`` correctly uses the ``NFP`` input when given. Also adds this as a similar input option to ``MagneticField.save_mgrid()``.
-- Significantly improves convergence of inverse stream maps in bounce integrals.
 - Fixes some bugs that hampered robustness of ``desc.geometry.FourierRZToroidalSurface.constant_offset_surface``, particularly when the given grid had stellarator symmetry or when NFP=1.
 - Fixes possible bug in computing normalizations when both kinetic and pressure profiles are assigned. Also adds warnings whenever an pressure is added to a kinetic-constrained equilibrium and vice-versa to alert user to ambiguous equilibrium setups.
 - Adds error in ``MercierStability`` to guard against situation where if a grid with a point at ``rho=0`` were used, NaN would be computed, as``MercierStability`` is undefined on-axis.
 
 Performance Improvements
 
 - `ProximalProjection.grad` uses a single VJP on the objective instead of multiple JVP followed by a manual VJP. This should be more efficient for expensive objectives.
-- 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.
 
 Deprecations
 

desc/backend.py

@@ -36,12 +36,12 @@
                     + "installed JAX with GPU support?"
                 )
                 set_device("cpu")
-            x = jnp.linspace(0, 5)
+            x = jnp.linspace(0, 5, 2)
             y = jnp.exp(x)
         use_jax = True
     except ModuleNotFoundError:
         jnp = np
-        x = jnp.linspace(0, 5)
+        x = jnp.linspace(0, 5, 2)
         y = jnp.exp(x)
         use_jax = False
         set_device(kind="cpu")

desc/compute/__init__.py

@@ -23,8 +23,7 @@
 
 """
 
-# just need to import all the submodules here to register everything in the
-# data_index
+import numpy as np
 
 from ..utils import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec
 from . import (
@@ -55,42 +54,160 @@
     profile_names,
 )
 
+# just need to import all the submodules here to register everything in the
+# data_index
+
 
 # Rather than having to recursively compute the full dependencies every time we
 # compute something, it's easier to just do it once for all quantities when we first
 # import the compute module.
-def _build_data_index():
+def _build_data_index():  # noqa: C901
+    """For each quantity in data_index, build the full set of dependencies.
+
+    This function first performs a sort of the quantities such that the ones with
+    no dependencies come first, then the ones that depend only on those, etc. Then
+    it iterates through the quantities in that order, building the full dependency
+    set for each one by taking the union of its direct dependencies and the full
+    dependencies of those dependencies, which have already been computed by the time
+    we get to this quantity.
+
+    The first sorting is important to avoid deep recursion when building the full
+    dependency sets. Since the first elements of the order have no dependencies,
+    we can build their full dependency sets by simple union operation.
+
+    Note: This function is originally written by Claude Code and reviewed by Yigit
+    Gunsur Elmacioglu.
+    """
+
+    def _collect_deps(p, all_deps):
+        """Collect transforms, params, profiles from a list of dependency keys.
+
+        For each key in the data_index, we call this function with full set of
+        dependencies of that key, and it collects the transforms, params, and
+        profiles needed by all those dependencies in a single pass.
+        """
+        transforms = {}
+        params = []
+        profiles = []
+        for k in all_deps:
+            k_deps = data_index[p][k]["dependencies"]
+            for tkey, tval in k_deps["transforms"].items():
+                if tkey not in transforms:
+                    transforms[tkey] = []
+                transforms[tkey] += tval
+            params += k_deps["params"]
+            profiles += k_deps["profiles"]
+        transforms = {k: np.unique(v, axis=0).tolist() for k, v in transforms.items()}
+        profiles = sorted(set(profiles))
+        return transforms, params, profiles
 
     for p in data_index:
-        for key in data_index[p]:
+        # --- Step 1: Topological sort via iterative Depth-First Search ---
+        # We need to process quantities with no dependencies before the
+        # quantities that depend on them. This way, when we process key K,
+        # all of K's dependencies already have their full_dependencies and
+        # full_with_axis_dependencies cached, and we can build K's full
+        # dependency set with a simple set union instead of deep recursion.
+        order = []
+        visited = set()
+        for start in data_index[p]:
+            if start in visited:
+                continue
+            stack = [(start, False)]
+            while stack:
+                node, processed = stack.pop()
+                if processed:
+                    if node not in visited:
+                        visited.add(node)
+                        order.append(node)
+                    continue
+                if node in visited:
+                    continue
+                # Mark for post-processing after all deps are visited.
+                stack.append((node, True))
+                node_deps = data_index[p][node]["dependencies"]
+                for dep in node_deps["data"]:
+                    if dep not in visited:
+                        stack.append((dep, False))
+                for dep in node_deps["axis_limit_data"]:
+                    if dep not in visited:
+                        stack.append((dep, False))
+
+        # --- Step 2: Build full dependency sets incrementally ---
+        # Because we iterate in topological order, every dependency of the
+        # current key already has its full_dependencies cached. So the full
+        # transitive data deps of key K is just:
+        #   union of (each direct dep D) + (D's already-cached full data deps)
+        # No recursion needed — O(number of direct deps) per key.
+
+        # The deps are stored in topological order (not alphabetical) so that
+        # iterating over them computes quantities in valid dependency order.
+        topo_index = {key: i for i, key in enumerate(order)}
+        for key in order:
+            d = data_index[p][key]
+            deps_info = d["dependencies"]
+            direct_data = deps_info["data"]
+
+            # Full data deps without axis limit contributions.
+            # Apply union operation
+            full_data_set = set()
+            for dep in direct_data:
+                full_data_set.add(dep)
+                full_data_set.update(data_index[p][dep]["full_dependencies"]["data"])
+            deps_no_axis = sorted(full_data_set, key=topo_index.__getitem__)
+
+            transforms, params, profiles = _collect_deps(p, [key] + deps_no_axis)
             full = {
-                "data": get_data_deps(key, p, has_axis=False, basis="rpz"),
-                "transforms": get_derivs(key, p, has_axis=False, basis="rpz"),
-                "params": get_params(key, p, has_axis=False, basis="rpz"),
-                "profiles": get_profiles(key, p, has_axis=False, basis="rpz"),
+                "data": deps_no_axis,
+                "transforms": transforms,
+                "params": params,
+                "profiles": profiles,
             }
-            data_index[p][key]["full_dependencies"] = full
-            full_with_axis_data = get_data_deps(key, p, has_axis=True)
-            if len(full["data"]) >= len(full_with_axis_data):
-                # Then this quantity and all its dependencies do not need anything
+            # Cache now so later keys can use it.
+            d["full_dependencies"] = full
+
+            # Full data deps including axis limit data contributions.
+            # axis_limit_data lists extra quantities needed to evaluate limits
+            # at the magnetic axis; these are only relevant when has_axis=True.
+            axis_limit_data = deps_info["axis_limit_data"]
+            full_data_axis_set = set()
+            for dep in direct_data:
+                full_data_axis_set.add(dep)
+                full_data_axis_set.update(
+                    data_index[p][dep]["full_with_axis_dependencies"]["data"]
+                )
+            for dep in axis_limit_data:
+                full_data_axis_set.add(dep)
+                full_data_axis_set.update(
+                    data_index[p][dep]["full_with_axis_dependencies"]["data"]
+                )
+            deps_with_axis = sorted(full_data_axis_set, key=topo_index.__getitem__)
+
+            if len(deps_no_axis) >= len(deps_with_axis):
+                # This quantity and all its dependencies do not need anything
                 # extra to evaluate its limit at the magnetic axis.
-                # The dependencies in the `full` dictionary and the `full_with_axis`
-                # dictionary will be identical, so we assign the same reference to
-                # avoid storing a copy.
+                # Assign the same reference to avoid storing a copy.
                 full_with_axis = full
             else:
+                transforms_a, params_a, profiles_a = _collect_deps(
+                    p, [key] + deps_with_axis
+                )
                 full_with_axis = {
-                    "data": full_with_axis_data,
-                    "transforms": get_derivs(key, p, has_axis=True, basis="rpz"),
-                    "params": get_params(key, p, has_axis=True, basis="rpz"),
-                    "profiles": get_profiles(key, p, has_axis=True, basis="rpz"),
+                    "data": deps_with_axis,
+                    "transforms": transforms_a,
+                    "params": params_a,
+                    "profiles": profiles_a,
                 }
+                # transforms, params, and profiles can be the same for both full and
+                # full_with_axis, so check if they are and if so, dereference the copy
+                # to save memory.
                 for _key, val in full_with_axis.items():
                     if full[_key] == val:
-                        # Nothing extra was needed to evaluate this quantity's limit.
-                        # One is a copy of the other; dereference to save memory.
+                        # Nothing extra was needed for this field.
+                        # Dereference the copy to save memory.
                         full_with_axis[_key] = full[_key]
-            data_index[p][key]["full_with_axis_dependencies"] = full_with_axis
+
+            d["full_with_axis_dependencies"] = full_with_axis
 
 
 _build_data_index()

desc/integrals/bounce_integral.py

@@ -147,12 +147,6 @@ def plot(self, l, m, pitch_inv=None, **kwargs):
         """Plot B and bounce points on the specified field line."""
 
 
-default_quad = get_quadrature(
-    leggauss(32),
-    (automorphism_sin, grad_automorphism_sin),
-)
-
-
 class Bounce2D(Bounce):
     """Computes bounce integrals using pseudo-spectral methods.
 
@@ -301,7 +295,14 @@ def __init__(
         is_reshaped = is_reshaped or is_fourier
         vander = setdefault(vander, {})
 
-        self._quad = get_quadrature(setdefault(quad, default_quad), automorphism)
+        if quad is None:
+            quad = get_quadrature(
+                leggauss(32), (automorphism_sin, grad_automorphism_sin)
+            )
+        else:
+            quad = get_quadrature(quad, automorphism)
+        self._quad = jax.lax.stop_gradient(quad)
+
         self._NFP = grid.NFP
         self._num_t = grid.num_theta
         self._modes_z, self._modes_t = rfft2_modes(
@@ -1468,9 +1469,13 @@ def __init__(
     ):
         """Returns an object to compute bounce integrals."""
         assert grid.is_meshgrid
-        quad = setdefault(quad, default_quad)
-
-        self._quad = get_quadrature(quad, automorphism)
+        if quad is None:
+            quad = get_quadrature(
+                leggauss(32), (automorphism_sin, grad_automorphism_sin)
+            )
+        else:
+            quad = get_quadrature(quad, automorphism)
+        self._quad = jax.lax.stop_gradient(quad)
         self._data = {
             "|b^zeta|": jnp.abs(data["B^zeta"]) * Lref / data["|B|"],
             "|B|": data["|B|"] / Bref,