Infinite-\(n\) ideal ballooning stability optimization

This tutorial demonstrates how to evaluate and optimize an equilibrium for infinite-\(n\) ideal-ballooning stability with DESC. The infinite-\(n\) ideal ballooning equation is

\[\mathbf{B} \cdot \nabla \left( \frac{|\nabla \alpha|^2}{B^2} \mathbf{B} \cdot \nabla X \right) + \frac{2}{B^2} \frac{d \mu_0 p}{d\psi} \left[\mathbf{B} \times (\mathbf{b} \cdot \nabla \mathbf{b})\right] \cdot \nabla \alpha\, X = \lambda \frac{|\nabla \alpha|^2}{B^2} X, \quad \lambda = \gamma^2\]

where \(X\) is the ballooning eigenfunction whereas \(\gamma\) is the ballooning eigenvalue. When \(\lambda > 0\), the mode is ballooning-unstable, otherwise it is ballooning-stable. The equation is solved subject to Dirichlet boundary conditions on the eigenfunction \(X\)

\[X(\zeta = \zeta_1) = 0, X(\zeta = \zeta_2) = 0\]

where \(\zeta_1\) and \(\zeta_2\) are endpoints of the domain in ballooning space (a transformed covering space where \(\zeta_{\mathrm{DESC}} \in [0, 2\pi]\) to \(\zeta \in [-\infty, \infty]\)). Note that the ballooning \(\zeta\) is a different coordinate than periodic \(\zeta_{\mathrm{DESC}}\) due to their range of applicability.

Numerically, this equation becomes a 1D EVP (eigenvalue problem) along a field line.

\[\frac{d}{d\zeta} \left(g \frac{dX}{d\zeta} \right) + c X = \lambda f X\]

where

\[\begin{split}\begin{align*} \text{g} &= (\mathbf{b} \cdot \nabla \zeta) \frac{|\nabla \alpha|^2}{B}, \\ \text{c} &= \frac{1}{B^2} \frac{d(\mu_0 p)}{d\psi} \frac{2}{(\mathbf{b} \cdot \nabla \zeta)} (\mathbf{b} \times (\mathbf{b} \cdot \nabla \mathbf{b})) \cdot \nabla \alpha, \\ \text{f} &= \frac{1}{(\mathbf{b} \cdot \nabla \zeta)} \frac{|\nabla \alpha|^2}{B^3}, \end{align*}\end{split}\]

are functions of \(\zeta\) along a field line.

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 matplotlib.pyplot as plt
import numpy as np

import desc
from desc.grid import Grid, LinearGrid
from desc.optimize import Optimizer
from desc.objectives import (
    ForceBalance,
    AspectRatio,
    FixIota,
    FixPressure,
    FixPsi,
    PrincipalCurvature,
    BallooningStability,
    ObjectiveFunction,
    FixBoundaryR,
    FixBoundaryZ,
    GenericObjective,
)

plt.rcParams["font.size"] = 14

Evaluating ballooning stability of the initial equilibrium

In this section, we will show how you can take a DESC or VMEC equilibrium and evaluate its ballooning stability using DESC You have to specify the normalized distance \(\rho\) and the field line label \(\alpha\) on which you want to solve the ballooning equation.

DESC will solve the ballooning equation on multiple field lines on each flux surface and output the maximum ballooning growth rate on each flux surface.

\[\lambda_{\rho, \mathrm{max}} = \mathrm{max}_{\alpha}(\lambda_{\rho, \alpha}), \quad \forall \quad \alpha \in [0, \pi]\]

# Importing the HELIOTRON DESC equilibrium
eq0 = desc.examples.get("HELIOTRON")
# reduce resolution so it runs a bit quicker
eq0.change_resolution(L=12, M=6, N=2, L_grid=18, M_grid=12, N_grid=4)
eq0.surface = eq0.get_surface_at(rho=1)

/CODES/DESC/desc/utils.py:572: UserWarning: Reducing radial (L) resolution can make plasma boundary inconsistent. Recommend calling `eq.surface = eq.get_surface_at(rho=1.0)`
  warnings.warn(msg, err)

