In this tutorial, we will show how to optimize for the effective ripple in DESC. The computation involves integration over ripple wells whose structure determines the optimal resolution for the optimization. So we will also breifly show how to visualize the ripples and accordingly pick resolution parameters. The same tutorial can be used to optimize for fast ion confinement with Γ_c. To do so, replace the objective
EffectiveRipplewithGammaC.Note that there is still work in progress to improve the performance in DESC by an order of magnitude. See the GitHub issues linked in the objective docstring if you would like to contribute.
Neoclassical transport in banana regime
A 3D stellarator magnetic field admits ripple wells that lead to enhanced radial drift of trapped particles. In the banana regime, neoclassical (thermal) transport from ripple wells can become the dominant transport channel. The effective ripple (ε) proxy estimates the neoclassical transport coefficients in the banana regime. To ensure low neoclassical transport, a stellarator is typically optimized so that ε < 0.02.
Fast ion confinement
The energetic particle confinement metric γ_c quantifies whether the contours of the second adiabatic invariant close on the flux surfaces. In the limit where the poloidal drift velocity dominates the radial drift velocity, the contours lie parallel to flux surfaces. The optimization metric Γ_c averages γ_c² over the distribution of trapped particles on each flux surface. The radial electric field has a negligible effect, since fast particles have high energy with collisionless orbits, so it isassumed to be zero.
References
Evaluation of 1/ν neoclassical transport in stellarators. V. V. Nemov, S. V. Kasilov, W. Kernbichler, M. F. Heyn. Phys. Plasmas 1 December 1999; 6 (12): 4622–4632.
Poloidal motion of trapped particle orbits in real-space coordinates. V. V. Nemov, S. V. Kasilov, W. Kernbichler, G. O. Leitold. Phys. Plasmas 1 May 2008; 15 (5): 052501.
[1]:
# If DESC is not installed as described in the installation documentation,# then these lines may be needed to run this notebook.## import sys# import os# sys.path.insert(0, os.path.abspath("."))# sys.path.append(os.path.abspath("../../../"))
If you have access to a GPU, uncomment the following two lines before any DESC or JAX related imports. You should see about an order of magnitude speed improvement with only these two lines of code!
[2]:
# from desc import set_device# set_device("gpu")
As mentioned in DESC Documentation on performance tips, one can use compilation cache directory to reduce the compilation overhead time. Note: One needs to create jax-caches folder manually.
[3]:
# import jax# jax.config.update("jax_compilation_cache_dir", "../jax-caches")# jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)# jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
[ ]:
import numpy as npfrom matplotlib import pyplot as pltfrom desc.integrals import Bounce2Dfrom desc.examples import getfrom desc.grid import LinearGridfrom desc.optimize import Optimizerfrom desc.objectives import ( ForceBalance, FixPsi, FixBoundaryR, FixBoundaryZ, GenericObjective, FixPressure, FixIota, AspectRatio, EffectiveRipple, ObjectiveFunction,)
Documentation
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,
The equilibrium resolution determines the spectral resolution of the FourierZernike series fit to the boundary.
The grid determines the flux surfaces to compute on and the resolution of FFTs.
The parameters
XandYdetermine the spectral resolution of the map between coordinates that parameterize the boundary and field line coordinates.The parameter
Y_Bdetermines the resolution for the bounce point finding algorithm. Feel free to reduce this until the plots of \(\vert B\vert\) along field lines do not change. If \(\vert B\vert\) is high frequency, then a larger value will be needed (usually much larger thanY).
Plotting ripple wells
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.
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=W * num_transitwithW=10is a reasonable upper bound. By making this extra effort, the optimization will beY_B/Wtimes more performant. If one were to select something much less than10, as shown in the next example, then it should be clear from the plot that some ripple wells areignored, which is not desirable.
[5]:
def plot_wells( eq, grid, theta, Y_B=None, num_transit=3, num_well=None, num_pitch=10,): """Plotting tool to help user set tighter upper bound on ``num_well``. Parameters ---------- eq : Equilibrium Equilibrium to compute on. grid : Grid Tensor-product grid in (ρ, θ, ζ) with uniformly spaced nodes (θ, ζ) ∈ [0, 2π) × [0, 2π/NFP). Number of poloidal and toroidal nodes preferably rounded down to powers of two. Determines the flux surfaces to compute on and resolution of FFTs. theta : jnp.ndarray Shape (num rho, X, Y). DESC coordinates θ sourced from the Clebsch coordinates ``FourierChebyshevSeries.nodes(X,Y,rho,domain=(0,2*jnp.pi))``. Use the ``Bounce2D.compute_theta`` method to obtain this. ``X`` and ``Y`` are preferably rounded down to powers of two. Y_B : int Desired resolution for algorithm to compute bounce points. Default is double ``Y``. num_transit : int Number of toroidal transits to follow field line. In an axisymmetric device, field line integration over a single poloidal transit is sufficient to capture a surface average. For a 3D configuration, more transits will approximate surface averages on an irrational magnetic surface better, with diminishing returns. num_well : int Maximum number of wells to detect for each pitch and field line. Giving ``None`` will detect all wells 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 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. The ``check_points`` or ``plot`` methods in ``desc.integrals.Bounce2D`` are useful to select a reasonable value. num_pitch: int Number of pitch angles. Returns ------- plots Matplotlib (fig, ax) tuples for the 1D plot of each field line. """ data = eq.compute(Bounce2D.required_names + ["min_tz |B|", "max_tz |B|"], grid=grid) bounce = Bounce2D(grid, data, theta, Y_B, num_transit=num_transit) pitch_inv, _ = Bounce2D.get_pitch_inv_quad( grid.compress(data["min_tz |B|"]), grid.compress(data["max_tz |B|"]), num_pitch, ) points = bounce.points(pitch_inv, num_well) plots = bounce.check_points(points, pitch_inv) return plots
We will plot the ripple wells for precise QH from DESC examples folder on multiple flux surfaces.
[6]:
# ---------- Precise QH ----------# Computing at higher resolution than necessary.eq0 = get("precise_QH")rho = np.linspace(0.01, 1, 5)grid = LinearGrid(rho=rho, M=eq0.M_grid, N=eq0.N_grid, NFP=eq0.NFP, sym=False)# ---------- How to pick resolution? ----------# Plotting for 3 toroidal transits to see by eye# Seems like these resolutions are more than sufficient.# We will use more pitch angles for the integration.X, Y = 16, 32theta = Bounce2D.compute_theta(eq0, X, Y, rho)num_transit = 3Y_B = 32plot_wells( eq0, grid, theta, Y_B=Y_B, num_transit=num_transit, num_well=10 * num_transit,);
[7]:
plot_wells( eq0, grid, theta, Y_B=Y_B, num_transit=num_transit, # Here we see some wells are ignored if num_well is too low. num_well=1 * num_transit,);
Calculating effective ripple for Precise QH
[8]:
num_transit = 20num_well = 10 * num_transitnum_quad = 32num_pitch = 45data = eq0.compute( "effective ripple", grid=grid, theta=theta, Y_B=Y_B, num_transit=num_transit, num_well=num_well, num_quad=num_quad, num_pitch=num_pitch, # Can also specify ``pitch_batch_size`` which determines the # number of pitch values to compute simultaneously. # Reduce this if insufficient memory. If insufficient memory is detected # early then the code will exit and return ε = 0 everywhere. If not detected # early then typical OOM errors will occur.)eps = grid.compress(data["effective ripple"])fig, ax = plt.subplots()ax.plot(rho, eps, marker="o")ax.set(xlabel=r"$\rho$", ylabel=r"$\epsilon$", title="Precise QH")plt.tight_layout()plt.show()
Calculating effective ripple for Heliotron
Let us do a high resolution computation so that we are certain the optimization is successful when we compare to the optimized result later.
[9]:
eq0 = get("HELIOTRON")rho = np.linspace(0.01, 1, 10)grid = LinearGrid(rho=rho, M=eq0.M_grid, N=eq0.N_grid, NFP=eq0.NFP, sym=False)X = 32Y = 64Y_B = 133num_transit = 20num_well = 30 * num_transitnum_quad = 64data = eq0.compute( "effective ripple", grid=grid, theta=Bounce2D.compute_theta(eq0, X, Y, rho=rho), Y_B=Y_B, num_transit=num_transit, num_well=num_well, num_quad=num_quad,)eps = grid.compress(data["effective ripple"])fig, ax = plt.subplots()ax.plot(rho, eps, marker="o")ax.set(xlabel=r"$\rho$", ylabel=r"$\epsilon$", title="Heliotron NFP=19")plt.tight_layout()plt.show()
Optimizing Heliotron
[ ]:
eq1 = eq0.copy()k = 1print()print("---------------------------------------")print(f"Optimizing boundary modes M, N <= {k}")print("---------------------------------------")modes_R = np.vstack( ( [0, 0, 0], eq1.surface.R_basis.modes[np.max(np.abs(eq1.surface.R_basis.modes), 1) > k, :], ))modes_Z = eq1.surface.Z_basis.modes[np.max(np.abs(eq1.surface.Z_basis.modes), 1) > k, :]constraints = ( ForceBalance(eq=eq1), FixBoundaryR(eq=eq1, modes=modes_R), FixBoundaryZ(eq=eq1, modes=modes_Z), FixPressure(eq=eq1), FixIota(eq=eq1), FixPsi(eq=eq1),)curvature_grid = LinearGrid( rho=np.array([1.0]), M=eq1.M_grid, N=eq1.N_grid, NFP=eq1.NFP, sym=eq1.sym)ripple_grid = LinearGrid( rho=np.linspace(0.2, 1, 3), M=eq1.M_grid, N=eq1.N_grid, NFP=eq1.NFP, sym=False)objective = ObjectiveFunction( ( EffectiveRipple( eq1, grid=ripple_grid, X=16, Y=32, Y_B=133, num_transit=10, num_well=25 * 10, num_quad=32, num_pitch=45, jac_chunk_size=1, # to reduce the memory usage ), AspectRatio(eq1, bounds=(8, 11), weight=1e3), GenericObjective( "curvature_k2_rho", eq1, grid=curvature_grid, bounds=(-128, 10), weight=2e3 ), ))optimizer = Optimizer("proximal-lsq-exact")(eq1,), _ = optimizer.optimize( eq1, objective, constraints, ftol=1e-4, xtol=1e-6, gtol=1e-6, maxiter=5, verbose=3, options={"initial_trust_ratio": 2e-3},)print("Optimization complete!")
---------------------------------------Optimizing boundary modes M, N <= 1---------------------------------------Building objective: Effective rippleBuilding objective: aspect ratioPrecomputing transformsTimer: Precomputing transforms = 74.9 msBuilding objective: genericTimer: Objective build = 2.33 secBuilding objective: forcePrecomputing transformsTimer: Precomputing transforms = 538 msTimer: Objective build = 1.12 secTimer: Proximal projection build = 10.7 secBuilding objective: lcfs RBuilding objective: lcfs ZBuilding objective: fixed pressureBuilding objective: fixed iotaBuilding objective: fixed PsiTimer: Objective build = 645 msTimer: Linear constraint projection build = 2.29 secNumber of parameters: 8Number of objectives: 251Timer: Initializing the optimization = 13.7 secStarting optimizationUsing method: proximal-lsq-exact Iteration Total nfev Cost Cost reduction Step norm Optimality 0 1 5.009e-01 9.640e-01 1 2 4.719e-01 2.902e-02 5.598e-03 9.210e-01 2 3 4.043e-01 6.754e-02 1.060e-02 8.320e-01 3 4 3.247e-01 7.962e-02 1.371e-02 6.201e-01 4 5 1.826e-01 1.421e-01 3.413e-02 4.133e-01 5 7 1.455e-01 3.708e-02 1.054e-02 4.318e-01Warning: Maximum number of iterations has been exceeded. Current function value: 1.455e-01 Total delta_x: 7.093e-02 Iterations: 5 Function evaluations: 7 Jacobian evaluations: 6Timer: Solution time = 10.7 minTimer: Avg time per step = 1.79 min============================================================================================================== Start --> EndTotal (sum of squares): 5.009e-01 --> 1.455e-01,Maximum absolute Effective ripple ε: 6.967e-01 --> 3.680e-01 ~Minimum absolute Effective ripple ε: 4.967e-01 --> 1.936e-01 ~Average absolute Effective ripple ε: 5.709e-01 --> 3.017e-01 ~Maximum absolute Effective ripple ε: 6.967e-01 --> 3.680e-01 (normalized)Minimum absolute Effective ripple ε: 4.967e-01 --> 1.936e-01 (normalized)Average absolute Effective ripple ε: 5.709e-01 --> 3.017e-01 (normalized)Aspect ratio: 1.048e+01 --> 1.094e+01 (dimensionless)Maximum Generic objective value: -6.864e-01 --> -6.904e-01 (m^{-1})Minimum Generic objective value: -5.858e+00 --> -6.371e+00 (m^{-1})Average Generic objective value: -1.566e+00 --> -1.638e+00 (m^{-1})Maximum Generic objective value: -6.864e-01 --> -6.904e-01 (normalized)Minimum Generic objective value: -5.858e+00 --> -6.371e+00 (normalized)Average Generic objective value: -1.566e+00 --> -1.638e+00 (normalized)Maximum absolute Force error: 5.705e+03 --> 5.130e+03 (N)Minimum absolute Force error: 2.188e-02 --> 7.922e-03 (N)Average absolute Force error: 7.113e+01 --> 5.279e+01 (N)Maximum absolute Force error: 4.588e-04 --> 4.126e-04 (normalized)Minimum absolute Force error: 1.760e-09 --> 6.371e-10 (normalized)Average absolute Force error: 5.720e-06 --> 4.246e-06 (normalized)R boundary error: 0.000e+00 --> 0.000e+00 (m)Z boundary error: 0.000e+00 --> 0.000e+00 (m)Fixed pressure profile error: 0.000e+00 --> 0.000e+00 (Pa)Fixed iota profile error: 0.000e+00 --> 0.000e+00 (dimensionless)Fixed Psi error: 0.000e+00 --> 0.000e+00 (Wb)==============================================================================================================Optimization complete![11]:
data = eq1.compute( "effective ripple", grid=grid, theta=Bounce2D.compute_theta(eq1, X, Y, rho=rho), num_transit=num_transit, num_well=num_well, num_quad=num_quad,)eps_opt = grid.compress(data["effective ripple"])
[12]:
fig, ax = plt.subplots()ax.plot(rho, eps, marker="o", label="original")ax.plot(rho, eps_opt, marker="o", label="optimized")ax.set(xlabel=r"$\rho$", ylabel=r"$\epsilon$", title="Heliotron NFP=19")ax.legend();
[13]:
from desc.plotting import plot_comparisonplot_comparison(eqs=[eq0, eq1], labels=["original", "optimized"]);
