Simplified Bout computation in free boundary

CHANGELOG.md

@@ -24,6 +24,7 @@ New Features
 - Adds batching feature to singular integrals.
 - ``desc.objectives.CoilSetMinDistance`` and ``desc.objectives.PlasmaCoilSetMinDistance`` now include the option to use a softmin which can give smoother gradients. They also both now have a ``dist_chunk_size`` option to break up the distance calculation into smaller pieces to save memory
 - Adds a new function ``desc.coils.initialize_helical_coils`` for creating an initial guess for stage 2 helical coil optimization.
+- Adds ``chunk_size`` option to compute magnetic field methods to increase performance.
 - Adds ``desc.vmec_utils.make_boozmn_output `` for writing boozmn.nc style output files
 for compatibility with other codes which expect such files from the Booz_Xform code.
 - Renames compute quantity ``sqrt(g)_B`` to ``sqrt(g)_Boozer_DESC`` to more accurately reflect what the quantiy is (the jacobian from (rho,theta_B,zeta_B) to (rho,theta,zeta)), and adds a new function to compute ``sqrt(g)_Boozer`` which is the jacobian from (rho,theta_B,zeta_B) to (R,phi,Z).

README.rst

@@ -97,7 +97,7 @@ Contribute
 - `Contributing guidelines <https://github.com/PlasmaControl/DESC/blob/master/CONTRIBUTING.rst>`_
 - `Issue Tracker <https://github.com/PlasmaControl/DESC/issues>`_
 - `Source Code <https://github.com/PlasmaControl/DESC/>`_
-- `Documentation <https://desc-docs.readthedocs.io/>`_
+- `Documentation <https://desc-docs.readthedocs.io/en/latest/>`_
 
 .. |License| image:: https://img.shields.io/github/license/PlasmaControl/desc?color=blue&logo=open-source-initiative&logoColor=white
     :target: https://github.com/PlasmaControl/DESC/blob/master/LICENSE

desc/batching.py

@@ -4,8 +4,7 @@
 ----------
 The ``_batch_and_remainder``, ``_evaluate_in_chunks``, ``jacrev_chunked``,
 and ``jacfwd_chunked`` functions are adapted from JAX.
-https://github.com/jax-ml/jax/blob/main/jax/_src/lax/control_flow/loops.py#L2139.
-https://github.com/jax-ml/jax/blob/main/jax/_src/lax/control_flow/loops.py#L2209.
+https://github.com/jax-ml/jax/blob/main/jax/_src/lax/control_flow/loops.py.
 https://github.com/jax-ml/jax/blob/main/jax/_src/api.py.
 
 The original copyright notice is as follows
@@ -42,7 +41,7 @@
     tree_transpose,
 )
 
-from desc.backend import jax, jnp, scan
+from desc.backend import jax, jnp, scan, vmap
 from desc.utils import errorif
 
 if jax.__version_info__ >= (0, 4, 16):
@@ -191,7 +190,7 @@ def vmap_chunked(
     reduction=None,
     chunk_reduction=_identity,
 ):
-    """Behaves like jax.vmap but uses scan to chunk the computations in smaller chunks.
+    """Behaves like ``vmap`` but uses scan to chunk the computations in smaller chunks.
 
     Parameters
     ----------
@@ -217,10 +216,7 @@ def vmap_chunked(
 
     """
     in_axes, argnums = _parse_in_axes(in_axes)
-    if isinstance(argnums, int):
-        argnums = (argnums,)
-
-    f = jax.vmap(f, in_axes=in_axes)
+    f = vmap(f, in_axes=in_axes)
     if chunk_size is None:
         return lambda *args, **kwargs: chunk_reduction(f(*args, **kwargs))
     return partial(
@@ -244,12 +240,11 @@ def batch_map(
 ):
     """Compute ``chunk_reduction(fun(fun_input))`` in batches.
 
-    This utility is like ``desc.backend.imap(fun,fun_input,batch_size)`` 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.
+    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
     ----------
@@ -265,7 +260,8 @@ def batch_map(
         Should take two arguments and return one output, e.g. ``jnp.add``.
     chunk_reduction : callable
         Chunk-wise reduction operation.
-        Should apply ``reduction`` along the mapped axis, e.g. ``jnp.add.reduce``.
+        Should typically apply ``reduction`` along the mapped axis,
+        e.g. ``jnp.add.reduce``.
 
     Returns
     -------