# Flux surfaces on which to evaluate ballooning stability
surfaces = np.array([0.01, 0.1, 0.2, 0.4, 0.6, 0.8, 1.0])
# Field lines on which to evaluate ballooning stability
alpha = np.linspace(0, np.pi, 8, endpoint=False)
# Number of toroidal transits of the field line
nturns = 3
# Number of point along a field line in ballooning space
N0 = nturns * 200
# range of the ballooning coordinate zeta
zeta = np.linspace(-np.pi * nturns, np.pi * nturns, N0)
# we need to make a special grid in field aligned coordinates, which we do here
# coordinates="raz" tells desc that this grid is in rho,alpha,zeta coordinates.
grid = Grid.create_meshgrid([surfaces, alpha, zeta], coordinates="raz")
data = eq0.compute(
    ["ideal ballooning lambda", "ideal ballooning eigenfunction"], grid=grid
)
print("Growth rate and eigenfunction calculation finished!")
eigenvals = data["ideal ballooning lambda"]
eigenfuns = data["ideal ballooning eigenfunction"]

lambda_max0 = np.zeros(surfaces.size)
eigenfunc_max0 = np.zeros((surfaces.size, N0))
for j in range(surfaces.size):
    idxmax = np.argmax(eigenvals[j])
    alpha_idx, zeta0_idx, eigval_idx = np.unravel_index(idxmax, eigenvals[j].shape)
    # max eigenvalues
    lambda_max0[j] = eigenvals[j, alpha_idx, zeta0_idx, eigval_idx]
    # eigenfunction corresponding to the max eigenvalue
    X0 = eigenfuns[j, alpha_idx, zeta0_idx, :, eigval_idx]
    sign_max = np.sign(X0[np.argmax(np.abs(X0))])
    eigenfunc_max0[j, 1:-1] = X0 / np.max(np.abs(X0)) * sign_max

Growth rate and eigenfunction calculation finished!

Plotting the maximum ballooning growth rate and the corresponding eigenfunction

Now, we plot \(\lambda_{\mathrm{max}}\) on different flux surfaces.

plt.plot(surfaces, lambda_max0, "-or", ms=4)
plt.xlabel(r"$\rho$", fontsize=18)
plt.ylabel(r"$\lambda_{\mathrm{max}}$", fontsize=18)
plt.xticks(fontsize=16)
plt.yticks(fontsize=16)

plt.figure()
plt.plot(zeta, eigenfunc_max0[3])  # plotting eigenfunction on rho=0.4
plt.xlabel(r"$\zeta$", fontsize=18)
plt.ylabel(r"$X_{\mathrm{max}}$", fontsize=18)
plt.xticks(fontsize=16)
plt.yticks(fontsize=16);

Note that the ballooning instability is driven by a pressure gradient. In the edge, if \(p = 0\) but \(dp/d\rho \neq 0\), we can still have a finite ballooning growth rate. As a corollary, for vacuum equilibria or on equilibria with \(dp/d\rho = 0\), the ballooning growth rate \(\lambda \leq 0\)

Aside: Newcomb’s metric, A faster proxy for ballooning stability

To expedite the stability calculation even further, we can use the Newcomb metric instead. It uses the Sturm oscillation theorem to determine whether an equilibrium is stable or unstable. However, the growth rate information or eigenfunction information is not calculated. But since the calculation is faster than an eigenvalue solve, we can determine the stability of an equilibrium.

This idea is explained in much more detail in Appendix D of Gaur et al.

data = eq0.compute(["Newcomb ballooning metric"], grid=grid, data=data)

Plotting Newcomb’s metric

A newcomb metric value greater than zero implies stability whereas any value less than zero implies instability.

plt.plot(surfaces, data["Newcomb ballooning metric"], "-or", ms=4)
plt.xlabel(r"$\rho$", fontsize=18)
plt.ylabel("Newcomb metric", fontsize=18)
plt.xticks(fontsize=16)
plt.yticks(fontsize=16);

Optimizing for ballooning stability

Next, we explain how to take an existing ballooning unstable equilibirum and optimize it for ideal ballooning stability using DESC. The process is similar to other tutorials where you add the ballooning objective to the list of objective functions.

# save a copy of original for comparison
eq1 = eq0.copy()
# Number of point along a field line per transit
nzetaperturn = 200
# Determine which modes to unfix
k = 2

print("\n---------------------------------------")
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(
    M=eq1.M_grid,
    N=eq1.N_grid,
    rho=np.array([1.0]),
    NFP=eq1.NFP,
    sym=eq1.sym,
    axis=False,
)

