Neoclassical transport and fast ions

• 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 EffectiveRipple with GammaC.

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 ε < \(10^{-2}\).

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 is assumed 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.

• Spectrally accurate, reverse-mode differentiable bounce-averaging algorithm and its applications. Kaya Unalmis et al. Journal of Plasma Physics.

# 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!

# 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.

# 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 np
from matplotlib import pyplot as plt

from desc.integrals import Bounce2D
from desc.examples import get
from desc.grid import LinearGrid, Grid
from desc.optimize import Optimizer
from desc.objectives import (
    ForceBalance,
    FixPsi,
    FixBoundaryR,
    FixBoundaryZ,
    GenericObjective,
    FixPressure,
    FixIota,
    AspectRatio,
    EffectiveRipple,
    ObjectiveFunction,
)

An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.

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. Briefly,

• 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 X and Y determine the spectral resolution of the map between coordinates that parameterize the boundary and field line coordinates.

Frequency spectrum of the map to the mesh of field lines

• First, let us view the spectrum of the map between coordinates that parameterize the boundary and coordinates that parameterize the magnetic field line trajectory.

• This should give you some intuition on how to pick the resolution for the parameters X and Y.

• For example, let us view this spectrum on a few different flux surfaces of the precise QH equilibrium from the DESC examples folder.

eq = get("precise_QH")
rho = np.linspace(0.01, 1, 2)
angle = Bounce2D.angle(eq, X=32, Y=32, rho=rho)
for l in range(rho.size):
    fig = Bounce2D.plot_angle_spectrum(angle, l)

• If you need higher resolution, you should make sure the maximum error allowed by the Newton solve with the tol argument to .angle() is consistent.

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_field_periods with W=3 is a reasonable upper bound. By making this extra effort, the optimization will be more performant. If one were to select something much smaller, as shown in the next example, then it should be clear from the plot that some ripple wells are ignored, which is not desirable.

• Making a good choice for num_well is important for performance in optimization.

