diff --git a/.gitignore b/.gitignore index df621252f0..551d089d5f 100644 --- a/.gitignore +++ b/.gitignore @@ -30,3 +30,5 @@ uv.lock # Environments .env .venv + +.vscode/ diff --git a/CHANGELOG.md b/CHANGELOG.md index 1030fd7c1b..3c5621a65b 100644 --- a/CHANGELOG.md +++ b/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,7 +64,6 @@ 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. @@ -56,8 +71,6 @@ Bug Fixes 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 diff --git a/desc/backend.py b/desc/backend.py index bfd6908108..a5dbd9a17b 100644 --- a/desc/backend.py +++ b/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") diff --git a/desc/compute/__init__.py b/desc/compute/__init__.py index 1627dfe63f..343ca29002 100644 --- a/desc/compute/__init__.py +++ b/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() diff --git a/desc/integrals/bounce_integral.py b/desc/integrals/bounce_integral.py index 02ca47d65e..d14f209bc9 100644 --- a/desc/integrals/bounce_integral.py +++ b/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,