objective = ObjectiveFunction(
    (
        BallooningStability(
            eq=eq1,
            rho=np.array([0.8]),
            alpha=alpha,
            nturns=nturns,
            nzetaperturn=nzetaperturn,
            weight=2,
        ),
        AspectRatio(
            eq=eq1,
            bounds=(8, 11),
            weight=1,
        ),
        GenericObjective(
            f="curvature_k2_rho",
            thing=eq1,
            grid=Curvature_grid,
            bounds=(-np.inf, 0),
            weight=2,
        ),
    )
)


optimizer = Optimizer("proximal-lsq-exact")
(eq1,), _ = optimizer.optimize(
    eq1,
    objective,
    constraints,
    ftol=1e-4,
    xtol=1e-6,
    gtol=1e-6,
    maxiter=5,  # increase maxiter to 50 for a better result
    verbose=3,
    options={"initial_trust_ratio": 2e-3},
)
print("Optimization complete!")

---------------------------------------
Optimizing boundary modes M, N <= 2
---------------------------------------
Building objective: ideal ballooning lambda
Building objective: aspect ratio
Precomputing transforms
Timer: Precomputing transforms = 28.4 ms
Building objective: Generic
Timer: Objective build = 1.26 sec
Building objective: force
Precomputing transforms
Timer: Precomputing transforms = 767 ms
Timer: Objective build = 1.45 sec
Timer: Objective build = 1.51 ms
Timer: Eq Update LinearConstraintProjection build = 4.39 sec
Timer: Proximal projection build = 20.9 sec
Building objective: lcfs R
Building objective: lcfs Z
Building objective: fixed pressure
Building objective: fixed iota
Building objective: fixed Psi
Timer: Objective build = 578 ms
Timer: LinearConstraintProjection build = 1.67 sec
Number of parameters: 24
Number of objectives: 119
Timer: Initializing the optimization = 23.3 sec

Starting optimization
Using method: proximal-lsq-exact
Solver options:
------------------------------------------------------------
Maximum Function Evaluations       : 26
Maximum Allowed Total Δx Norm      : inf
Scaled Termination                 : True
Trust Region Method                : qr
Initial Trust Radius               : 1.893e-01
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          1.926e+02                                    1.963e+01
       1              2          8.279e+01      1.098e+02      5.610e-02      1.287e+01
       2              3          6.403e+01      1.876e+01      1.759e-02      1.132e+01
       3              4          5.443e+01      9.601e+00      3.516e-02      1.043e+01
       4              5          4.639e+01      8.041e+00      4.194e-02      9.632e+00
       5              6          3.210e+01      1.429e+01      3.839e-02      8.012e+00
Warning: Maximum number of iterations has been exceeded.
         Current function value: 3.210e+01
         Total delta_x: 1.516e-01
         Iterations: 5
         Function evaluations: 6
         Jacobian evaluations: 6
Timer: Solution time = 1.23 min
Timer: Avg time per step = 12.3 sec
==============================================================================================================
                                                                 Start  -->   End