def plot_wells(
    eq,
    grid,
    angle,
    *,
    Y_B=None,
    spline=True,
    num_pitch=10,
    alpha=None,
    num_field_periods=3,
    num_well=None,
):
    """Plotting tool to help user set tighter upper bound on ``num_well``.

    Also prints error statistics for the bounce points.

    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.
    angle : jnp.ndarray
        Shape (num rho, X, Y).
        Angle returned by ``Bounce2D.angle``.
    Y_B : int
        Desired resolution for algorithm to compute bounce points.
        A reference value is ``(grid.num_theta+grid.num_zeta)//2``.

        If the option ``spline`` is ``True``, the bounce points are found with
        𝒪(Y_B⁻¹²) error. In this case, the final error will be of order
        𝒪(Y_B⁻¹⁸) in bounce integrals with (v_∥)¹ and
        𝒪(Y_B⁻⁶)  in bounce integrals with (v_∥)⁻¹.

        If the option ``spline`` is ``False``, the bounce points are found such
        that the bounce integrals have exponential accuracy in this parameter.
    spline : bool
        Whether to use cubic splines to compute initial guess for bounce points
        instead of Chebyshev series. Default is ``True``. It can be preferable
        to set to ``False`` on equilibria with high ``NFP``, (such cases make
        smaller ``Y_B`` feasible), or on GPUs where eigenvalue solves are fast.
    num_pitch: int
        Number of pitch angles.
    alpha : jnp.ndarray
        Shape (num α, ).
        Starting field line poloidal labels.
        Default is single field line.
        On irrational magnetic surfaces, it is sufficient to integrate along a
        single field line. On a rational or near-rational surface in
        non-axisymmetric configurations, it is necessary to integrate along
        multiple field lines until the surface is covered sufficiently.
    num_field_periods : int
        Number of field periods to follow field line.
        In axisymmetric configurations, integration along the field line for a
        single poloidal transit between two global maxima of B is sufficient for
        convergence. For a 3D configuration, the magnetic surface should be covered
        sufficiently.
    num_well : int
        Maximum number of wells to detect for each pitch and field line.
        Giving ``-1`` 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ι+C`` where ``A``, ``C`` 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ι+C)*num_transit`` is preferable.
        The ``check_points`` or ``plot`` methods in ``desc.integrals.Bounce2D``
        are useful to select a reasonable value.

    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=grid,
        data=data,
        angle=angle,
        Y_B=Y_B,
        spline=spline,
        alpha=alpha,
        num_field_periods=num_field_periods,
        quad=(np.zeros(0), np.zeros(0)),
    )
    pitch_inv, _ = Bounce2D.pitch_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 plot the magnetic field norm over 12 field periods to educate our guess for the num_well parameter.

grid = LinearGrid(rho=rho, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=False)
angle = Bounce2D.angle(eq, X=16, Y=16, rho=rho)
num_field_periods = 12

plot_wells(
    eq,
    grid,
    angle,
    num_field_periods=num_field_periods,
    num_well=3 * num_field_periods,
);

Error statistics for the given bounce points:
After 1 iteration(s) | ζ₁₂(w) error mean = 1e-08 | std. dev. = 3e-08 | max = 3e-07

• Let us plot these again with a lower choice for num_well.

• We observe some wells are not detected. Notice the triangular markers that indicate the detection of a bounce point are missing.

• These plots are better viewed in interactive Python sessions where you can zoom in etc. An interactive session is usually started automatically if running a Python file instead of a Jupyter notebook cell.

plot_wells(
    eq,
    grid,
    angle,
    num_field_periods=num_field_periods,
    num_well=num_field_periods // 3,
);

Error statistics for the given bounce points:
After 1 iteration(s) | ζ₁₂(w) error mean = 2e-08 | std. dev. = 3e-08 | max = 3e-07

Calculating effective ripple for Precise QH

• If your machine has insufficient memory, then you should specify the parameter pitch_batch_size to some value less than num_pitch. Recall this determines the number of pitch values to compute simultaneously.

• Some machines detect insufficient memory early, exit early, and return ε = 0 everywhere. On most machines you should get a typical OOM error.

rho = np.linspace(0.01, 1, 10)
grid = LinearGrid(rho=rho, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=False)
num_field_periods = 32
data = eq.compute(
    "effective ripple",
    grid,
    angle=Bounce2D.angle(eq, X=16, Y=16, rho=rho),
    num_field_periods=num_field_periods,
    num_well=3 * num_field_periods,
    num_quad=16,
    num_pitch=15,
)
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

rho = np.linspace(0.01, 1, 2)
eq0 = get("HELIOTRON")
angle = Bounce2D.angle(eq0, X=40, Y=20, rho=rho)
for l in range(rho.size):
    fig = Bounce2D.plot_angle_spectrum(angle, l)

plot_wells(
    eq0,
    LinearGrid(rho=rho, M=eq0.M_grid, N=eq0.N_grid, NFP=eq0.NFP, sym=False),
    angle,
    num_field_periods=15,
    num_well=2 * 15,
);

Error statistics for the given bounce points:
After 1 iteration(s) | ζ₁₂(w) error mean = 5e-11 | std. dev. = 7e-11 | max = 7e-10

rho = np.linspace(0.01, 1, 10)
angle = Bounce2D.angle(eq0, X=32, Y=20, rho=rho)
grid = LinearGrid(rho=rho, M=eq0.M_grid, N=eq0.N_grid, NFP=eq0.NFP, sym=False)
num_field_periods = 398
num_well = 2 * num_field_periods
num_quad = 48
num_pitch = 45
data = eq0.compute(
    "effective ripple",
    grid,
    angle=angle,
    num_field_periods=num_field_periods,
    num_well=num_well,
    num_quad=num_quad,
    num_pitch=num_pitch,
)
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()

Let us also compute this quantity using an older lower-order accurate algorithm and compare the results.

# This cell will consume significant memory.
knots_per_field_period = (grid.num_zeta + grid.num_theta) // 2
low_order_eps_grid = Grid.create_meshgrid(
    [
        rho,
        np.array([0]),
        np.linspace(
            0,
            num_field_periods * 2 * np.pi / eq0.NFP,
            num_field_periods * knots_per_field_period,
        ),
    ],
    coordinates="raz",
)
low_order_eps = eq0.compute(
    "old effective ripple",
    low_order_eps_grid,
    num_well=num_well,
    num_quad=num_quad,
    num_pitch=num_pitch,
)["old effective ripple"]
low_order_eps = low_order_eps_grid.compress(low_order_eps)
del low_order_eps_grid
del knots_per_field_period

fig, ax = plt.subplots()
ax.set(
    xlabel=r"$\rho$",
    ylabel=r"$\epsilon$",
    title="Heliotron NFP=19, effective ripple comparison.",
)
ax.plot(rho, eps, marker="o", label="effective ripple")
ax.plot(rho, low_order_eps, marker="o", label="low order accurate effective ripple")
plt.legend()
plt.tight_layout()
plt.show()

Optimizing Heliotron

eq1 = eq0.copy()
k = 1
print()
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, 5), M=eq1.M_grid, N=eq1.N_grid, NFP=eq1.NFP, sym=False
)
objective = ObjectiveFunction(
    (
        EffectiveRipple(
            eq1,
            grid=ripple_grid,
            X=32,
            Y=20,
            num_field_periods=190,
            num_well=250,
            num_quad=32,
            num_pitch=45,
        ),
        AspectRatio(eq1, bounds=(8, 11), weight=1e2),
        GenericObjective(
            "curvature_k2_rho",
            eq1,
            grid=curvature_grid,
            bounds=(-128, 10),
            weight=2e2,
        ),
    )
)
# does an expensive qr/svd solve at each step
optimizer = Optimizer("proximal-lsq-exact")
(eq1,), _ = optimizer.optimize(
    eq1,
    objective,
    constraints,
    ftol=1e-4,
    xtol=1e-6,
    gtol=1e-6,
    maxiter=7,
    verbose=3,
    options={"initial_trust_ratio": 2e-3},
)
print("Optimization complete!")

---------------------------------------
Optimizing boundary modes M, N <= 1
---------------------------------------
Building objective: Effective ripple
Building objective: aspect ratio
Precomputing transforms
Timer: Precomputing transforms = 109 ms
Building objective: Generic
Timer: Objective build = 4.37 sec
Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 156 ms
Timer: Objective build = 1.14 sec
Timer: Objective build = 3.10 ms
Timer: Eq Update LinearConstraintProjection build = 6.17 sec
Timer: Proximal projection build = 47.3 sec
Building objective: lcfs R
Building objective: lcfs Z
Building objective: fixed pressure
Building objective: fixed iota
Building objective: fixed Psi
Timer: Objective build = 1.46 sec
Timer: LinearConstraintProjection build = 1.88 sec
Number of parameters: 8
Number of objectives: 253
Timer: Initializing the optimization = 50.7 sec

Starting optimization
Using method: proximal-lsq-exact
Solver options:
------------------------------------------------------------
Maximum Function Evaluations       : 36
Maximum Allowed Total Δx Norm      : inf
Scaled Termination                 : True
Trust Region Method                : qr
Initial Trust Radius               : 1.677e-02
Maximum Trust Radius               : inf
Minimum Trust Radius               : 2.220e-16
Trust Radius Increase Ratio        : 2.000e+00
Trust Radius Decrease Ratio        : 2.500e-01
Trust Radius Increase Threshold    : 7.500e-01
Trust Radius Decrease Threshold    : 2.500e-01
------------------------------------------------------------

   Iteration     Total nfev        Cost      Cost reduction    Step norm     Optimality
       0              1          4.105e-01                                    8.483e-01
       1              2          3.847e-01      2.579e-02      4.576e-03      7.533e-01
       2              3          3.434e-01      4.136e-02      8.062e-03      6.567e-01
       3              4          2.743e-01      6.902e-02      1.538e-02      5.024e-01
       4              5          1.881e-01      8.630e-02      2.817e-02      3.293e-01
       5              7          1.608e-01      2.721e-02      2.136e-02      2.597e-01
       6              9          1.483e-01      1.250e-02      7.956e-03      2.419e-01
       7             10          1.257e-01      2.265e-02      1.889e-02      1.754e-01
Warning: Maximum number of iterations has been exceeded.
         Current function value: 1.257e-01
         Total delta_x: 8.343e-02
         Iterations: 7
         Function evaluations: 10
         Jacobian evaluations: 8
Timer: Solution time = 14.2 min
Timer: Avg time per step = 1.78 min
==============================================================================================================
                                                                 Start  -->   End
Total (sum of squares):                                      4.105e-01  -->   1.257e-01,
Maximum absolute Effective ripple ε:                         4.496e-01  -->   2.994e-01 ~
Minimum absolute Effective ripple ε:                         2.495e-01  -->   1.268e-01 ~
Average absolute Effective ripple ε:                         3.982e-01  -->   2.158e-01 ~
Maximum absolute Effective ripple ε:                         4.496e-01  -->   2.994e-01 (normalized)
Minimum absolute Effective ripple ε:                         2.495e-01  -->   1.268e-01 (normalized)
Average absolute Effective ripple ε:                         3.982e-01  -->   2.158e-01 (normalized)
Aspect ratio:                                                1.048e+01  -->   1.100e+01 (dimensionless)
Maximum Generic objective value:                            -6.864e-01  -->  -6.858e-01 (m^{-1})
Minimum Generic objective value:                            -5.858e+00  -->  -6.339e+00 (m^{-1})
Average Generic objective value:                            -1.566e+00  -->  -1.628e+00 (m^{-1})
Maximum Generic objective value:                            -6.864e-01  -->  -6.858e-01 (normalized)
Minimum Generic objective value:                            -5.858e+00  -->  -6.339e+00 (normalized)
Average Generic objective value:                            -1.566e+00  -->  -1.628e+00 (normalized)
Maximum absolute Force error:                                5.503e+03  -->   6.476e+03 (N)
Minimum absolute Force error:                                1.430e-02  -->   5.463e-03 (N)
Average absolute Force error:                                7.043e+01  -->   4.982e+01 (N)
Maximum absolute Force error:                                4.426e-04  -->   5.208e-04 (normalized)
Minimum absolute Force error:                                1.150e-09  -->   4.394e-10 (normalized)
Average absolute Force error:                                5.664e-06  -->   4.006e-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!

data = eq1.compute(
    "effective ripple",
    grid,
    angle=Bounce2D.angle(eq1, X=40, Y=20, rho=rho),
    num_field_periods=num_field_periods,
    num_well=num_well,
    num_quad=num_quad,
    num_pitch=num_pitch,
)
eps_opt = grid.compress(data["effective ripple"])

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();

from desc.plotting import plot_comparison

plot_comparison(eqs=[eq0, eq1], labels=["original", "optimized"]);