From 4097aa8242a68996f22d437a7d331d0ba8573cf6 Mon Sep 17 00:00:00 2001
From: unalmis <kayaunalmis@proton.me>
Date: Thu, 19 Dec 2024 16:17:01 -0500
Subject: [PATCH] Fix math comment

---
 desc/compute/_deprecated.py       |  2 +-
 desc/compute/_neoclassical.py     |  6 +++---
 desc/integrals/bounce_integral.py | 12 +++---------
 3 files changed, 7 insertions(+), 13 deletions(-)

diff --git a/desc/compute/_deprecated.py b/desc/compute/_deprecated.py
index bcce0db5b..35142f268 100644
--- a/desc/compute/_deprecated.py
+++ b/desc/compute/_deprecated.py
@@ -121,7 +121,7 @@ def _epsilon_32_1D(params, transforms, profiles, data, **kwargs):
     )
 
     def eps_32(data):
-        """(∂ψ/∂ρ)⁻² B₀⁻² ∫ dλ λ⁻² ∑ⱼ Hⱼ²/Iⱼ."""
+        """(∂ψ/∂ρ)⁻² B₀⁻³ ∫ dλ λ⁻² ∑ⱼ Hⱼ²/Iⱼ."""
         # B₀ has units of λ⁻¹.
         # Nemov's ∑ⱼ Hⱼ²/Iⱼ = (∂ψ/∂ρ)² (λB₀)³ ``(H**2 / I).sum(axis=-1)``.
         # (λB₀)³ d(λB₀)⁻¹ = B₀² λ³ d(λ⁻¹) = -B₀² λ dλ.
diff --git a/desc/compute/_neoclassical.py b/desc/compute/_neoclassical.py
index f2ef2b495..bbc3fe629 100644
--- a/desc/compute/_neoclassical.py
+++ b/desc/compute/_neoclassical.py
@@ -6,7 +6,7 @@
 
 from desc.backend import imap, jit, jnp
 
-from ..batching import _chunk_vmapped_function
+from ..batching import _eval_fun_in_chunks
 from ..integrals.bounce_integral import Bounce2D
 from ..integrals.quad_utils import chebgauss2
 from ..utils import safediv
@@ -121,7 +121,7 @@ def _foreach_pitch(fun, pitch_inv, batch_size):
         fun(pitch_inv)
         if (batch_size is None or batch_size >= (pitch_inv.size - 1))
         # else imap(fun, pitch_inv, batch_size=batch_size).squeeze(axis=-1) # noqa: E800
-        else _chunk_vmapped_function(fun, chunk_size=batch_size)(pitch_inv)
+        else _eval_fun_in_chunks(fun, batch_size, (0,), pitch_inv)
     )
 
 
@@ -199,7 +199,7 @@ def _epsilon_32(params, transforms, profiles, data, **kwargs):
     )
 
     def eps_32(data):
-        """(∂ψ/∂ρ)⁻² B₀⁻² ∫ dλ λ⁻² ∑ⱼ Hⱼ²/Iⱼ."""
+        """(∂ψ/∂ρ)⁻² B₀⁻³ ∫ dλ λ⁻² ∑ⱼ Hⱼ²/Iⱼ."""
         # B₀ has units of λ⁻¹.
         # Nemov's ∑ⱼ Hⱼ²/Iⱼ = (∂ψ/∂ρ)² (λB₀)³ ``(H**2 / I).sum(axis=-1)``.
         # (λB₀)³ d(λB₀)⁻¹ = B₀² λ³ d(λ⁻¹) = -B₀² λ dλ.
diff --git a/desc/integrals/bounce_integral.py b/desc/integrals/bounce_integral.py
index 6adebf166..c38c321cf 100644
--- a/desc/integrals/bounce_integral.py
+++ b/desc/integrals/bounce_integral.py
@@ -1045,9 +1045,10 @@ def __init__(
         # Compute local splines.
         # Note it is simple to do FFT across field line axis, and spline
         # Fourier coefficients across ζ to obtain Fourier-CubicSpline of functions.
-        # The point of Bounce2D is to do such a 2D interpolation but also do so
-        # without rebuilding DESC transforms each time an objective is computed.
+        # The point of Bounce2D is to do such a 2D interpolation without
+        # rebuilding DESC transforms each time an objective is computed.
         self._zeta = grid.compress(grid.nodes[:, 2], surface_label="zeta")
+        # Shape is (num alpha, num rho, N - 1, -1).
         self._B = jnp.moveaxis(
             CubicHermiteSpline(
                 x=self._zeta,
@@ -1059,13 +1060,6 @@ def __init__(
             source=(0, 1),
             destination=(-1, -2),
         )
-        # Shape (num alpha, num rho, N - 1, -1).
-        # Polynomial coefficients of the spline of |B| in local power basis.
-        # Last axis enumerates the coefficients of power series. For a polynomial
-        # given by ∑ᵢⁿ cᵢ xⁱ, coefficient cᵢ is stored at ``B[...,n-i]``.
-        # Third axis enumerates the polynomials that compose a particular spline.
-        # Second axis enumerates flux surfaces.
-        # First axis enumerates field lines of a particular flux surface.
         self._dB_dz = polyder_vec(self._B)
 
     @staticmethod