From 274352cb4a9ccf96009c0149d5b0e35e8d209ec8 Mon Sep 17 00:00:00 2001 From: unalmis Date: Thu, 19 Dec 2024 06:31:54 -0500 Subject: [PATCH] Fix docs --- desc/compute/_neoclassical.py | 2 +- desc/integrals/_bounce_utils.py | 8 +-- desc/integrals/bounce_integral.py | 57 +++++++++---------- desc/objectives/_fast_ion.py | 20 +++---- desc/objectives/_neoclassical.py | 20 +++---- docs/api_objectives.rst | 11 +++- .../notebooks/tutorials/EffectiveRipple.ipynb | 6 +- 7 files changed, 65 insertions(+), 59 deletions(-) diff --git a/desc/compute/_neoclassical.py b/desc/compute/_neoclassical.py index 0ca82231c..f2ef2b495 100644 --- a/desc/compute/_neoclassical.py +++ b/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. diff --git a/desc/integrals/_bounce_utils.py b/desc/integrals/_bounce_utils.py index b79d3497c..4680269a4 100644 --- a/desc/integrals/_bounce_utils.py +++ b/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 diff --git a/desc/integrals/bounce_integral.py b/desc/integrals/bounce_integral.py index cc7bc0dbf..6adebf166 100644 --- a/desc/integrals/bounce_integral.py +++ b/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). diff --git a/desc/objectives/_fast_ion.py b/desc/objectives/_fast_ion.py index b4cddbd20..bb6e020db 100644 --- a/desc/objectives/_fast_ion.py +++ b/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( diff --git a/desc/objectives/_neoclassical.py b/desc/objectives/_neoclassical.py index 02501ed6c..7324d3fd9 100644 --- a/desc/objectives/_neoclassical.py +++ b/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( diff --git a/docs/api_objectives.rst b/docs/api_objectives.rst index 145a2898d..64603268e 100644 --- a/docs/api_objectives.rst +++ b/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 diff --git a/docs/notebooks/tutorials/EffectiveRipple.ipynb b/docs/notebooks/tutorials/EffectiveRipple.ipynb index 72feff581..ef4119f38 100644 --- a/docs/notebooks/tutorials/EffectiveRipple.ipynb +++ b/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." ] }, {