Tweak documentation as requested in pull request review

desc/grid.py

@@ -638,7 +638,8 @@ def meshgrid_reshape(self, x, order):
             vec = True
             shape += (-1,)
         x = x.reshape(shape, order="F")
-        x = jnp.swapaxes(x, 1, 0)  # now shape rtz/raz etc
+        # swap to change shape from trz/arz to rtz/raz etc.
+        x = jnp.swapaxes(x, 1, 0)
         newax = tuple(self.coordinates.index(c) for c in order)
         if vec:
             newax += (3,)

desc/integrals/basis.py

@@ -27,6 +27,11 @@ def _in_epigraph_and(is_intersect, df_dy_sign, /):
         Boolean array indicating whether element is an intersect
         and satisfies the stated condition.
 
+    Examples
+    --------
+    See ``desc/integrals/bounce_utils.py::bounce_points``.
+    This is used there to ensure the domains of integration are magnetic wells.
+
     """
     # The pairs ``y1`` and ``y2`` are boundaries of an integral only if ``y1 <= y2``.
     # For the integrals to be over wells, it is required that the first intersect

desc/integrals/bounce_integral.py

@@ -26,12 +26,12 @@
 class Bounce1D(IOAble):
     """Computes bounce integrals using one-dimensional local spline methods.
 
-    The bounce integral is defined as ∫ f(ℓ) dℓ, where
+    The bounce integral is defined as ∫ f(λ, ℓ) dℓ, where
         dℓ parameterizes the distance along the field line in meters,
-        f(ℓ) is the quantity to integrate along the field line,
-        and the boundaries of the integral are bounce points ζ₁, ζ₂ s.t. λ|B|(ζᵢ) = 1,
-        where λ is a constant proportional to the magnetic moment over energy
-        and |B| is the norm of the magnetic field.
+        f(λ, ℓ) is the quantity to integrate along the field line,
+        and the boundaries of the integral are bounce points ℓ₁, ℓ₂ s.t. λ|B|(ℓᵢ) = 1,
+        where λ is a constant defining the integral proportional to the magnetic moment
+        over energy and |B| is the norm of the magnetic field.
 
     For a particle with fixed λ, bounce points are defined to be the location on the
     field line such that the particle's velocity parallel to the magnetic field is zero.
@@ -299,18 +299,18 @@ def integrate(
         check=False,
         plot=False,
     ):
-        """Bounce integrate ∫ f(ℓ) dℓ.
+        """Bounce integrate ∫ f(λ, ℓ) dℓ.
 
-        Computes the bounce integral ∫ f(ℓ) dℓ for every field line and pitch.
+        Computes the bounce integral ∫ f(λ, ℓ) dℓ for every field line and pitch.
 
         Parameters
         ----------
         integrand : callable
             The composition operator on the set of functions in ``f`` that maps the
-            functions in ``f`` to the integrand f(ℓ) in ∫ f(ℓ) dℓ. It should accept the
-            arrays in ``f`` as arguments as well as the additional keyword arguments:
-            ``B`` and ``pitch``. A quadrature will be performed to approximate the
-            bounce integral of ``integrand(*f,B=B,pitch=pitch)``.
+            functions in ``f`` to the integrand f(λ, ℓ) in ∫ f(λ, ℓ) dℓ. It should
+            accept the arrays in ``f`` as arguments as well as the additional keyword
+            arguments: ``B`` and ``pitch``. A quadrature will be performed to
+            approximate the bounce integral of ``integrand(*f,B=B,pitch=pitch)``.
         pitch_inv : jnp.ndarray
             Shape (M, L, P).
             1/λ values to compute the bounce integrals. 1/λ(α,ρ) is specified by
