Preparing merge into bounce branch

desc/basis.py

@@ -1131,10 +1131,6 @@ def evaluate(
             m = m[midx]
             n = n[nidx]
 
-        # TODO: in map_clebsch_root findign
-        #       lambda should be fixed to rho and zeta
-        #       so lambda is slimmed to 1d fourier series for fixed rho zeta.
-        #       cache radial and toroidal for rootfinding
         radial = zernike_radial(r[:, np.newaxis], lm[:, 0], lm[:, 1], dr=derivatives[0])
         poloidal = fourier(t[:, np.newaxis], m, dt=derivatives[1])
         toroidal = fourier(z[:, np.newaxis], n, NFP=self.NFP, dt=derivatives[2])

desc/integrals/basis.py

@@ -7,10 +7,10 @@
 
 from desc.backend import dct, flatnonzero, idct, irfft, jnp, put, rfft
 from desc.integrals.interp_utils import (
+    _filter_distinct,
     cheb_from_dct,
     cheb_pts,
     chebroots_vec,
-    filter_distinct,
     fourier_pts,
     harmonic,
     idct_non_uniform,
@@ -50,8 +50,8 @@ def _subtract(c, k):
 def _in_epigraph_and(is_intersect, df_dy_sign):
     """Set and epigraph of function f with the given set of points.
 
-    Return only intersects where the straight line path between adjacent
-    intersects resides in the epigraph of a continuous map ``f``.
+    Used to return only intersects where the straight line path between
+    adjacent intersects resides in the epigraph of a continuous map ``f``.
 
     Warnings
     --------
@@ -430,7 +430,7 @@ def intersect2d(self, k=0.0, eps=_eps):
 
         # Intersects must satisfy y ∈ [-1, 1].
         # Pick sentinel such that only distinct roots are considered intersects.
-        y = filter_distinct(y, sentinel=-2.0, eps=eps)
+        y = _filter_distinct(y, sentinel=-2.0, eps=eps)
         is_intersect = (jnp.abs(y.imag) <= eps) & (jnp.abs(y.real) <= 1.0)
         # Ensure y is in domain of arcos; choose 1 because kernel probably cheaper.
         y = jnp.where(is_intersect, y.real, 1.0)
@@ -471,17 +471,16 @@ def intersect1d(self, k=0.0, num_intersect=None, pad_value=0.0):
         -------
         z1, z2 : (jnp.ndarray, jnp.ndarray)
             Shape broadcasts with (..., *self.cheb.shape[:-2], num_intersect).
-            ``z1``, ``z2`` holds intersects satisfying ∂f/∂y <= 0, ∂f/∂y >= 0,
-            respectively. The points are grouped and ordered such that the
-            straight line path between the intersects in ``z1`` and ``z2``
-            resides in the epigraph of f.
+            ``z1`` and ``z2`` are intersects satisfying ∂f/∂y <= 0 and ∂f/∂y >= 0,
+            respectively. The points are grouped and ordered such that the straight
+            line path between ``z1`` and ``z2`` resides in the epigraph of f.
 
         """
         errorif(
             self.N < 2,
             NotImplementedError,
-            "This method requires the Chebyshev spectral resolution of at "
-            f"least 2, but got N={self.N}.",
+            "This method requires a Chebyshev spectral resolution of N > 1, "
+            f"but got N = {self.N}.",
         )
 
         # Add axis to use same k over all Chebyshev series of the piecewise object.
@@ -498,9 +497,9 @@ def intersect1d(self, k=0.0, num_intersect=None, pad_value=0.0):
         # polynomials is a left intersection i.e. ``is_z1`` because the subset of
         # pitch values that generate this edge case has zero measure. By ignoring
         # this, for those subset of pitch values the integrations will be done in
-        # the hypograph of |B| rather than the epigraph, which will be integrated
-        # to zero. If we decide later to not ignore this, the technique to solve
-        # this is to disqualify intersects within ``_eps`` from ``domain[-1]``.
+        # the hypograph of |B|, which will yield zero. If in far future decide to
+        # not ignore this, note the solution is to disqualify intersects within
+        # ``_eps`` from ``domain[-1]``.
         is_z1 = (df_dy_sign <= 0) & is_intersect
         is_z2 = (df_dy_sign >= 0) & _in_epigraph_and(is_intersect, df_dy_sign)
 
@@ -532,10 +531,9 @@ def check_intersect1d(self, z1, z2, k, plot=True, **kwargs):
         ----------
         z1, z2 : jnp.ndarray
             Shape must broadcast with (*self.cheb.shape[:-2], W).
-            ``z1``, ``z2`` holds intersects satisfying ∂f/∂y <= 0, ∂f/∂y >= 0,
-            respectively. The points are grouped and ordered such that the
-            straight line path between the intersects in ``z1`` and ``z2``
-            resides in the epigraph of f.
+            ``z1`` and ``z2`` are intersects satisfying ∂f/∂y <= 0 and ∂f/∂y >= 0,
+            respectively. The points are grouped and ordered such that the straight
+            line path between ``z1`` and ``z2`` resides in the epigraph of f.
         k : jnp.ndarray
             Shape must broadcast with *self.cheb.shape[:-2].
             k such that fₓ(yᵢ) = k.

desc/integrals/bounce_integral.py

@@ -10,7 +10,7 @@
     bounce_points,
     bounce_quadrature,
     get_alpha,
-    interp_to_argmin_B_soft,
+    interp_to_argmin_g,
     plot_ppoly,
 )
 from desc.integrals.interp_utils import interp_rfft2, irfft2_non_uniform, polyder_vec
@@ -444,8 +444,8 @@ def bounce_points(self, pitch, num_well=None):
         z1, z2 : (jnp.ndarray, jnp.ndarray)
             Shape (P, L, num_well).
             ζ coordinates of bounce points. The points are grouped and ordered such
-            that the straight line path between the intersects in ``z1`` and ``z2``
-            resides in the epigraph of |B|.
+            that the straight line path between ``z1`` and ``z2`` resides in the
+            epigraph of |B|.
 
         """
         return self._B.intersect1d(1 / jnp.atleast_2d(pitch), num_well)
@@ -458,8 +458,8 @@ def check_bounce_points(self, z1, z2, pitch, plot=True, **kwargs):
         z1, z2 : (jnp.ndarray, jnp.ndarray)
             Shape (P, L, num_well).
             ζ coordinates of bounce points. The points are grouped and ordered such
-            that the straight line path between the intersects in ``z1`` and ``z2``
-            resides in the epigraph of |B|.
+            that the straight line path between ``z1`` and ``z2`` resides in the
+            epigraph of |B|.
         pitch : jnp.ndarray
             Shape (P, L).
             λ values to evaluate the bounce integral at each field line. λ(ρ) is
@@ -777,8 +777,8 @@ def bounce_points(self, pitch, num_well=None):
         z1, z2 : (jnp.ndarray, jnp.ndarray)
             Shape (P, L * M, num_well).
             ζ coordinates of bounce points. The points are grouped and ordered such
-            that the straight line path between the intersects in ``z1`` and ``z2``
-            resides in the epigraph of |B|.
+            that the straight line path between ``z1`` and ``z2`` resides in the
+            epigraph of |B|.
 
             If there were less than ``num_wells`` wells detected along a field line,
             then the last axis, which enumerates bounce points for  a particular field
@@ -801,8 +801,8 @@ def check_bounce_points(self, z1, z2, pitch, plot=True, **kwargs):
         z1, z2 : (jnp.ndarray, jnp.ndarray)
             Shape (P, L * M, num_well).
             ζ coordinates of bounce points. The points are grouped and ordered such
-            that the straight line path between the intersects in ``z1`` and ``z2``
-            resides in the epigraph of |B|.
+            that the straight line path between ``z1`` and ``z2`` resides in the
+            epigraph of |B|.
         pitch : jnp.ndarray
             Shape must broadcast with (P, L * M).
             λ values to evaluate the bounce integral at each field line. λ(ρ,α) is
@@ -877,7 +877,7 @@ def integrate(
             wells detected along a field line than the size of the last axis of the
             returned arrays, then that axis is padded with zero.
         method : str
-            Method of interpolation for functions contained in ``f``.
+            Method of interpolation.
             See https://interpax.readthedocs.io/en/latest/_api/interpax.interp1d.html.
             Default is cubic C1 local spline.
         batch : bool
@@ -910,13 +910,13 @@ def integrate(
             check=check,
         )
         if weight is not None:
-            result *= interp_to_argmin_B_soft(
-                g=weight,
+            result *= interp_to_argmin_g(
+                h=weight,
                 z1=z1,
                 z2=z2,
                 knots=self._zeta,
-                B=self.B,
-                dB_dz=self._dB_dz,
+                g=self.B,
+                dg_dz=self._dB_dz,
                 method=method,
             )
         assert result.shape[-1] == setdefault(num_well, (self._zeta.size - 1) * 3)