Total (sum of squares):                                      1.822e+02  -->   3.210e+01,
Maximum absolute Ideal ballooning lambda:                    9.544e+00  -->   4.006e+00 ~
Minimum absolute Ideal ballooning lambda:                    9.544e+00  -->   4.006e+00 ~
Average absolute Ideal ballooning lambda:                    9.544e+00  -->   4.006e+00 ~
Maximum absolute Ideal ballooning lambda:                    9.544e+00  -->   4.006e+00 (normalized)
Minimum absolute Ideal ballooning lambda:                    9.544e+00  -->   4.006e+00 (normalized)
Average absolute Ideal ballooning lambda:                    9.544e+00  -->   4.006e+00 (normalized)
Aspect ratio:                                                1.049e+01  -->   1.054e+01 (dimensionless)
Maximum Generic objective value:                            -6.936e-01  -->  -4.708e-01 (m^{-1})
Minimum Generic objective value:                            -5.727e+00  -->  -1.157e+01 (m^{-1})
Average Generic objective value:                            -1.567e+00  -->  -1.600e+00 (m^{-1})
Maximum Generic objective value:                            -6.936e-01  -->  -4.708e-01 (normalized)
Minimum Generic objective value:                            -5.727e+00  -->  -1.157e+01 (normalized)
Average Generic objective value:                            -1.567e+00  -->  -1.600e+00 (normalized)
Maximum absolute Force error:                                2.589e+07  -->   1.590e+05 (N)
Minimum absolute Force error:                                2.997e+00  -->   1.916e+00 (N)
Average absolute Force error:                                1.034e+05  -->   5.980e+03 (N)
Maximum absolute Force error:                                2.078e+00  -->   1.276e-02 (normalized)
Minimum absolute Force error:                                2.406e-07  -->   1.538e-07 (normalized)
Average absolute Force error:                                8.300e-03  -->   4.801e-04 (normalized)
R boundary error:                                            2.506e-19  -->   1.253e-18 (m)
Z boundary error:                                            2.657e-19  -->   1.329e-18 (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!

desc.plotting.plot_comparison([eq0, eq1]);

Comparing growth rates of the initial and optimized equilibria

After finishing an optimization, we compare the maximum ballooning growth rate \(\lambda_{\mathrm{max}}\) between the initial and optimized equilibria.

data = eq1.compute(["ideal ballooning lambda"], grid=grid)
lambda_max1 = data["ideal ballooning lambda"].max(axis=(-1, -2, -3))

plt.plot(surfaces, lambda_max0, "-or", ms=4)
plt.plot(surfaces, lambda_max1, "-og", ms=4)
plt.legend(["initial", "optimized"], fontsize=16)
plt.xlabel(r"$\rho$", fontsize=18)
plt.ylabel(r"$\lambda_{\mathrm{max}}$", fontsize=18)
plt.xticks(fontsize=16)
plt.yticks(fontsize=16);

Relation to \(\delta\! f\) gyrokinetics

The coefficients used by our ballooning solver are the exact same set used by any \(\delta\! f\) gyrokinetic solver. Below we will specify the list and plot seven coefficients used by the GX gyrokinetics solver

\[\mathrm{gradpar} = \mathbf{b}\cdot \mathbf{\nabla}\zeta\]

\[g^{\alpha\alpha} = \lvert \mathbf{\nabla} \alpha \rvert^2\]

\[g^{\rho\alpha} = \mathbf{\nabla} \alpha \cdot \mathbf{\nabla} \rho\]

\[g^{\rho\rho} = \lvert \mathbf{\nabla} \rho \rvert^2\]

\[\mathrm{cvdrift} = \frac{1}{B}[\mathbf{b} \times (\mathbf{b}\cdot \mathbf{\nabla}\mathbf{b})] \cdot \mathbf{\nabla} \alpha = \frac{1}{B^3}[\mathbf{b} \times \nabla \left(\mu_0 p + B^2/2\right)] \cdot \mathbf{\nabla} \alpha\]

\[\mathrm{gbdrift} = \frac{1}{B^2}(\mathbf{b} \times \mathbf{\nabla}B) \cdot \mathbf{\nabla} \alpha\]

\[\mathrm{cvdrift_0} = \frac{2 \rho}{B^2}(\mathbf{b} \times \mathbf{\nabla}B) \cdot \mathbf{\nabla} \rho\]

grid = Grid.create_meshgrid([surfaces[-1], alpha[-1], zeta], coordinates="raz")
data = eq1.compute(
    ["B^zeta", "g^aa", "g^ra", "g^rr", "cvdrift", "cvdrift0", "gbdrift"], grid=grid
)
fig, ax = plt.subplots(7, sharex=True, figsize=(8, 16))

ax[0].plot(zeta, data["B^zeta"] / data["|B|"], "-or", ms=1.5)
ax[1].plot(zeta, data["g^aa"], "-or", ms=1.5)
ax[2].plot(zeta, data["g^ra"], "-or", ms=1.5)
ax[3].plot(zeta, data["g^rr"], "-or", ms=1.5)
ax[4].plot(zeta, data["cvdrift"], "-or", ms=1.5)
ax[5].plot(zeta, data["cvdrift0"], "-or", ms=1.5)
ax[6].plot(zeta, data["gbdrift"], "-or", ms=1.5)

ax[0].set_ylabel("$\\mathrm{gradpar}$")
ax[1].set_ylabel("$g^{\\alpha \\alpha}$")
ax[2].set_ylabel("$g^{\\rho \\alpha}$")
ax[3].set_ylabel("$g^{\\rho \\rho}$")
ax[4].set_ylabel("$\\mathrm{cvdrift}$")
ax[5].set_ylabel("$\\mathrm{cvdrift}_0$")
ax[6].set_ylabel("$\\mathrm{gbdrift}$")
ax[6].set_xlabel("$\\zeta$");