@@ -325,7 +325,7 @@ def integrate(
         weight : jnp.ndarray
             Shape (M, L, N).
             If supplied, the bounce integral labeled by well j is weighted such that
-            the returned value is w(j) ∫ f(ℓ) dℓ, where w(j) is ``weight``
+            the returned value is w(j) ∫ f(λ, ℓ) dℓ, where w(j) is ``weight``
             interpolated to the deepest point in that magnetic well. Use the method
             ``self.reshape_data`` to reshape the data into the expected shape.
         num_well : int or None
@@ -410,9 +410,9 @@ def plot(self, m, l, pitch_inv=None, /, **kwargs):
         """
         B, dB_dz = self.B, self._dB_dz
         if B.ndim == 4:
-            B = B[m, l]
-            dB_dz = dB_dz[m, l]
-        elif B.ndim == 3:
+            B = B[m]
+            dB_dz = dB_dz[m]
+        if B.ndim == 3:
             B = B[l]
             dB_dz = dB_dz[l]
         if pitch_inv is not None:

desc/integrals/bounce_utils.py

@@ -13,8 +13,8 @@
     polyval_vec,
 )
 from desc.integrals.quad_utils import (
-    _composite_linspace,
     bijection_from_disc,
+    composite_linspace,
     grad_bijection_from_disc,
 )
 from desc.utils import (
@@ -54,7 +54,7 @@ def get_pitch_inv(min_B, max_B, num, relative_shift=1e-6):
     min_B = (1 + relative_shift) * min_B
     max_B = (1 - relative_shift) * max_B
     # Samples should be uniformly spaced in |B| and not λ (GitHub issue #1228).
-    pitch_inv = jnp.moveaxis(_composite_linspace(jnp.stack([min_B, max_B]), num), 0, -1)
+    pitch_inv = jnp.moveaxis(composite_linspace(jnp.stack([min_B, max_B]), num), 0, -1)
     assert pitch_inv.shape == (*min_B.shape, num + 2)
     return pitch_inv
 
@@ -295,7 +295,7 @@ def _bounce_quadrature(
     check=False,
     plot=False,
 ):
-    """Bounce integrate ∫ f(ℓ) dℓ.
+    """Bounce integrate ∫ f(λ, ℓ) dℓ.
 
     Parameters
     ----------
@@ -312,10 +312,10 @@ def _bounce_quadrature(
         epigraph of |B|.
     integrand : callable
         The composition operator on the set of functions in ``f`` that maps the
-        functions in ``f`` to the integrand f(ℓ) in ∫ f(ℓ) dℓ. It should accept the
-        arrays in ``f`` as arguments as well as the additional keyword arguments:
-        ``B`` and ``pitch``. A quadrature will be performed to approximate the
-        bounce integral of ``integrand(*f,B=B,pitch=pitch)``.
+        functions in ``f`` to the integrand f(λ, ℓ) in ∫ f(λ, ℓ) dℓ. It should
+        accept the arrays in ``f`` as arguments as well as the additional keyword
+        arguments: ``B`` and ``pitch``. A quadrature will be performed to
+        approximate the bounce integral of ``integrand(*f,B=B,pitch=pitch)``.
     pitch_inv : jnp.ndarray
         Shape (..., P).
         1/λ values to compute the bounce integrals.

desc/integrals/quad_utils.py

@@ -19,7 +19,7 @@ def bijection_from_disc(x, a, b):
 
 
 def grad_bijection_from_disc(a, b):
-    """Gradient of affine bijection from disc."""
+    """Gradient wrt ``x`` of ``bijection_from_disc``."""
     dy_dx = 0.5 * (b - a)
     return dy_dx
 
@@ -220,7 +220,7 @@ def get_quadrature(quad, automorphism):
     return x, w
 
 
-def _composite_linspace(x, num):
+def composite_linspace(x, num):
     """Returns linearly spaced values between every pair of values in ``x``.
 
     Parameters

desc/utils.py

@@ -742,16 +742,4 @@ def atleast_nd(ndmin, ary):
     return jnp.array(ary, ndmin=ndmin) if jnp.ndim(ary) < ndmin else ary
 
 
-def atleast_3d_mid(ary):
-    """Like np.atleast_3d but if adds dim at axis 1 for 2d arrays."""
-    ary = jnp.atleast_2d(ary)
-    return ary[:, jnp.newaxis] if ary.ndim == 2 else ary
-
-
-def atleast_2d_end(ary):
-    """Like np.atleast_2d but if adds dim at axis 1 for 1d arrays."""
-    ary = jnp.atleast_1d(ary)
-    return ary[:, jnp.newaxis] if ary.ndim == 1 else ary
-
-
 PRINT_WIDTH = 60  # current longest name is BootstrapRedlConsistency with pre-text

tests/test_integrals.py

@@ -921,7 +921,7 @@ class TestBounce1DQuadrature:
         ],
     )
     def test_bounce_quadrature(self, is_strong, quad, automorphism):
-        """Test bounce integral matches singular elliptic integrals."""
+        """Test quadrature matches singular (strong and weak) elliptic integrals."""
         p = 1e-4
         m = 1 - p
         # Some prime number that doesn't appear anywhere in calculation.
@@ -971,9 +971,12 @@ def _fixed_elliptic(integrand, k, deg):
         x, w = get_quadrature(leggauss(deg), (automorphism_sin, grad_automorphism_sin))
         Z = bijection_from_disc(x, a[..., np.newaxis], b[..., np.newaxis])
         k = k[..., np.newaxis]
-        quad = np.dot(integrand(Z, k), w) * grad_bijection_from_disc(a, b)
+        quad = integrand(Z, k).dot(w) * grad_bijection_from_disc(a, b)
         return quad
 
+    # TODO: add the analytical test that converts incomplete elliptic integrals to
+    #  complete ones using the Reciprocal Modulus transformation
+    #  https://dlmf.nist.gov/19.7#E4.
     @staticmethod
     def elliptic_incomplete(k2):
         """Calculate elliptic integrals for bounce averaged binormal drift.
@@ -1179,6 +1182,7 @@ def dg_dz(z):
         assert result.shape == z1.shape
         np.testing.assert_allclose(h_min, result, rtol=1e-3)
 
+    # TODO: stellarator geometry test with ripples
     @staticmethod
     def drift_analytic(data):
         """Compute analytic approximation for bounce-averaged binormal drift.

tests/test_quad_utils.py

@@ -6,11 +6,11 @@
 
 from desc.backend import jnp
 from desc.integrals.quad_utils import (
-    _composite_linspace,
     automorphism_arcsin,
     automorphism_sin,
     bijection_from_disc,
     bijection_to_disc,
+    composite_linspace,
     grad_automorphism_arcsin,
     grad_automorphism_sin,
     grad_bijection_from_disc,
@@ -26,7 +26,7 @@ def test_composite_linspace():
     B_min_tz = np.array([0.1, 0.2])
     B_max_tz = np.array([1, 3])
     breaks = np.linspace(B_min_tz, B_max_tz, num=5)
-    b = _composite_linspace(breaks, num=3)
+    b = composite_linspace(breaks, num=3)
     for i in range(breaks.shape[0]):
         for j in range(breaks.shape[1]):
             assert only1(np.isclose(breaks[i, j], b[:, j]).tolist())