Fix docs

desc/compute/_neoclassical.py

@@ -35,7 +35,7 @@
         JAX this will have worse performance.
         Specifying a number that tightly upper bounds the number of wells will
         increase performance. In general, an upper bound on the number of wells
-        per toroidal transit is ``Aι+B`` where ``A``,``B`` are the poloidal and
+        per toroidal transit is ``Aι+B`` where ``A``, ``B`` are the poloidal and
         toroidal Fourier resolution of B, respectively, in straight-field line
         PEST coordinates, and ι is the rotational transform normalized by 2π.
         A tighter upper bound than ``num_well=(Aι+B)*num_transit`` is preferable.

desc/integrals/_bounce_utils.py

@@ -152,7 +152,7 @@ def bounce_points(
         but due to current limitations in JAX this will have worse performance.
         Specifying a number that tightly upper bounds the number of wells will
         increase performance. In general, an upper bound on the number of wells
-        per toroidal transit is ``Aι+B`` where ``A``,``B`` are the poloidal and
+        per toroidal transit is ``Aι+B`` where ``A``, ``B`` are the poloidal and
         toroidal Fourier resolution of B, respectively, in straight-field line
         PEST coordinates, and ι is the rotational transform normalized by 2π.
         A tighter upper bound than ``num_well=(Aι+B)*num_transit`` is preferable.
@@ -330,9 +330,9 @@ def _bounce_quadrature(
         Shape (N, ).
         Unique ζ coordinates where the arrays in ``data`` and ``f`` were evaluated.
     integrand : callable or list[callable]
-        The composition operator on the set of functions in ``data`` that
-        maps that determines ``f`` in ∫ f(ρ,α,λ,ℓ) dℓ. It should accept a dictionary
-        which stores the interpolated data and the keyword argument ``pitch``.
+        The composition operator on the set of functions in ``data``
+        that determines ``f`` in ∫ f(ρ,α,λ,ℓ) dℓ. It should accept a dictionary
+        which stores the interpolated data and the arguments ``B`` and ``pitch``.
     pitch_inv : jnp.ndarray
         Shape (num alpha, num rho, num pitch).
         1/λ values to compute the bounce integrals. 1/λ(α,ρ) is specified by

desc/integrals/bounce_integral.py

@@ -87,8 +87,8 @@ class Bounce2D(Bounce):
     the particle's guiding center trajectory traveling in the direction of increasing
     field-line-following coordinate ζ.
 
-    Overview
-    --------
+    Notes
+    -----
     Magnetic field line with label α, defined by B = ∇ψ × ∇α, is determined from
       α : ρ, θ, ζ ↦ θ + λ(ρ,θ,ζ) − ι(ρ) [ζ + ω(ρ,θ,ζ)]
     Interpolate Fourier-Chebyshev series to DESC poloidal coordinate.
@@ -113,27 +113,6 @@ class Bounce2D(Bounce):
     --------
     Bounce1D : Uses one-dimensional local spline methods for the same task.
 
-
-    Comparison to Bounce1D
-    ----------------------
-    ``Bounce2D`` solves the dominant cost of optimization objectives relying on
-    ``Bounce1D``: interpolating DESC's 3D transforms to an optimization-step
-    dependent grid that is dense enough for function approximation with local
-    splines. This is sometimes referred to as off-grid interpolation in literature;
-    it is often a bottleneck.
-
-    * The function approximation done here requires DESC transforms on a fixed
-      grid with typical resolution, using FFTs to compute the map between
-      coordinate systems. This enables evaluating functions along field lines
-      without root-finding.
-    * The faster convergence of spectral methods requires a less dense
-      grid to interpolate onto from DESC's 3D transforms.
-    * 2D interpolation enables tracing the field line for many toroidal transits.
-    * The drawback is that evaluating a Fourier series with resolution F at Q
-      non-uniform quadrature points takes 𝒪([F+Q] log[F] log[1/ε]) time
-      whereas cubic splines take 𝒪(C Q) time. However, as NFP increases,
-      F decreases whereas C increases. Also, Q >> F and Q >> C.
-
     """
 
     # Notes
@@ -217,6 +196,24 @@ class Bounce2D(Bounce):
     # Fast transforms are used where possible. Fast multipoint methods are not
     # implemented. For non-uniform interpolation, MMTs are used. It will be
     # worthwhile to use the inverse non-uniform fast transforms.
+    #
+    # ``Bounce2D`` solves the dominant cost of optimization objectives relying on
+    # ``Bounce1D``: interpolating DESC's 3D transforms to an optimization-step
+    # dependent grid that is dense enough for function approximation with local
+    # splines. This is sometimes referred to as off-grid interpolation in literature;
+    # it is often a bottleneck.
+    #
+    # * The function approximation done here requires DESC transforms on a fixed
+    #   grid with typical resolution, using FFTs to compute the map between
+    #   coordinate systems. This enables evaluating functions along field lines
+    #   without root-finding.
+    # * The faster convergence of spectral methods requires a less dense
+    #   grid to interpolate onto from DESC's 3D transforms.
+    # * 2D interpolation enables tracing the field line for many toroidal transits.
+    # * The drawback is that evaluating a Fourier series with resolution F at Q
+    #   non-uniform quadrature points takes 𝒪([F+Q] log[F] log[1/ε]) time
+    #   whereas cubic splines take 𝒪(C Q) time. However, as NFP increases,
+    #   F decreases whereas C increases. Also, Q >> F and Q >> C.
 
     required_names = ["B^zeta", "|B|", "iota"]
 
@@ -458,7 +455,7 @@ def points(self, pitch_inv, *, num_well=None):
             but due to current limitations in JAX this will have worse performance.
             Specifying a number that tightly upper bounds the number of wells will
             increase performance. In general, an upper bound on the number of wells
-            per toroidal transit is ``Aι+B`` where ``A``,``B`` are the poloidal and
+            per toroidal transit is ``Aι+B`` where ``A``, ``B`` are the poloidal and
             toroidal Fourier resolution of B, respectively, in straight-field line
             PEST coordinates, and ι is the rotational transform normalized by 2π.
             A tighter upper bound than ``num_well=(Aι+B)*num_transit`` is preferable.
@@ -582,8 +579,8 @@ def integrate(
         Parameters
         ----------
         integrand : callable or list[callable]
-            The composition operator on the set of functions in ``data`` that
-            maps that determines ``f`` in ∫ f(ρ,α,λ,ℓ) dℓ. It should accept a dictionary
+            The composition operator on the set of functions in ``data``
+            that determines ``f`` in ∫ f(ρ,α,λ,ℓ) dℓ. It should accept a dictionary
             which stores the interpolated data and the arguments ``B`` and ``pitch``.
         pitch_inv : jnp.ndarray
             Shape (num rho, num pitch).
@@ -968,7 +965,7 @@ class Bounce1D(Bounce):
 
     See Also
     --------
-    Bounce2D : Uses two-dimensional pseudo-spectral techniques for the same task.
+    Bounce2D : Uses pseudo-spectral methods for the same task.
 
     Examples
     --------
@@ -1107,7 +1104,7 @@ def points(self, pitch_inv, *, num_well=None):
             but due to current limitations in JAX this will have worse performance.
             Specifying a number that tightly upper bounds the number of wells will
             increase performance. In general, an upper bound on the number of wells
-            per toroidal transit is ``Aι+B`` where ``A``,``B`` are the poloidal and
+            per toroidal transit is ``Aι+B`` where ``A``, ``B`` are the poloidal and
             toroidal Fourier resolution of B, respectively, in straight-field line
             PEST coordinates, and ι is the rotational transform normalized by 2π.
             A tighter upper bound than ``num_well=(Aι+B)*num_transit`` is preferable.
@@ -1193,8 +1190,8 @@ def integrate(
         Parameters
         ----------
         integrand : callable or list[callable]
-            The composition operator on the set of functions in ``data`` that
-            maps that determines ``f`` in ∫ f(ρ,α,λ,ℓ) dℓ. It should accept a dictionary
+            The composition operator on the set of functions in ``data``
+            that determines ``f`` in ∫ f(ρ,α,λ,ℓ) dℓ. It should accept a dictionary
             which stores the interpolated data and the arguments ``B`` and ``pitch``.
         pitch_inv : jnp.ndarray
             Shape (num alpha, num rho, num pitch).

desc/objectives/_fast_ion.py

@@ -47,6 +47,15 @@ class GammaC(_Objective):
     https://doi.org/10.1088/1741-4326/ac2994.
     Equation 16.
 
+    Notes
+    -----
+    Performance will improve significantly by resolving these GitHub issues.
+      * ``1154`` Improve coordinate mapping performance
+      * ``1294`` Nonuniform fast transforms
+      * ``1303`` Patch for differentiable code with dynamic shapes
+      * ``1206`` Upsample data above midplane to full grid assuming stellarator symmetry
+      * ``1034`` Optimizers/objectives with auxiliary output
+
     Parameters
     ----------
     eq : Equilibrium
@@ -78,7 +87,7 @@ class GammaC(_Objective):
         JAX this will have worse performance.
         Specifying a number that tightly upper bounds the number of wells will
         increase performance. In general, an upper bound on the number of wells
-        per toroidal transit is ``Aι+B`` where ``A``,``B`` are the poloidal and
+        per toroidal transit is ``Aι+B`` where ``A``, ``B`` are the poloidal and
         toroidal Fourier resolution of B, respectively, in straight-field line
         PEST coordinates, and ι is the rotational transform normalized by 2π.
         A tighter upper bound than ``num_well=(Aι+B)*num_transit`` is preferable.
@@ -103,15 +112,6 @@ class GammaC(_Objective):
         to be evaluated by measuring decrease in Γ_c at a fixed number of toroidal
         transits.
 
-    Notes
-    -----
-    Performance will improve significantly by resolving these GitHub issues.
-      * ``1154`` Improve coordinate mapping performance
-      * ``1294`` Nonuniform fast transforms
-      * ``1303`` Patch for differentiable code with dynamic shapes
-      * ``1206`` Upsample data above midplane to full grid assuming stellarator symmetry
-      * ``1034`` Optimizers/objectives with auxiliary output
-
     """
 
     __doc__ = __doc__.rstrip() + collect_docs(

desc/objectives/_neoclassical.py

@@ -46,6 +46,15 @@ class EffectiveRipple(_Objective):
     V. V. Nemov, S. V. Kasilov, W. Kernbichler, M. F. Heyn.
     Phys. Plasmas 1 December 1999; 6 (12): 4622–4632.
 
+    Notes
+    -----
+    Performance will improve significantly by resolving these GitHub issues.
+      * ``1154`` Improve coordinate mapping performance
+      * ``1294`` Nonuniform fast transforms
+      * ``1303`` Patch for differentiable code with dynamic shapes
+      * ``1206`` Upsample data above midplane to full grid assuming stellarator symmetry
+      * ``1034`` Optimizers/objectives with auxiliary output
+
     Parameters
     ----------
     eq : Equilibrium
@@ -77,7 +86,7 @@ class EffectiveRipple(_Objective):
         JAX this will have worse performance.
         Specifying a number that tightly upper bounds the number of wells will
         increase performance. In general, an upper bound on the number of wells
-        per toroidal transit is ``Aι+B`` where ``A``,``B`` are the poloidal and
+        per toroidal transit is ``Aι+B`` where ``A``, ``B`` are the poloidal and
         toroidal Fourier resolution of B, respectively, in straight-field line
         PEST coordinates, and ι is the rotational transform normalized by 2π.
         A tighter upper bound than ``num_well=(Aι+B)*num_transit`` is preferable.
@@ -91,15 +100,6 @@ class EffectiveRipple(_Objective):
         Number of pitch values with which to compute simultaneously.
         If given ``None``, then ``batch_size`` defaults to ``num_pitch``.
 
-    Notes
-    -----
-    Performance will improve significantly by resolving these GitHub issues.
-      * ``1154`` Improve coordinate mapping performance
-      * ``1294`` Nonuniform fast transforms
-      * ``1303`` Patch for differentiable code with dynamic shapes
-      * ``1206`` Upsample data above midplane to full grid assuming stellarator symmetry
-      * ``1034`` Optimizers/objectives with auxiliary output
-
     """
 
     __doc__ = __doc__.rstrip() + collect_docs(

docs/api_objectives.rst

@@ -35,6 +35,16 @@ Equilibrium
     desc.objectives.HelicalForceBalance
 
 
+Fast ion confinement
+--------------------
+.. autosummary::
+    :toctree: _api/objectives
+    :recursive:
+    :template: class.rst
+
+    desc.objectives.GammaC
+
+
 Geometry
 --------
 .. autosummary::
@@ -61,7 +71,6 @@ Neoclassical
     :template: class.rst
 
     desc.objectives.EffectiveRipple
-    desc.objectives.GammaC
 
 
 Omnigenity

docs/notebooks/tutorials/EffectiveRipple.ipynb

@@ -78,6 +78,7 @@
    "source": [
     "## Documentation\n",
     "Please read the full documentation of the methods to understand what the input parameters do. In Jupyter Lab, you can click on the code and press ``Shift+Tab`` to pull up the documentation. Breifly,\n",
+    "\n",
     "- The equilibrium resolution determines the spectral resolution of the FourierZernike series fit to the boundary.\n",
     "- The grid determines the flux surfaces to compute on and the resolution of FFTs.\n",
     "- The parameters ``X`` and ``Y`` determine the spectral resolution of the map between coordinates that parameterize the boundary and field line coordinates.\n",
@@ -86,7 +87,6 @@
     "## Plotting ripple wells\n",
     "\n",
     "- Here we plot $\\vert B\\vert$ along field lines to see the structure of the ripple wells. This is beneficial to choose the resolution for the optimization.\n",
-    "\n",
     "- Due to limitations in JAX, it is recommended to plot the field lines and pick a reasonable, yet preferably tight, upper bound on the number of ripple wells. From the plots, we see that ``num_well=10 * num_transit`` is a reasonable upper bound. By making this extra effort, the optimization will be ``Y_B/10`` times more performant. If one were to select something much less than ``10``, as shown in the next example, then it should be clear from the plot that some ripple wells are ignored, which is not desirable."
    ]
   },
@@ -134,7 +134,7 @@
     "        JAX this will have worse performance.\n",
     "        Specifying a number that tightly upper bounds the number of wells will\n",
     "        increase performance. In general, an upper bound on the number of wells\n",
-    "        per toroidal transit is ``Aι+B`` where ``A``,``B`` are the poloidal and\n",
+    "        per toroidal transit is ``Aι+B`` where ``A``, ``B`` are the poloidal and\n",
     "        toroidal Fourier resolution of B, respectively, in straight-field line\n",
     "        PEST coordinates, and ι is the rotational transform normalized by 2π.\n",
     "        A tighter upper bound than ``num_well=(Aι+B)*num_transit`` is preferable.\n",
@@ -374,7 +374,7 @@
    "source": [
     "## Calculating effective ripple for Heliotron\n",
     "\n",
-    "Let us do a high resolution compuation so that we are certain the optimization is successful when we compare to the optimized result later."
+    "Let us do a high resolution computation so that we are certain the optimization is successful when we compare to the optimized result later."
    ]
   },
   {