GX: a GPU-native gyrokinetic turbulence code for tokamak and stellarator design

Abstract

GX is a code designed to solve the nonlinear gyrokinetic system for low-frequency turbulence in magnetized plasmas, particularly tokamaks and stellarators. In GX, our primary motivation and target is a fast gyrokinetic solver that can be used for fusion reactor design and optimization along with wide-ranging physics exploration. This has led to several code and algorithm design decisions, specifically chosen to prioritize time to solution. First, we have used a discretization algorithm that is pseudospectral in the entire phase space, including a Laguerre–Hermite pseudospectral formulation of velocity space, which allows for smooth interpolation between coarse gyrofluid-like resolutions and finer conventional gyrokinetic resolutions and efficient evaluation of a model collision operator. Additionally, we have built GX to natively target graphics processors (GPUs), which are among the fastest computational platforms available today. Finally, we have taken advantage of the reactor-relevant limit of small $\rho _*$ by using the radially local flux-tube approach. In this paper we present details about the gyrokinetic system and the numerical algorithms used in GX to solve the system. We then present several numerical benchmarks against established gyrokinetic codes in both tokamak and stellarator magnetic geometries to verify that GX correctly simulates gyrokinetic turbulence in the small $\rho _*$ limit. Moreover, we show that the convergence properties of the Laguerre–Hermite spectral velocity formulation are quite favourable for nonlinear problems of interest. Coupled with GPU acceleration, which we also investigate with scaling studies, this enables GX to be able to produce useful turbulence simulations in minutes on one (or a few) GPUs and higher fidelity results in a few hours using several GPUs. GX is open-source software that is ready for fusion reactor design studies.

1. Introduction

Tokamaks and stellarators are fusion reactor concepts that use magnetic fields to reduce the rate of loss of particles, momentum and energy from the reactor's toroidal confinement volume. In general, the confining magnetic fields are produced both by magnets that surround the toroidal confinement region, and by currents flowing through the confined plasma. To a great extent, the detailed spatial patterns of magnetic fields determine whether the operating conditions of a tokamak or stellarator permit sustained thermonuclear fusion reactions. Consequently, a key step in the design of any tokamak or stellarator reactor is the evaluation of reactor performance as a function of magnetic geometry.

Turbulent transport is one major factor that limits reactor performance. Turbulence in the fuel column is driven by the steep gradients of density, velocity and temperature that exist between the hot, sometimes rapidly flowing, dense fuel core and the relatively cold, diffuse environment. Under optimal conditions, the turbulent fluctuations are small-amplitude perturbations of the steady-state plasma density, temperature, etc. but these fluctuations nonetheless enhance losses, reduce insulation (by rapidly mixing cold plasma from the outer region with hot plasma in the reactor core) and generally limit reactor performance (Abel et al. 2008). Decades of meticulous experiments and direct observations of plasma fluctuations in many tokamaks and stellarators have established a good understanding of when and where the turbulence arises, how intense it is, its dominant wavelengths and frequencies, and so on. The evaluation of the expected performance of any proposed tokamak or stellarator reactor design therefore essentially requires some way to estimate properties of the plasma turbulence that will arise as a result of the details of the proposed magnetic field geometry.

More than two decades of extensive validation campaigns have established that gyrokinetic theory quantitatively describes this turbulence. Many gyrokinetic simulation codes (Parker, Lee & Santoro 1993; Kotschenreuther, Rewoldt & Tang 1995; Dorland et al. 2000; Jenko 2000; Lin et al. 2000; Jost et al. 2001; Candy & Waltz 2003; Idomura, Tokuda & Kishimoto 2003; Watanabe & Sugama 2005; Wang et al. 2006; Jolliet et al. 2007; Idomura et al. 2008; Peeters et al. 2009; Candy, Belli & Bravenec 2016; Barnes, Parra & Landreman 2019; Lanti et al. 2020) have been used to study, test and ultimately validate our understanding of how plasma turbulence depends on the magnetic geometry (see e.g. Belli, Hammett & Dorland (2008), Marinoni et al. (2009), Mynick, Xanthopoulos & Boozer (2009) and White (2019)). Because there are many algorithmic trade-offs to consider when setting up a turbulence simulation code, no single gyrokinetic code is ‘best’ for these purposes. In this paper we present GX, a fast gyrokinetic solver where our primary motivations and targets are fusion reactor design, optimization and wide-ranging physics exploration at reactor scales.

What does this mean in practice? It means primarily that we have prioritized time to solution (together with quantitative uncertainty estimates) over comprehensiveness and full realism, and that we aim to provide simple runtime control parameters to shift smoothly from very fast but relatively less accurate calculations such as are of most value in ‘outer loops’ of a design activity, to more comprehensive, ‘standard’ results that are (much) more expensive to evaluate. Moreover, instead of retrofitting an existing CPU-based code, we started from scratch, with every algorithmic design decision made to target graphics processors (GPUs), which are among the fastest computational platforms available today, with NVIDIA or AMD GPUs providing most of the computer power for seven of the top 10 fastest supercomputers in the world (Strohmaier et al. 2022). Finally, since we are interested in reactors specifically, we focused on the reactor limit of small $\rho _* \equiv \rho _i/R$, where $\rho _i$ is the typical radius of gyration of an ion in the plasma, and $R$ is the major radius of the device. In the small $\rho _*$ limit, there are major simplifications to be had, corresponding physically to the opening of a wide ‘gap’ between typical radial correlation lengths $\lambda _{\rm corr}\sim \rho _i$ of the turbulence and the radial gradient scale lengths $L\sim R$. In contrast, present-day devices are not in the small $\rho _*$ limit, so turbulent eddies can be large enough to sample variations in the driving gradients and other relevant plasma properties. Additionally, because stellarators are non-axisymmetric, they have irreducible geometric variation within a given flux surface that makes the small $\rho _*$ limit more challenging. The turbulent correlation length $\lambda _{\rm corr}$ must also be small compared with the distance that geometric differences within the flux surface change significantly. This limit exists for small enough $\rho _*$, but it is roughly more challenging by a factor of the aspect ratio $R/a$, where $a$ is the minor radius of the torus. To account for the ‘mesoscale’ phenomena that may result and to complete the many quantitative validation studies that have been undertaken, many teams chose algorithms that handle radially non-local physics (often called ‘global’ codes). Codes that take advantage of the simpler, radially local, small $\rho _*$ limit are known as ‘flux-tube’ codes, and our new code GX is in this category. Several flux-tube turbulence calculations can then be coupled together with a transport solver to obtain transport time scale profile evolution on the device scale (Barnes et al. 2009) or steady-state profiles (Candy et al. 2009). This multiscale approach is asymptotically valid in the limit of vanishing $\rho _*$ (Abel et al. 2013). Recent work has also shown a novel approach to include radially global effects within the flux-tube formulation (Parra & Barnes 2015; St-Onge, Barnes & Parra 2022).

To verify that GX correctly simulates gyrokinetic turbulence in the small $\rho _*$ limit, we present several tokamak and stellarator benchmarks of GX with established flux-tube gyrokinetic simulation codes. Note that while numerical verification (benchmarking) is code-specific, the extensive history of experimental validation of the gyrokinetic model is inherited by any verified gyrokinetic code and need not be repeated. GX is open-source scientific software¹ that is therefore ready for immediate design studies. Below, we show that GX can produce useful turbulence simulations in minutes using one (or a few) GPUs and higher fidelity results in a few hours using several GPUs.

Our algorithmic approach to solving the gyrokinetic system is based in pseudospectral methods. There is a long history of the use of Fourier pseudospectral methods in turbulence simulations to discretize the configuration space, especially in flux-tube gyrokinetic codes like GS2 (Kotschenreuther et al. 1995; Dorland et al. 2000), (non-global) GENE (Jenko & Dorland 2001), CGYRO (Candy et al. 2016) and stella (Barnes et al. 2019). These methods have the advantage of spectrally accurate evaluation of derivatives. Building on this, we have developed a pseudospectral formulation for discretizing the velocity space in the gyrokinetic system using a Laguerre–Hermite spectral basis (Mandell, Dorland & Landreman 2018). Projecting the (perturbed) gyrokinetic distribution function onto this basis produces spectral modes that correspond to (gyro)fluid moment quantities (like density, parallel momentum, etc.), so that projection of the gyrokinetic equation corresponds to a gyrofluid system with arbitrarily many moments. In the lowest-resolution limit the model retains critical conservation laws (density, momentum, energy and free energy) and corresponds precisely to the well-established gyrofluid models of Dorland & Hammett (1993), Beer & Hammett (1996) and Snyder & Hammett (2001) when their Landau-fluid closures are employed (Hammett & Perkins 1990; Beer & Hammett 1996). At moderate velocity-space resolution, one can extend the Landau-fluid closure approach to accommodate more moments (Smith 1997) or use ‘hypercollisions’ to close the moment hierarchy. At high velocity-space resolution, the model can retain a similar number of degrees of freedom as typically used in conventional gyrokinetic Eulerian models. A key advantage of our approach is therefore the flexibility to interpolate between these limits, depending on the desired balance of accuracy and speed for a particular calculation. Additionally, our Fourier–Laguerre–Hermite pseudospectral algorithm is a good fit for modern GPU computing: the algorithm relies heavily on fast transform methods that are well-optimized on GPUs, and the memory requirements of spectral algorithms are low enough to fit a problem onto one or a few GPUs. Several gyrokinetic codes have been ported to GPUs in recent years to target modern heterogeneous computing platforms (Madduri et al. 2011; D'Azevedo et al. 2018; Sfiligoi, Candy & Kostuk 2018; Germaschewski et al. 2021; Ohana et al. 2021; Belli et al. 2022).

Similar works on Laguerre–Hermite spectral methods for drift kinetics (Jorge, Ricci & Loureiro 2017) and gyrokinetics (Jorge, Frei & Ricci 2019; Frei, Jorge & Ricci 2020; Frei et al. 2021; Frei, Ernst & Ricci 2022a; Frei, Hoffmann & Ricci 2022b; Hoffmann, Frei & Ricci 2023a,b) have shown promising results and can be viewed as complementary to our work. In particular, extensive linear benchmarks and convergence studies have been performed (Frei et al. 2022b, 2023), showing that the Laguerre–Hermite approach can be used successfully to model ion temperature gradient (ITG), trapped electron, kinetic ballooning and microtearing modes. Jorge et al. (2019), Frei et al. (2021, 2022a,b) develop, implement and study advanced gyrokinetic collision operators by leveraging the Laguerre–Hermite spectral approach. Nonlinear studies in Z-pinch (Hoffmann et al. 2023b) and full toroidal (Hoffmann et al. 2023a) geometry have also now been performed, showing that accurate nonlinear results, including the Dimits shift (Dimits et al. 2000), can be obtained with fewer basis functions than typically necessary for convergence of linear instabilities. Additionally, Frei et al. (2020) developed a Laguerre–Hermite projection of the full-$f$ gyrokinetic system to target the periphery of tokamaks, and Frei, Mencke & Ricci (2024) implemented the full-$f$ system in slab geometry for modelling linear devices.

The remainder of this paper is organized as follows. In § 2 we describe the gyrokinetic model. We describe the Fourier–Laguerre–Hermite pseudospectral formulation of the gyrokinetic system in § 4. Section 3 describes the magnetic geometry and coordinates used for tokamaks and stellarators. The time discretization scheme is described in § 5. In § 6, several benchmarks are presented comparing GX with established flux-tube gyrokinetic codes (GS2, GENE and stella) for linear and nonlinear cases in tokamak and stellarator geometries. Section 7 details the convergence properties of the Laguerre–Hermite basis for nonlinear calculations, along with details about GPU performance and multi-GPU scaling. Finally, we present the conclusions and future work in § 8.

2. Gyrokinetic model equations

In GX we solve the $\delta f$ gyrokinetic system. The literature of gyrokinetic theory is vast; our notation and philosophy follow four particular references: Antonsen & Lane (1980), Frieman & Chen (1982), Barnes et al. (2010) and Abel et al. (2013). The velocity space in gyrokinetic theory is two-dimensional because the particles gyrate very rapidly around their guiding centre position, averaging over any field variations that are encountered in one gyration period so that we can ignore the instantaneous position of the particle around its gyro-orbit. We choose to write the gyrokinetic equation in $(v_\parallel,\mu )$ coordinates, where $v_\parallel$ is the speed in the direction of the magnetic field and $\mu \equiv v_\perp ^2/2B$ is the magnetic moment, where $v_\perp$ is the speed perpendicular to the magnetic field and $B$ is the local magnetic field strength. Using $\mu$ as a coordinate takes advantage of the fact that $\mu$ is a gyrokinetic adiabatic invariant. These coordinates are not optimal for resolving the boundary layers in velocity that occur in linear gyrokinetic theory where particle orbits are ‘trapped’ in magnetic wells on one side of the boundary layer and ‘passing’ (untrapped) on the other. To resolve the sharp boundary layers it is better to use energy and pitch angle as one's velocity space coordinates (as is done in the GS2 code, for example, which optimizes its velocity space grid to enhance resolution near the trapped–passing boundary), but this leads to a considerably more complicated code implementation, particularly for stellarator applications, because of the need to carefully handle turning points in each magnetic well. However, in a sufficiently turbulent plasma, these boundary layers will be broadened, reducing the need for coordinates optimized for a sharp boundary. Thus, we have selected the coordinates that lead to simpler, more efficient code because we are particularly interested in calculating nonlinear turbulence properties (as opposed to linear stability conditions). These coordinates are also used in GENE and stella.

Following (144) of Abel et al. (2013), but expressing the equations in $(v_\parallel,\mu )$ velocity coordinates, neglecting strong equilibrium flows and normalizing all quantities to non-dimensional form (see appendix A), we have

(2.1)\begin{align} & \frac{\partial h_s}{\partial t} + \left[v_{ts} v_\parallel {\hat{\boldsymbol b}} + \langle{ {\boldsymbol v}_{\chi} \rangle}_{\boldsymbol R} + \frac{\tau_s}{Z_s}{\boldsymbol v}_d \right] \boldsymbol{\cdot} \boldsymbol{\nabla} h_s - v_{ts} \mu ({\hat{\boldsymbol b}}\boldsymbol{\cdot}\boldsymbol{\nabla} B) \frac{\partial h_s}{\partial v_\parallel} \nonumber\\ & \quad = \frac{Z_s}{\tau_s}F_{Ms} \frac{\partial \langle{ \chi \rangle}_{\boldsymbol R}}{\partial t} - \langle{ {\boldsymbol v}_{\chi} \rangle}_{\boldsymbol R} \boldsymbol{\cdot} \boldsymbol{\nabla}|_E F_{Ms} + C(h_s). \end{align}

Here, we have split the total distribution function as $F_s = F_{0s} + \delta f_s = F_{Ms}(1-Z_s\varPhi /\tau _s)+h_s$, with $(Z_s/\tau _s) F_{Ms} \varPhi ({\boldsymbol r},t)$ and $h_s({\boldsymbol R},v_\parallel,\mu,t)$ the Boltzmann and non-Boltzmann parts of the $\delta f_s$ perturbation, respectively. The equilibrium distribution function is a Maxwellian in energy $E$ (which we take to have no flows),

(2.2)\begin{equation} F_{0s} = F_{Ms} = \frac{n_s}{(2{\rm \pi} v_{ts}^2)^{3/2}}\, {\rm e}^{{-}E} = \frac{n_s}{(2{\rm \pi} v_{ts}^2)^{3/2}}\, {\rm e}^{{-}v_\parallel^2/2- \mu B}, \end{equation}

and the fluctuating gyroaveraged distribution function satisfies $h_s\ll F_{Ms}$. Note that in (2.1), the gradient of $F_M$ at constant energy is denoted $\boldsymbol {\nabla }|_E F_M$. We also have the following dimensionless species parameters: equilibrium density $n_s$; equilibrium temperature $\tau _s$; mass $m_s$, charge $Z_s$; thermal velocity $v_{ts}=\sqrt {\tau _s/m_s}$.

The distribution function describes the probability of finding a particle of species $s$ with gyrocentre position ${\boldsymbol R}$, velocity parallel to the magnetic field $v_\parallel$ and magnetic moment $\mu =v_\perp ^2/2B$, where $v_\perp$ is the speed in the plane perpendicular to the magnetic field. The equilibrium magnetic field ${\boldsymbol B}$ has magnitude $B=B({\boldsymbol r})$ and direction $\hat {\boldsymbol b} = {\boldsymbol B}/B$, and magnetic fluctuations are given by $\delta {\boldsymbol B} = \delta {\boldsymbol B}_\perp + \delta B_\parallel {\boldsymbol {\hat {b}}}$, where the perpendicular and parallel components can be expressed in terms of the vector potential ${\boldsymbol A} = {\boldsymbol A}_\perp + A_\parallel {\boldsymbol {\hat {b}}}$ via $\delta {\boldsymbol B}_\perp = {[\boldsymbol {\nabla } \times {\boldsymbol A}]_\perp \simeq } \boldsymbol {\nabla } A_\parallel \times {\boldsymbol {\hat {b}}}$ and $\delta B_\parallel = {\boldsymbol {\hat {b}}}\boldsymbol {\cdot }(\boldsymbol {\nabla } \times {\boldsymbol A}_\perp )$, respectively. The gyrokinetic potential $\chi ({\boldsymbol r},t)$ is composed of the electrostatic potential $\varPhi ({\boldsymbol r}, t)$ and the vector potential ${\boldsymbol A}({\boldsymbol r}, t)$,

(2.3)\begin{equation} \chi({\boldsymbol r}, {\boldsymbol v}, t) = \varPhi({\boldsymbol r}, t) - v_{ts} {\boldsymbol v}\boldsymbol{\cdot} {\boldsymbol A}({\boldsymbol r}, t), \end{equation}

and is a function of particle position $\boldsymbol r$, where ${\boldsymbol r}= {\boldsymbol R} + \boldsymbol {\rho }$ so that $\boldsymbol {\rho }$ is the gyroradius vector that rotates at the gyrofrequency $\varOmega$ and points from the gyrocentre (at $\boldsymbol R$) to the particle (at $\boldsymbol r$). In (2.1) gyroangle dependence is eliminated by gyroaveraging the potentials at fixed gyrocentre position ${\boldsymbol R}$, denoted by

(2.4)\begin{equation} \langle{ \chi \rangle}_{\boldsymbol R} = \langle{ \varPhi \rangle}_{\boldsymbol R} - v_{ts}v_\parallel\langle{ A_\parallel \rangle}_{\boldsymbol R} - v_{ts}\langle{ {\boldsymbol v}_\perp\boldsymbol{\cdot}{\boldsymbol A}_\perp \rangle}_{\boldsymbol R}, \end{equation}

along with the corresponding gyroaveraged fluctuating velocity $\langle { {\boldsymbol v}_{\chi } \rangle }_{\boldsymbol R} = {\boldsymbol {\hat {b}}}\times \boldsymbol {\nabla } \langle { \chi \rangle }_{\boldsymbol R}$. For most applications, the plasma $\beta$ is low and the magnetic drift velocity is ${\boldsymbol v}_d ={\hat {\boldsymbol b}}\times (v_\parallel ^2 {\hat {\boldsymbol b}}\boldsymbol {\cdot }\boldsymbol {\nabla }{\hat {\boldsymbol b}} + \mu \boldsymbol {\nabla } B)/B$. GX normally uses this low-$\beta$ approximation, but the full expression for the magnetic drifts can be selected at runtime.

The collision operator is denoted by $C(h)$. Here we will follow Mandell et al. (2018) and take the Dougherty model collision operator (Dougherty 1964). Expressed in Fourier space, the like-species Dougherty collision operator is given by

(2.5)\begin{align} C_{ss}(h_s) & = \nu_{ss} \left\{ \left[ \frac{\partial }{\partial v_\parallel} \left(\frac{\partial }{\partial v_\parallel} + v_\parallel \right) + 2 \frac{\partial }{\partial \mu} \left( \frac{\mu }{B} \frac{\partial }{\partial \mu} + \mu \right) + k_\perp^2\rho_s^2 \right] h_s \right. \nonumber\\ & \quad + \left.\vphantom{\frac{\partial }{\partial v_\parallel}}( \bar{T}_s [ (v_\parallel^2 - 1) + 2 (\mu B - 1) ] J_{0 s} + [ \bar{u}_{{\parallel} s} v_\parallel J_{0 s} + \bar{u}_{{\perp} s} v_\perp J_{1 s} ]) F_{Ms} \right\}, \end{align}

with the field-particle terms given by

(2.6)\begin{gather} \bar{u}_{{\parallel} s} \equiv \int {\rm d}^3{\boldsymbol v} J_{0 s} v_\parallel h_s, \end{gather}

(2.7)\begin{gather}\bar{u}_{{\perp} s} \equiv \int {\rm d}^3{\boldsymbol v} J_{1 s} v_\perp h_s, \end{gather}

(2.8)\begin{gather}\bar{T}_s = \frac{1 }{3} (\bar{T}_{{\parallel} s} + 2 \bar{T}_{{\perp} s}) \equiv \frac{1 }{3} \int {\rm d}^3{\boldsymbol v} [ (v_\parallel^2 - 1) + 2 (\mu B - 1) ] J_{0 s} h_s. \end{gather}

Here we have the Bessel functions $J_{0 s}=J_{0}(\alpha _s)$ and $J_{1 s}=J_{1}(\alpha _s)$ that result from the Fourier transform of gyroaverage operators, with $\alpha _s = \sqrt {2\mu B b_s}$, $b_s = k_\perp ^2 \rho _s^2$ and $\rho _s = m_s v_{ts}/(Z_s B)$ is the normalized gyroradius for species $s$.

The Dougherty collision operator is a good physical model of like-particle collisions, which are important to gyrokinetic dynamics. It captures the physics of the collision operator presented in Abel et al. (2008), except that our collision frequency does not have velocity dependence, and there is no difference between the rates of pitch-angle scattering and slowing down. It is also attractive because its Laguerre–Hermite transform is sparse. More realistic collision operators could be included in our model in the future, such as those in Frei et al. (2021, 2022a), which employs a similar Laguerre–Hermite spectral velocity approach as ours. The Dougherty model operator can also be extended to multiple species (Francisquez et al. 2022).

The electromagnetic potentials are determined by the gyrokinetic equivalent of Maxwell's equations. The electrostatic potential is determined by the quasineutrality equation,

(2.9)\begin{equation} \sum_s \frac{ Z_s^2 n_s}{\tau_s} \varPhi = \sum_s Z_s n_s \int {\rm d}^3{\boldsymbol v}\ \langle{ h_s \rangle}_{\boldsymbol r},\end{equation}

where note that here the gyroaverage on the right-hand side is taken at constant particle position ${\boldsymbol r}$. The parallel magnetic vector potential is determined by the parallel component of Ampère's law,

(2.10)\begin{equation} {-}\boldsymbol{\nabla}_\perp^2 A_\parallel{=}\frac{\beta_\mathrm{ref}}{2}\sum_s Z_s n_s v_{ts} \int {\rm d}^3{\boldsymbol v}\ v_\parallel \langle{ h_s \rangle}_{\boldsymbol r}, \end{equation}

with $\beta _\mathrm {ref} = 8{\rm \pi} n_\mathrm {ref} T_\mathrm {ref}/B_{\mathrm {N}}^2$ the plasma beta of the reference species. The perpendicular component of Ampère's law determines the perpendicular component of the vector potential, but is more conveniently expressed in terms of the parallel magnetic fluctuation, $\delta B_\parallel$,

(2.11)\begin{equation} \boldsymbol{\nabla}_\perp^2 \delta B_\parallel{=}{-}\frac{\beta_\mathrm{ref}}{2 {B}}\boldsymbol{\nabla}_\perp\boldsymbol{\nabla}_\perp : \sum_s n_s \tau_s \int {\rm d}^3{\boldsymbol v}\ \langle {\boldsymbol v}_\perp {\boldsymbol v}_\perp h_s\rangle_r, \end{equation}

where the double dot product is defined as $A : B = \sum _{ij} A_{ij} B_{ij}$.

Finally, note that we can define an auxiliary distribution function

(2.12)\begin{equation} g_s = h_s - \frac{Z_s}{\tau_s}\langle{ \chi \rangle}_{\boldsymbol R}F_{Ms}, \end{equation}

which eliminates the time derivative on the right-hand side of (2.1), resulting in

(2.13)\begin{equation} \frac{\partial g_s}{\partial t} + \left[v_{ts} v_\parallel {\hat{\boldsymbol b}} + \langle{ {\boldsymbol v}_{\chi} \rangle}_{\boldsymbol R} + \frac{\tau_s}{Z_s}{\boldsymbol v}_d \right] \boldsymbol{\cdot} \boldsymbol{\nabla} h_s - v_{ts} \mu ({\hat{\boldsymbol b}}\boldsymbol{\cdot}\boldsymbol{\nabla} B) \frac{\partial h_s}{\partial v_\parallel} ={-} \langle{ {\boldsymbol v}_{\chi} \rangle}_{\boldsymbol R} \boldsymbol{\cdot} \boldsymbol{\nabla}|_E F_{Ms} + C(h_s).\end{equation}

This is the form of the gyrokinetic equation that we will solve numerically.

3. Magnetic geometry and choice of coordinates

The geometric structure of a tokamak's magnetic field is axisymmetric, meaning there are no geometric variations that distinguish points at different toroidal angles. Most (but not all) gyrokinetic studies have focused on the axisymmetric limit. Stellarator magnetic fields are not axisymmetric. We have selected coordinates that are convenient for either case. Our spatial coordinates are aligned with the local background magnetic field so that we are able to efficiently resolve perturbations whose parallel wavelengths $\lambda _\parallel$ are $O(\rho _*^{-1})$ longer than the wavelengths $\lambda _\perp$, perpendicular to the local magnetic field; these field-line-following coordinates reduce the number of grid points required to resolve the turbulence by a factor of $\rho _*$. In the case of stellarators, the lack of axisymmetry implies that even when one can describe the turbulence successfully as radially local, it might happen that the geometric variations in the binormal direction (within the magnetic surface) are not numerically well-separated from $\lambda _\perp$, even for reactor-relevant values of $\rho _*$. We have chosen not to address this possibility; we assume that $\rho _*$ is small enough that the turbulence within each flux tube in a given flux surface is not affected by turbulence occurring in a different flux tube within the same surface.

In this section, we will start with the general form of a divergence-free magnetic field, define key important quantities, various coordinate systems and briefly explain how they are used to obtain the geometry-dependent coefficients in the gyrokinetic model. Please refer to appendix E for specific mathematical details.

3.1. Axisymmetric configurations (tokamaks)

We start from the Clebsch form (D'haeseleer et al. 1991) for the equilibrium magnetic field:

(3.1)\begin{equation} \boldsymbol{B} = \boldsymbol{\nabla} \alpha\times \boldsymbol{\nabla} \psi.\end{equation}

We restrict our attention to solutions whose magnetic field lines lie on closed nested toroidal surfaces, known as flux surfaces. Here, flux surfaces are labelled by $\psi$, a radial-like coordinate. On each surface, the binormal coordinate $\alpha$ is a field-line label such that a line of constant $\alpha$ gives the path of a magnetic field line.

For tokamaks, we choose to define the coordinate $\psi$ to be the enclosed poloidal flux divided by $2{\rm \pi}$,

(3.2)\begin{equation} \psi = \frac{\varPsi_\mathrm{pol}}{2{\rm \pi}} = \frac{1}{(2{\rm \pi})^2}\int_V \, {\rm d} \tau {\boldsymbol B}\boldsymbol{\cdot}\boldsymbol{\nabla}\theta, \end{equation}

and the binormal coordinate $\alpha$ is chosen to be

(3.3)\begin{equation} \alpha = \phi - q(\psi)\theta, \end{equation}

where $\phi$ and $\theta$ are ‘straight-field-line’ toroidal and poloidal angles, respectively, and

(3.4)\begin{equation} q(\psi) = \frac{1}{(2{\rm \pi})^2}\int_{0}^{2{\rm \pi}} {\rm d} \phi \int_{0}^{2{\rm \pi}} {\rm d} \theta\, \frac{\boldsymbol{B}\boldsymbol{\cdot} \boldsymbol{\boldsymbol{\nabla}}\phi}{\boldsymbol{B}\boldsymbol{\cdot}\boldsymbol{\boldsymbol{\nabla}}\theta},\end{equation}

is the safety factor.

3.2. Non-axisymmetric configurations (stellarators)

For stellarators, we find it more convenient to write the magnetic field in Clebsch form as

(3.5)\begin{equation} {\boldsymbol B} = \boldsymbol{\nabla} \psi \times \boldsymbol{\nabla} \alpha, \end{equation}

with the toroidal flux chosen as the radial-like coordinate,

(3.6)\begin{equation} \psi = \frac{\varPsi_\mathrm{tor}}{2{\rm \pi}} = \frac{1}{(2{\rm \pi})^2}\int_V \, {\rm d} \tau {\boldsymbol B}\boldsymbol{\cdot}\boldsymbol{\nabla}\phi, \end{equation}

and the binormal coordinate chosen to be

(3.7)\begin{equation} \alpha = \theta - \iota(\psi) \phi. \end{equation}

Here, $\iota (\psi ) = 1/q(\psi )$ is the rotational transform, and once again $\phi$ and $\theta$ are ‘straight-field-line’ toroidal and poloidal angles, respectively.

3.3. Field-aligned coordinate system

Since GX is a local flux-tube code, it is advantageous to use a field-aligned coordinate system. Therefore, we introduce the field-aligned, flux-tube coordinates $(x, y, z)$ used by Beer, Cowley & Hammett (1995). Here, $x = x(\psi )$ is a radial coordinate, $y= y(\alpha )$ is a binormal coordinate, and $z=z(\theta )$ is a field-line-following coordinate that parametrizes distance along the equilibrium magnetic field using the poloidal angle $\theta$. The coordinates $x$ and $y$ are normalized forms of $\psi$ and $\alpha$ such that

(3.8)\begin{gather} x = \frac{{\rm d} x}{{\rm d}\psi}(\psi-\psi_0), \end{gather}

(3.9)\begin{gather}y = \frac{{\rm d} y}{{\rm d}\alpha}(\alpha-\alpha_0), \end{gather}

where $\psi _0$ and $\alpha _0$ are the values of equilibrium $\psi$ and $\alpha$ at the centre of the flux tube. The coordinates $(x, y, z)$ have units of length. In the local flux-tube, $x$ and $y$ vary on the length scale of the gyroradius whereas $z$ varies on the length scale of the machine size. We require an equispaced coordinate $z$ along the field line, which is needed to utilize the Fourier spectral treatment of the parallel derivative terms in the gyrokinetic equation. The procedure for obtaining an equispaced, equal-arc $z$ coordinate is provided in appendix E.

Upon defining the field-aligned coordinate system, we can calculate the geometric coefficients required to compute the various terms in the gyrokinetic equation, as described briefly in appendix E. For a three-dimensional equilibrium, GX can read the output data generated by the VMEC code (Hirshman & Whitson 1983). For a two-dimensional axisymmetric equilibrium, GX can use a local equilibrium defined using a Miller parametrization (Miller et al. 1998). The user also has the option to work with simpler geometries like a slab or a cylindrical Z-pinch. Besides these native geometry options, GX can also read geometry data generated by the geometry module in the GS2 code; for example, this enables the capability to use experimentally relevant geometry derived from an EFIT equilibrium reconstruction.

4. Fourier–Laguerre–Hermite pseudospectral formulation

Here we present the Fourier–Laguerre–Hermite formulation of the electromagnetic $\delta f$ gyrokinetic system. Following Mandell et al. (2018), and extending to include electromagnetic fluctuations, we solve the gyrokinetic equation by taking a Fourier transform in the perpendicular directions (which we will call $x$ and $y$; the parallel direction will be $z$), a Laguerre transform in $\mu B$, and a Hermite transform in $v_\parallel$. Note that there are no major disadvantages to the choice of $\mu B$ versus $\mu$ as a coordinate, and our choice to use $\mu B$ is natural because the Laguerre weight function becomes $\exp (-\mu B)$ just like in the Maxwellian distribution.

We define the Fourier–Laguerre–Hermite transform of the distribution function $h_{s}$ for species $s$ as

(4.1)\begin{align} \mathcal{L}_\ell \mathcal{H}_m \mathcal{F}_{\boldsymbol k_\perp}h_s & = 2{\rm \pi} \int_0^\infty {\rm d} \mu B \psi_\ell(\mu B)\int_{-\infty}^\infty {\rm d} v_\parallel \phi_m(v_\parallel) \int {{\rm d} x}\, {{\rm d} y}\, {\rm e}^{-{\rm i} {\boldsymbol k_\perp \boldsymbol{\cdot} R}} h_s(x,y,z,v_\parallel, \mu) \nonumber\\ & \equiv H_{\ell,m}^s({\boldsymbol k_\perp},z), \end{align}

(4.2)\begin{equation} H_{\ell,m}^s({\boldsymbol k_\perp},z) = 2{\rm \pi} \int_0^\infty {\rm d} \mu B \psi_\ell(\mu B)\int_{-\infty}^\infty {\rm d} v_\parallel \phi_m(v_\parallel) \int {{\rm d} x}\, {{\rm d} y}\, {\rm e}^{-{\rm i} {\boldsymbol k_\perp \boldsymbol{\cdot} R}} h_s(x,y,z,v_\parallel, \mu), \end{equation}

where ${\boldsymbol k_\perp } =k_x\boldsymbol {\nabla } x + k_y \boldsymbol {\nabla } y$ is the perpendicular wavenumber vector, $\psi _\ell (\mu B) = (-1)^\ell {\rm L}_{\ell }(\mu B)$ with

(4.3)\begin{equation} {\mathrm{L}_{\ell}(x)=\frac{{\rm e}^x}{\ell!}\frac{\mathrm{d}^\ell}{\mathrm{d} x^\ell} x^\ell \, {\rm e}^{{-}x}}, \end{equation}

the Laguerre polynomials, and $\phi _m (v_\parallel ) = {{\rm He}_{m}({v}_\parallel )}/{\sqrt {m!}}$ with

(4.4)\begin{equation} {\mathrm{He}_m(x) = ({-}1)^m\, {\rm e}^{x^2/2}\frac{\mathrm{d}^m}{\mathrm{d} x^m}\, {\rm e}^{{-}x^2/2}}, \end{equation}

the probabilists’ Hermite polynomials. In Fourier space, gyroaveraging operations (whether at constant $\boldsymbol R$ or constant $\boldsymbol r$) simply become multiplications by Bessel functions, so that the Fourier transform of the gyroaveraged potential is

(4.5)\begin{equation} \mathcal{F}_{\boldsymbol k_\perp} \langle{ \chi \rangle}_{\boldsymbol R} = J_{0s}\varPhi({\boldsymbol k_\perp},z) -v_{ts}{v_\parallel}J_{0s} A_\parallel({\boldsymbol k_\perp},z) + \frac{\tau_s}{Z_s}2\mu B \frac{J_{1s}}{\alpha_s} {\frac{\delta B_\parallel({\boldsymbol k_\perp},z)}{B}}. \end{equation}

Weighting the gyroaveraged potential by a Maxwellian and Laguerre–Hermite transforming results in

(4.6)\begin{equation} \mathcal{L}_\ell\mathcal{H}_m \mathcal{F}_{\boldsymbol k_\perp} \langle{ \chi \rangle}_{\boldsymbol R}F_{Ms} = \mathcal{J}_\ell^s \varPhi \delta_{m 0} - v_{ts}\mathcal{J}_\ell^s A_\parallel \delta_{m1} + \frac{\tau_s}{Z_s}(\mathcal{J}_\ell^s + \mathcal{J}_{\ell-1}^s){\frac{\delta B_\parallel}{B}} \delta_{m0}, \end{equation}

with

(4.7)\begin{equation} \mathcal{J}_\ell^s \equiv \int_0^\infty {\rm d}\mu B \psi_\ell J_{0s}(\sqrt{2\mu B b_s})\,{\rm e}^{-\mu B} = \frac{1}{\ell!} \left(-\frac{b_s}{2}\right)^\ell \,{\rm e}^{{-}b_s/2} \quad (\ell\geq0), \end{equation}

and $\mathcal {J}_\ell ^s\equiv 0$ when $\ell <0$ or when $\ell \geq N_\ell$, with $N_\ell$ the number of evolved Laguerre moments. The Fourier–Laguerre–Hermite transform of the auxiliary distribution function $g$ is then

(4.8)\begin{equation} \mathcal{L}_\ell \mathcal{H}_m \mathcal{F}_{\boldsymbol k_\perp}g_s \equiv G_{\ell,m}^s({\boldsymbol k_\perp},z) = H_{\ell,m}^s - \frac{Z_s}{\tau_s} \mathcal{J}_\ell^s \varPhi \delta_{m 0} + \frac{Z_s v_{ts}}{\tau_s}\mathcal{J}_\ell^s A_\parallel \delta_{m1} - (\mathcal{J}_\ell^s + \mathcal{J}_{\ell-1}^s){\frac{\delta B_\parallel}{B}} \delta_{m0}. \end{equation}

The Fourier–Laguerre–Hermite transform of (2.13) then gives

(4.9)\begin{align} & \frac{\partial G_{\ell,m}^s}{\partial t} + \mathcal{N}_{\ell,m}^s + v_{ts} \boldsymbol{\nabla}_\parallel ( \sqrt{m+1} H^s_{\ell,m+1} + \sqrt{m} H^s_{\ell,m-1} ) \nonumber\\ & \qquad +v_{ts} [ - (\ell+1) \sqrt{m+1} H^s_{\ell,m+1} - \ell \sqrt{m+1} H^s_{\ell-1,m+1} \nonumber\\ & \qquad+ \ell \sqrt{m} H^s_{\ell,m-1} + (\ell+1) \sqrt{m} H^s_{\ell+1,m-1} ] \boldsymbol{\nabla}_\parallel \ln B \nonumber\\ & \qquad + {\rm i}\frac{\tau_s}{Z_s} \omega_{d}^\kappa [ \sqrt{(m+1)(m+2)} H^s_{\ell,m+2} + (2m+1)H^s_{\ell,m} +\sqrt{m (m-1)} H^s_{\ell,m-2} ] \nonumber\\ & \qquad +{\rm i}\frac{\tau_s}{Z_s} \omega_{d}^{\boldsymbol{\nabla} B} [ (\ell+1) H^s_{\ell+1,m} + (2\ell+1)H^s_{\ell,m} +\ell H^s_{\ell-1,m} ] \nonumber\\ & \quad = \mathcal{D}_{\ell,m}^s + \mathcal{C}_{\ell,m}^{ss}. \end{align}

Parallel convection, including bounce motion induced by magnetic trapping in the equilibrium magnetic field, is described by the terms proportional to $\boldsymbol {\nabla }_\parallel \equiv {\hat {\boldsymbol b}} \boldsymbol {\cdot } \boldsymbol {\nabla } = ({\hat {\boldsymbol b}}\boldsymbol {\cdot }\boldsymbol {\nabla } z) \partial /\partial z$. The $\boldsymbol {\nabla }_\parallel H{^s_{\ell,m}}$ terms are evaluated spectrally by Fourier transforming in $z$,

(4.10)\begin{equation} \boldsymbol{\nabla}_\parallel H{^s_{\ell,m}} \equiv {({\hat{\boldsymbol b}}\boldsymbol{\cdot}\boldsymbol{\nabla} z)}\mathcal{F}_{k_{z}}^{{-}1}[ {\rm i} k_{z} \mathcal{F}_{k_{z}} H{^s_{\ell,m}} ],\end{equation}

with

(4.11)\begin{equation} \mathcal{F}_{k_{z}}H{^s_{\ell,m}} = \int {\rm d} z\, {\rm e}^{-{\rm i} k_{z} z} H{^s_{\ell,m}}(z). \end{equation}

Note that a spectral evaluation of this term requires an equispaced grid in $z$, and we also require the $z$ coordinate to be an equal-arclength coordinate so that $({\hat {\boldsymbol b}}\boldsymbol {\cdot }\boldsymbol {\nabla } z)$ is a constant (see § E.1) to avoid a convolution in (4.10). Additional details about this operation related to boundary conditions are given in § 4.2. Toroidicity gives rise to the terms proportional to the curvature drift operator ${\rm i} \omega _d^{\kappa } = (1/B){\hat {\boldsymbol b}}\times ({\hat {\boldsymbol b}}\boldsymbol {\cdot }\boldsymbol {\nabla }{\hat {\boldsymbol b}})\boldsymbol {\cdot } {{\rm i}\boldsymbol k_\perp }$ and the $\boldsymbol {\nabla } B$ drift operator ${\rm i} \omega _d^{\boldsymbol {\nabla } B} = (1/B^2){\hat {\boldsymbol b}}\times \boldsymbol {\nabla } B\boldsymbol {\cdot } {{\rm i}\boldsymbol k_\perp }$ is the $\boldsymbol {\nabla } B$. Drive terms from equilibrium gradients, denoted by $\mathcal {D}_{\ell,m}$, are given by

(4.12)\begin{align} & \mathcal{D}_{\ell,m=0}^s = {\rm i} \omega_* \left[\frac{1}{L_{ns}} \mathcal{J}_{\ell} ^s + \frac{1}{L_{Ts}} [ \ell \mathcal{J}_{\ell-1} ^s + 2\ell \mathcal{J}_{\ell} ^s + (\ell+1) \mathcal{J}_{\ell+1} ^s ] \right]\varPhi \nonumber\\ & \quad +{ \frac{\tau_s}{Z_s}}{\rm i} \omega_* \, \left[\frac{1}{L_{ns}} [\mathcal{J}_{\ell} ^s+\mathcal{J}_{\ell-1} ^s] + \frac{1}{L_{Ts}} [ \ell \mathcal{J}_{\ell-2} ^s \!+\! 3\ell\mathcal{J}_{\ell-1} ^s + (3\ell+1) \mathcal{J}_{\ell} ^s +\! (\ell+1) \mathcal{J}_{\ell+1} ^s ] \right]{\frac{\delta \!B_\parallel}{B}}, \nonumber\\ & \mathcal{D}^s_{\ell, m=1} ={-}v_{ts}{\rm i}\omega_*\left[\frac{1}{L_{ns}}\mathcal{J}_\ell^s + \frac{1}{L_{Ts}}[\ell \mathcal{J}_{\ell-1}^s+(2\ell+1)\mathcal{J}^s_\ell + (\ell+1)\mathcal{J}_{\ell+1}^s] \right]A_\parallel, \nonumber\\ & \mathcal{D}_{\ell,m=2}^s = \frac{1}{\sqrt{2}} {\rm i}\omega_* \frac{1}{L_{Ts}} \mathcal{J}_{\ell} ^s \varPhi + { \frac{\tau_s}{Z_s}}\frac{1}{\sqrt{2}} {\rm i}\omega_* \frac{1}{L_{Ts}} [\mathcal{J}_{\ell} ^s+\mathcal{J}_{\ell-1} ^s] {\frac{\delta B_\parallel}{B}}, \nonumber\\ & \mathcal{D}_{\ell, m=3}^s ={-}v_{ts} \sqrt{\frac{3}{2}}{\rm i}\omega_* \frac{1}{L_{Ts}}\mathcal{J}_\ell^s A_\parallel , \nonumber\\ & \mathcal{D}^s_{\ell, m>3} = 0 , \end{align}

where ${\rm i} \omega _*\equiv -\boldsymbol {\nabla } x \boldsymbol {\cdot } {\hat {\boldsymbol b}}\times {\rm i}{\boldsymbol k_\perp }={\rm i} k_y$, and $L_{ns}$ and $L_{Ts}$ are the normalized density and temperature gradient scale lengths, respectively.

The like-species collision terms are given by

(4.13)\begin{align} \mathcal{C}^{ss}_{\ell,0} & ={-} \nu_{ss} (b_s+2\ell + 0) H^s_{\ell,0} \nonumber\\ & \quad +\nu_{ss} (\sqrt{b_s} \, (\mathcal{J}_\ell^s + \mathcal{J}_{\ell-1}^s) \bar{u}_{{\perp} s} + 2 [ \ell \mathcal{J}_{\ell-1}^s + 2 \ell \mathcal{J}_\ell^s + (\ell+1)\mathcal{J}_{\ell+1}^s] \bar{T}_s ) , \nonumber\\ \mathcal{C}^{ss}_{\ell,1} & ={-} \nu_{ss} (b_s + 2 \ell + 1) H^s_{\ell,1} + \nu_{ss} \mathcal{J}_{\ell} ^s \bar{u}_{{\parallel} s}, \nonumber\\ \mathcal{C}^{ss}_{\ell,2} & ={-} \nu_{ss} (b + 2\ell + 2) H^s_{\ell,2} + \nu_{ss} \sqrt{2} \mathcal{J}_\ell^s \bar{T}_s, \nonumber\\ \mathcal{C}^{ss}_{\ell,m}& ={-} \nu_{ss} (b_s + 2 \ell +m) H^s_{\ell,m}, \quad (m>2), \end{align}

with the field-particle terms given by

(4.14)\begin{gather} \bar{u}_{{\parallel} s} = \int {\rm d}^3{\boldsymbol v} J_{0s} v_\parallel h_s = \sum_{\ell=0}^{N_\ell} \mathcal{J}_{\ell} ^s H_{\ell,1}^s, \end{gather}

(4.15)\begin{gather}\bar{u}_{{\perp} s} = \int {\rm d}^3{\boldsymbol v} J_{1s} v_\perp h_s = \sqrt{b_s} \sum_{\ell=0}^{N_\ell}(\mathcal{J}_\ell^s+\mathcal{J}_{\ell-1}^s)H^s_{\ell,0}, \end{gather}

(4.16)\begin{gather}\bar{T}_{{\parallel} s} = \int {\rm d}^3{\boldsymbol v} J_{0s} (v_\parallel^2-1) h_s = \sqrt{2} \sum_{\ell=0}^{N_\ell} \mathcal{J}_{\ell} ^s H^s_{\ell,2}, \end{gather}

(4.17)\begin{gather}\bar{T}_{{\perp} s} = \int {\rm d}^3{\boldsymbol v} J_{0s} (\mu B-1) h_s = \sum_{\ell=0}^{N_\ell} [\ell \mathcal{J}^s_{\ell-1}+2\ell\mathcal{J}^s_{\ell} +(\ell+1)\mathcal{J}^s_{\ell+1} ]H_{\ell,0}^s, \end{gather}

(4.18)\begin{gather}{\bar{T}_s = \frac{1}{3}(\bar{T}_\parallel{+} 2 \bar{T}_\perp)}. \end{gather}

The nonlinear terms are evaluated pseudospectrally in $(x,y,z,\mu B, m)$ space to avoid convolutions in both Fourier and Laguerre coefficients as (Mandell et al. 2018)

(4.19)\begin{equation} \mathcal{N}_{\ell,m}^s = \mathcal{L}_\ell \mathcal{H}_m \mathcal{F}_{\boldsymbol k_\perp} [\langle{ {\boldsymbol v}_{\chi} \rangle}_{\boldsymbol R}\boldsymbol{\cdot}\boldsymbol{\nabla} {h_s}], \end{equation}

where the complete pseudospectral expression for this term is given in appendix D.

Taking the Fourier–Laguerre–Hermite transform of the field equations, the quasineutrality equation, (2.9), becomes

(4.20)\begin{equation} \sum_s \frac{Z_s^2 n_s}{\tau_s}\varPhi = \sum_s Z_s n_s \int {\rm d}^3{\boldsymbol v} J_{0s} h_s = \sum_s Z_s n_s \sum_{\ell=0}^{N_\ell} \mathcal{J}_\ell^s H_{\ell,0}^s, \end{equation}

the parallel component of Ampère's law, (2.10), becomes

(4.21)\begin{equation} k_\perp^2 A_\parallel{=} \frac{\beta_\mathrm{ref}}{2}\sum_s Z_s n_s v_{ts} \int {\rm d}^3{\boldsymbol v} v_\parallel J_{0s} h_s= \frac{\beta_\mathrm{ref}}{2}\sum_s Z_s n_s v_{ts} \sum_{\ell=0}^{N_\ell} \mathcal{J}_\ell^s H^s_{\ell, 1} \end{equation}

and the perpendicular component of Ampère's law, (2.11), becomes

(4.22)\begin{equation} {\frac{\delta B_\parallel}{B}}={-}\frac{\beta_\mathrm{ref}}{2{B^2}}\sum_s n_s \tau_s \int {\rm d}^3{\boldsymbol v} 2\mu B \frac{J_{1s}}{\alpha_s} h_s ={-}\frac{\beta_\mathrm{ref}}{2{B^2}}\sum_s n_s\tau_s \sum_{\ell=0}^{N_\ell} (\mathcal{J}^s_\ell+\mathcal{J}^s_{\ell-1}) H^s_{\ell, 0}. \end{equation}

To solve these field equations numerically, it is more convenient to express them in terms of $G$ rather than $H$, which results in

(4.23)\begin{gather} \sum_s \frac{ Z_s^2n_s}{\tau_s}\left(1-\sum_{\ell=0}^{N_\ell}(\mathcal{J}^s_\ell)^2\right)\varPhi - \sum_s Z_s n_s \sum_{\ell=0}^{N_\ell}\mathcal{J}^s_\ell(\mathcal{J}^s_\ell+\mathcal{J}^s_{\ell-1}) {\frac{\delta B_\parallel}{B}} = \sum_s Z_s n_s \sum_{\ell=0}^{N_\ell} \mathcal{J}^s_\ell G^s_{\ell,0} , \end{gather}

(4.24)\begin{gather}\left(k_\perp^2 + \frac{\beta_\mathrm{ref}}{2}\sum_s \frac{Z_s^2 n_s}{m_s} \sum_{\ell=0}^{N_\ell} (\mathcal{J}^s_\ell)^2\right)A_\parallel{=} \frac{\beta_\mathrm{ref}}{2}\sum_s Z_s n_s v_{ts} \sum_{\ell=0}^{N_\ell} \mathcal{J}^s_\ell G^s_{\ell,1}, \end{gather}

(4.25)\begin{align} & \frac{\beta_\mathrm{ref}}{2{B^2}}\sum_s Z_s n_s \sum_{\ell=0}^{N_\ell}\mathcal{J}^s_\ell(\mathcal{J}^s_\ell+ \mathcal{J}^s_{\ell-1})\varPhi+ \left(1 + \frac{\beta_\mathrm{ref}}{2{B^2}}\sum_s n_s \tau_s \sum_{\ell=0}^{N_\ell}(\mathcal{J}^s_\ell + \mathcal{J}^s_{\ell-1})^2\right){\frac{\delta B_\parallel}{B}} \nonumber\\ & \quad={-}\frac{\beta_\mathrm{ref}}{2{B^2}}\sum_s n_s \tau_s \sum_{\ell=0}^{N_\ell}(\mathcal{J}^s_\ell + \mathcal{J}^s_{\ell-1})G^s_{\ell,0}. \end{align}

Note that while the sum $\sum (\mathcal {J}_\ell ^s)^2$ could be computed analytically in the $N_\ell \rightarrow \infty$ limit as $\varGamma _0(b) = I_0(b) \,{\rm e}^{-b}$, with $I_0$ the modified Bessel function of the first kind, doing so breaks the energetic consistency of the equations (Mandell et al. 2018); thus these terms should be evaluated as truncated sums.

4.1. Hyperdissipation and closure

Grid-scale dissipation terms are used to ensure numerical stability and robustness. For dissipation in configuration space, we employ a hyperviscosity term of the standard form

(4.26)\begin{equation} \left(\frac{\partial G_{\ell,m}}{\partial t}\right)_\textrm{hyperviscosity} ={-} D \left(\frac{k_\perp^2}{k_{{\perp},\mathrm{max}}^2}\right)^n G_{\ell,m}, \end{equation}

where $k_{\perp,\mathrm {max}}$ is the largest dealiased perpendicular wavenumber in the simulation, and typical values of the variable parameters for nonlinear simulations are $D=0.05$ and $n=2$. This provides a sink at the shortest wavelengths in the domain, which can be useful in nonlinear simulations to avoid spectral pileup. Unlike in other Eulerian gyrokinetic codes which include dissipation along the parallel direction via upwinding, (Kotschenreuther et al. 1995; Jenko & Dorland 2001; Candy et al. 2016; Barnes et al. 2019), the spectral discretization scheme used for the parallel direction in GX does not itself introduce dissipation in the parallel direction. The hypercollision operator and the parallel boundary conditions discussed below introduce some parallel dissipation. A standard parallel hyperdissipation operator of the form $k_\parallel ^{2n}$ could also be included, but we find that this is not necessary for the linear and nonlinear cases that we consider.

Fine-scale structure in velocity space must also be regulated. The Dougherty collision operator fulfils this purpose by acting increasingly strongly on higher Laguerre and Hermite moments, limiting their amplitude. Thus, for a given collisionality, there is a physical cutoff at some ${\ell _\mathrm {cut}}$ and ${m_\mathrm {cut}}$ beyond which fine scales in velocity space are completely wiped out by collisions. At this point, closure of the Laguerre–Hermite moment series by simple truncation (setting $G_{\ell,m}=0$ for $\ell \geq {\ell _\mathrm {cut}}$ or $m\geq {m_\mathrm {cut}}$) is justified. This is the simplest high-resolution closure, but not the only option. One can also obtain an asymptotically correct collisional closure by assuming the collision term becomes dominant in the unresolved moment equations (Zocco, Helander & Connor 2015; Loureiro et al. 2016; Jorge et al. 2017; Frei et al. 2022b).

In the limit of low collisionality or lower Laguerre–Hermite resolution, however, the closure situation is more complicated. Unresolved moments are not expected to be negligible at collisionalities of interest, so closure by truncation will generally give poor results. One possible approach is to follow the collisionless closure approach pioneered by Hammett & Perkins (1990), which was used and extended in the gyrofluid models of Dorland & Hammett (1993), Beer & Hammett (1996), Snyder & Hammett (2001) and Smith (1997) to model parallel, toroidal and nonlinear finite-Larmor-radius (FLR) phase mixing. However, the development of generalized closures in toroidal geometry for a system with an arbitrary number of moments is complicated and beyond the scope of the current work. Instead, we have found that employing a hypercollision operator that provides a sink at large $m$ allows well-behaved results even with simple closure by truncation at relatively low resolution. The operator takes the generic form

(4.27)\begin{equation} \left(\frac{\partial G_{\ell,m}}{\partial t}\right)_\textrm{hyp-coll} ={-}\nu_\mathrm{hyp} m^{p} G_{\ell, m} \quad (m>2). \end{equation}

We ensure that the operator conserves density, momentum and energy by only operating on moments with $m>2$. Similar hypercollision operators (in Hermite space only) have been used in other contexts (Parker 2015; Parker & Dellar 2015; Loureiro et al. 2016), and Parker (2015) also develops theory of a ‘hypercollisional plateau’ where the behaviour of the operator is insensitive to the choice of the variable parameters. In our operator, the coefficient $\nu _\mathrm {hyp}$ is chosen to approximately remove the same amount of fluctuation energy from the largest $m$s as would be removed by parallel phase mixing. The details are provided in § B, and the result is

(4.28)\begin{equation} \nu_\mathrm{hyp} = 2.5 f_\mathrm{hyp} \frac{p+1/2}{m_\mathrm{max}^{p+1/2}}|k_\parallel|v_{ts}, \end{equation}

where $f_\mathrm {hyp}$ is an adjustable coefficient that is typically set to unity. The $|k_\parallel |$ operator is evaluated spectrally by Fourier transforming in $z$. The exponent $p$ is chosen to be $p=N_m/2$, which ensures that only the highest several Hermite modes are strongly damped, roughly independent of $N_m$. While one could consider a similar hyperdissipation mechanism for large $\ell$, we find that Laguerre hypercollisions are unnecessary even when using relatively low Laguerre resolution ($N_\ell \sim 4$) for problems of interest because phase mixing in $v_\parallel$ dominates.

4.2. Boundary conditions

In the flux tube approach we assume statistical periodicity in the perpendicular $(x,y)$ plane. In the parallel direction, we use the ‘twist-and-shift’ boundary condition (Beer et al. 1995), which in the spectral representation can be accomplished by linking modes with different $k_x$ values into an extended $z$ domain, such that

(4.29)\begin{equation} {H_{\ell,m}(k_x, k_y, z) = H_{\ell,m}(k_x+\delta k, k_y, z+2{\rm \pi})}, \end{equation}

with $\delta k = 2{\rm \pi} k_y \hat {s}$. We can then perform a Fourier transform on each extended $z$ domain to evaluate the parallel streaming terms as given by (4.10).

A Fourier representation naturally enforces a periodic boundary condition on $h$ at the ends of each extended domain, but this is not the physically correct boundary condition for non-zonal modes, which should have a zero incoming boundary condition in the infinite ballooning limit. Following Kotschenreuther et al. (1995), we would like to approximate this by enforcing a zero incoming boundary condition on $h$ at the ends of the extended domain, but this boundary condition cannot be enforced directly in Hermite space. Instead, we use a wave-absorbing layer (see e.g. Durran (2010), and references within) near the ends of each extended $z$ domain for the non-zonal components of $h$ to damp out perturbations leaving and re-entering the domain due to the natural periodicity of the Fourier representation. Following Beer (1995), we use a damping filter that is zero in most of the domain and ramps up smoothly near the ends of the extended domain as $d(\hat {z}) = A[1-2\hat {z}^2/(1+\hat {z}^4)]$, where $\hat {z} = (z-z_\mathrm {end})/z_\mathrm {width}$ is a normalized distance from the ends $z_\mathrm {end}$. The width $z_\mathrm {width}$ and the scaling factor $A$ are adjustable parameters; in appendix C we have performed a convergence study and found $z_\mathrm {width} = L_z/8$ and $A = 0.1/\Delta t$ to work well, where $L_z$ the total length of the extended domain. The wave-absorbing damping operation can then be expressed as

(4.30)\begin{equation} \left(\frac{\partial G_{\ell, m}}{\partial t}\right)_\mathrm{ends} ={-}d(\hat{z}) H_{\ell,m} \quad (|\hat{z}|<1,\ k_y \neq 0). \end{equation}

The zonal modes $(k_y=0)$ are not filtered because these modes should be physically periodic, so the natural periodic boundary condition from the Fourier representation is the correct one. We have found this wave-absorbing layer to be particularly necessary for simulations with kinetic electrons.

The standard twist-and-shift boundary condition can be ill-suited for low magnetic shear regions, particularly in stellarators. Martin et al. (2018) have developed a generalization of the twist-and-shift boundary condition that alleviates some of these issues by using the integrated local shear (rather than global shear) in the ‘twist’. This generalized twist-and-shift boundary condition can be used in GX for stellarator geometries with stellarator symmetry. Investigation of a non-twisting flux-tube boundary condition (Ball & Brunner 2021) using GX is also in progress.

Martin et al. (2018) prescribe methods to choose the parallel domain length to enforce exact periodicity or continuous magnetic drifts at the ends of the flux tube. In addition to these options, we adopt an additional criterion for choosing the parallel domain length: to enforce a desired aspect ratio for the perpendicular computational domain. When using the standard twist-and-shift boundary condition, the aspect ratio is quantized as $L_x/L_y = J/(2{\rm \pi} |\hat {s}|)$, with $J$ a non-zero integer. This results in $L_x \propto L_y/\hat {s}$, which makes resolution requirements challenging for $\hat {s}\ll 1$. In the generalized twist-and-shift boundary condition of Martin et al. (2018), the perpendicular aspect ratio of the flux tube is given by

(4.31)\begin{equation} \frac{L_x}{L_y} = J\left[\frac{|\boldsymbol{\nabla} x|^2}{2|\boldsymbol{\nabla} x\boldsymbol{\cdot} \boldsymbol{\nabla} y|}\right]_{z_\mathrm{end}}, \end{equation}

where $z_\mathrm {end}$ is the end of the flux tube and $J$ is again a non-zero integer. In a stellarator, this quantity varies as a function of the parallel domain length (see figure 7 of Martin et al. (2018)). The generalized twist-and-shift boundary condition thus provides the freedom to choose $z_\mathrm {end}$ to obtain the desired aspect ratio (for some integer $J$). Using this prescription to enforce an aspect ratio of order unity can then alleviate the resolution requirements for geometries with small magnetic shear.

5. Timestepping schemes

GX uses explicit time integration methods to advance the system. Along with several standard Runge–Kutta (RK) and strong stability-preserving RK (SSP RK) schemes (Gottlieb, Shu & Tadmor 2001), we also provide an option for the SSP 10-stage fourth-order method of Ketcheson (2008) (a low-storage method ideal for the memory constraints of GPU computing). Additionally, the sparse banded structure of the linear system of (4.9) and the fact that the field equations only involve the lowest-order moments (instead of requiring integration over all velocity space) could allow the development of efficient implicit–explicit (IMEX) schemes for electron dynamics that avoid the timestep restriction of the fast electron motion; this will be a topic for future work. We have used the standard RK3 (third-order) scheme for all of the benchmark calculations presented in § 6.

In all of our timestepping schemes the timestep size $\Delta t$ is constrained by the stability region of the scheme. While in principle the linear eigenvalues of the system could be computed exactly and used to constrain the timestep to ensure that all eigenvalues lie within the stability region, we instead estimate the maximum (linear and nonlinear) frequencies on the grid in each of the $(x,y,z)$ directions via

(5.1)\begin{gather} \omega_{x,\mathrm{max}} \approx \max\left(k_{x,\mathrm{max}} \left(\frac{\tau_s}{Z_s}{\boldsymbol v}_d\boldsymbol{\cdot}\boldsymbol{\nabla} x\right)_\mathrm{max},\ k_{x,\mathrm{max}} ({\boldsymbol v}_E\boldsymbol{\cdot} \boldsymbol{\nabla} x)_\mathrm{max}\right), \end{gather}

(5.2)\begin{gather}\omega_{y,\mathrm{max}} \approx \max\left(k_{y,\mathrm{max}} \left(\frac{\tau_s}{Z_s}{\boldsymbol v}_d\boldsymbol{\cdot}\boldsymbol{\nabla} y\right)_\mathrm{max},\ k_{y,\mathrm{max}} ({\boldsymbol v}_E\boldsymbol{\cdot} \boldsymbol{\nabla} y)_\mathrm{max}\right), \end{gather}

(5.3)\begin{gather}\omega_{z,\mathrm{max}} \approx \max (({\hat{\boldsymbol b}}\boldsymbol{\cdot}\boldsymbol{\nabla} z)k_{z,\mathrm{max}} v_{{\parallel},\mathrm{max}} v_{t,\mathrm{max}}, \omega_{A,\mathrm{max}}), \end{gather}

with

(5.4)\begin{equation} \omega_{A,\mathrm{max}} = \frac{({\hat{\boldsymbol b}}\boldsymbol{\cdot}\boldsymbol{\nabla} z)k_{z,\mathrm{max}} v_{\rm te}}{\sqrt{\dfrac{\beta_\mathrm{ref}n_e \tau_e}{2}\dfrac{m_i}{m_e} + k_{{\perp},\mathrm{min}}^2 \rho_s^2}}, \end{equation}

the maximum shear Alfvén frequency on the grid. In the electrostatic limit ($\beta _\mathrm {ref} = 0$) this reduces to the high-frequency $\omega _H$ mode (Lee 1987). Here, $k_{z,\mathrm {max}}$ is the largest parallel wavenumber in the simulation and $k_{\perp,\mathrm {min}}$ is the smallest perpendicular wavenumber. To evaluate $v_{\parallel,\mathrm {max}}$ and $(\mu B)_\mathrm {max}$ we use the maximum value of the Gauss–Hermite and Gauss–Laguerre collocation grids, respectively, which depend on the number of spectral modes used in the calculation, $N_m$ and $N_\ell$.

The constraint on the maximum stable timestep is then approximately given by

(5.5)\begin{equation} (\omega_{x,\mathrm{max}} + \omega_{y,\mathrm{max}} + \omega_{z,\mathrm{max}})\Delta t \lesssim C s, \end{equation}

where $s$ is dependent on the stability region for the particular timestepping scheme (for example, $s=1.73$ for RK3, and $s=2.82$ for RK4), and $0< C\leq 1$ is a user-defined scaling factor. Formally, the stability constraint with $C=1$ is only valid for linear modes that do not grow or decay; thus we typically use $C\sim 0.5$–0.9 to ensure stability in the full nonlinear system with a spectrum of growing and damped modes.

6. Numerical benchmarks

In this section we show a variety of linear and nonlinear benchmarks of GX by comparing results with widely benchmarked gyrokinetic codes (GS2, stella, GENE) in both tokamak and stellarator geometry. In all cases, both linear and nonlinear, an initial value problem is solved. The numerical resolution and associated parameters used in each case are given in appendix G.

6.1. Tokamak benchmarks: Cyclone base case

For benchmarks in tokamak geometry, we choose a configuration based on the ‘Cyclone base case’ (CBC), a widely used benchmark case with concentric circular flux surfaces (Dimits et al. 2000). For this we use a Miller local equilibrium (Miller et al. 1998) with $R_0/a=2.78$, $r/a=0.5$, $q=1.4$, $\hat {s}=0.8$, $\kappa =1.0$, $\delta =0.0$. Unless otherwise noted, the normalized gradient scale lengths are taken to be $a/L_{ni} = a/L_{ne} = 0.8$, $a/L_{Ti} = a/L_{Te}= 2.49$, and we also take $T_i = T_e$. In cases with kinetic electrons we take $m_i/m_e = 3670$ (deuterium ions). This series of benchmarks is inspired by the set of tokamak benchmarks chosen for stella by Barnes et al. (2019).

6.1.1. Linear results

We first make a linear comparison taking a Boltzmann (adiabatic) electron response, so that the quasineutrality equation (2.9) reduces to

(6.1)\begin{equation} \int {\rm d}^3{\boldsymbol v}\ \langle{ h_s \rangle}_{\boldsymbol r} = \sum_{\ell=0}^{N_\ell}\mathcal{J}_\ell H_{\ell,0} =\frac{T_e}{T_i}[\varPhi - \langle\langle \varPhi\rangle\rangle] + \varPhi, \end{equation}

where $\langle \langle \varPhi \rangle \rangle$ denotes the flux-surface average of $\varPhi$. In figure 1 we show normalized growth rates (figure 1a) and real frequencies (figure 1b) as a function of the normalized binormal wavenumber $k_y \rho _i$. GX results are shown with open blue circles connected by dashed lines, while results from GS2 are shown with yellow squares. The agreement between the two codes is very good across the range of unstable $k_y$ values. Here we included a small number of collisions in both codes to provide velocity-space regularization, taking the normalized ion collision frequency to be $\nu _{\rm ii} = 10^{-2} v_{\rm ti}/a$. Hypercollisions are not used in GX in this case. We next make linear comparisons using kinetic electrons between GX, GS2 and stella. Figure 2(a) shows normalized growth rates and real frequencies as a function of $k_y \rho _i$ at ion scales in the electrostatic limit, with $\beta _\mathrm {ref}=10^{-5}$ (we retain electromagnetic fluctuations even at low $\beta$ to alleviate the timestep restriction from the electrostatic $\omega _H$ mode (Lee 1987)). In this case we take $\nu _{\rm ee} = 10^{-2} v_{\rm ti}/a$ and $\nu _{\rm ii} = 1.65\times 10^{-4} v_{\rm ti}/a$; for best comparison with GX, we use the stella's Dougherty collision model option, even though stella also includes a more accurate Fokker–Planck collision operator (which produces results that agree even more closely with GS2 in this case). Hypercollisions were again not used in GX. We observe excellent agreement between GX and stella for the ITG branch $(\omega >0)$, but the agreement is not as good for the trapped electron mode (TEM) branch $(\omega <0)$. If there is a weakness in our scheme it is accurately capturing linear TEM instability, which requires resolving sharp trapped–passing boundaries in velocity space (which is not optimal in $v_\parallel,\mu$ coordinates) and is sensitive to the details of the collision operator (our Dougherty model collision operator is not as accurate as the collision operators in other standard gyrokinetic codes like GS2). Indeed, even getting this level of agreement required 128 Hermite modes in GX (for more resolution details, see appendix G). A more detailed study of TEM modes is left to future work, where we will explore more accurate collision operators and methods to enhance convergence. Note that Frei et al. (2021, 2022a) have demonstrated the success of the Laguerre–Hermite method in resolving TEM modes with a more accurate collision operator. We also show growth rates and frequencies as a function of $k_y \rho _i$ at electron scales in figure 2(b), again in the electrostatic limit. There is once again strong agreement between GX and GS2 for these electron-temperature-gradient (ETG) modes.

Figure 1. Normalized growth rates (a) and real frequencies (b) as a function of the normalized binormal wavenumber $k_y{\rho _i}$ from GX (blue circles and dotted lines) and GS2 (yellow squares) for ‘Cyclone base case’ (CBC) parameters using a Boltzmann electron response.

Figure 2. Normalized growth rates (i) and real frequencies (ii) as a function of the normalized binormal wavenumber $k_y{\rho _i}$ from GX (blue circles and dotted lines), GS2 (yellow squares) and stella (green inverted triangles) for CBC parameters using a kinetic electron response in the electrostatic limit ($\beta _\mathrm {ref} = 10^{-5}$). Ion scales (ITG and TEM) are shown in (a) and electron scales (ETG) are shown in (b).

We finally perform an electromagnetic linear benchmark of the transition from ITG instability to kinetic ballooning mode (KBM) instability. In this benchmark we include perpendicular magnetic fluctuations via finite $A_\parallel$ but we neglect parallel magnetic fluctuations ($\delta B_\parallel =0$). Taking $k_y\rho _i=0.3$, figure 3 shows normalized growth rates and real frequencies as a function of $\beta _\mathrm {ref}$, the reference beta. As $\beta _\mathrm {ref}$ increases, both GX and GS2 agree well and show moderate stabilization of the ITG mode until the transition to KBM around $\beta _\mathrm {ref} = 1.3\,\%$. Additional work studying reactor-relevant electromagnetic instabilities (e.g. microtearing and toroidal Alfvén eigenmodes) with GX will be reported in future publications; successful modelling of collisionless microtearing with a Laguerre–Hermite formulation has been shown by Frei et al. (2023).

Figure 3. Normalized growth rates (a) and real frequencies (b) as a function of the reference beta $\beta _\mathrm {ref}$ from GX (blue circles and dotted lines) and GS2 (yellow squares) for CBC parameters using a kinetic electron response, taking $k_y \rho _i = 0.3$. Both codes show the transition from ITG instability at low $\beta _\mathrm {ref}$ to kinetic ballooning mode (KBM) instability at high $\beta _\mathrm {ref}$.

6.1.2. Nonlinear results

We first perform nonlinear CBC calculations using a Boltzmann electron response. Figure 4(a) shows time traces of the ion heat flux in gyro-Bohm units (computed using the expressions in appendix F), showing very good agreement in the time-averaged heat flux amongst GX and GS2. For GX, we show a case with moderate velocity resolution, $( N_\ell, N_m)=(8,16)$, along with a coarse resolution case with $( N_\ell, N_m)=(4,8)$. Additional details about velocity-space convergence for this case are given in § 7. We also show in figure 4(b) a scan of the ITG parameter, $a/L_{Ti}$, which again shows excellent agreement between GX and GS2 for all cases at both resolutions. The agreement at $a/L_T=1.5$ is particularly notable because this is in the Dimits shift regime, where the system is linearly unstable to ITG but turbulence is suppressed by zonal flow dynamics. This regime was especially troublesome for the Beer gyrofluid model (Beer & Hammett 1996; Dimits et al. 2000), but we see that by extending that model to more moments (and neglecting the collisionless closure schemes) via our Laguerre–Hermite approach we can recover the Dimits shift with $(N_\ell,N_m)=(4,8)$. Hoffmann et al. (2023a) have studied the convergence of the Laguerre–Hermite basis for this case and also found that $(N_\ell,N_m)\geq (4,8)$ is required to accurately capture the Dimits shift. Further from marginality, Hoffmann et al. (2023a) have found that even lower velocity resolution can be used with reasonable accuracy, which is consistent with the convergence study in § 7. Here, both GX and GS2 included a small number of collisions, taking the normalized ion collision frequency to be $\nu _{\rm ii} = 10^{-2} \, v_{\rm ti}/a$. Hypercollisions were included in the GX calculations, with details given in § G.3.

Figure 4. (a) Time traces of ion heat flux $Q_i$, normalized to gyro-Bohm units, $Q_{\rm GB}= n_i T_i v_{\rm ti}\rho _i^2/a^2$, for the CBC with a Boltzmann electron response from GX with moderate (blue) and coarse (green) velocity resolution, and GS2 (yellow). (b) Time-averaged gyro-Bohm-normalized ion heat flux, $Q_i/Q_{\rm GB}$, as a function of the normalized inverse ITG scale length, $a/L_{Ti}$.

Nonlinear calculations with kinetic electrons are presented in figure 5, which shows excellent agreement amongst GX and GS2 for the time average of the heat flux in both the ion and electron channels. For GX, we show a case with moderate velocity resolution, $( N_\ell, N_m)=(4,16)$, along with a higher resolution case with $( N_\ell, N_m)=(16,32)$. Additional details about velocity-space convergence for this case are given in § 7. Here we took $\beta _\mathrm {ref} = 10^{-3}$ along with $\nu _{\rm ee} = 10^{-2} v_{\rm ti}/a$ and $\nu _{\rm ii} = 1.65\times 10^{-4} v_{\rm ti}/a$ for collisions in all codes, and hypercollisions were again used in GX (see § G.4).

Figure 5. Time traces of ion (solid) and electron (dashed) heat flux, normalized to ion gyro-Bohm units, for the CBC with kinetic electrons from GX with high velocity resolution (blue) and moderate velocity resolution (green) and GS2 (yellow).

6.2. Stellarator benchmarks: W7-X

For benchmarks in stellarator geometry, we choose a magnetic configuration from W7-X that has recently been used by González-Jerez et al. (2022) to benchmark stella and GENE. The numerical equilibrium is obtained from VMEC (Hirshman & Whitson 1983), and we choose the so-called bean flux tube with $\alpha _0=0$. To study ITG-driven instabilities and turbulence we take $a/L_{Ti}=3$ and $a/L_n=1$. The generalized twist-and-shift boundary condition of Martin et al. (2018) is used.

6.2.1. Linear results

We first make a linear comparison between GX and stella, taking a Boltzmann electron response. Figure 6 shows the normalized growth rates (figure 6a) and real frequencies (figure 6b) for a range of $k_y$ modes, with strong agreement between the two codes. Both codes use Dougherty model collisions with $\nu _{\rm ii} = 0.01\,v_{\rm ti}/a$. This test corresponds to figure 5 of González-Jerez et al. (2022). As observed in González-Jerez et al. (2022), the discontinuity in the frequency is due to a change in mode structure for two branches of ITG. Both branches have very similar growth rates at $k_y \rho _i = 1.0$, resulting in a small but likely inconsequential disagreement between the codes in the frequency of the fastest-growing mode.

Figure 6. Normalized growth rates (a) and real frequencies (b) as a function of the normalized binormal wavenumber $k_y$ from GX (blue circles and dotted lines) and stella (yellow squares) for W7-X bean flux-tube geometry and ITG parameters using a Boltzmann electron response.

6.2.2. Nonlinear results

Moving on to the nonlinear version of this test case, figure 7 shows time traces of ion heat flux from GX along with the time traces from stella and GENE, which were taken directly from the supplementary data made available by González-Jerez et al. (2022) (and corresponding to figure 12 in that work; however, note the factor of $\sqrt {2}$ difference in the definition of $v_{\rm ti}$, and hence a factor of $2\sqrt {2}$ difference in $Q_{\rm GB}$, in GX relative to the definitions used in stella and GENE). Both stella and GENE used rather high velocity resolution for this case, with $N_{v_\parallel }=60$, $N_\mu =24$. For GX we show a case with moderate velocity-space resolution, $( N_\ell, N_m)=(8,16)$, and a coarse resolution case with $( N_\ell, N_m)=(4,8)$. Both of these agree well with the higher-resolution cases from both stella and GENE, which indicates that the promising Laguerre–Hermite convergence observed in the tokamak benchmarks carries over to stellarators. Additionally, the wallclock time for the GX cases was 90 min (on four GPUs) in the moderate resolution case and 28 min (on one GPU) in the coarse case.

Figure 7. Time traces of ion heat flux $Q_i$, normalized to gyro-Bohm units, $Q_{\rm GB}= n_i T_i v_{\rm ti}\rho _i^2/a^2$, for the W7-X bean flux-tube geometry with a Boltzmann electron response. Results are shown from GX with moderate (blue) and coarse (red) velocity resolution, stella (yellow) and GENE (green).

7. Velocity-space convergence and GPU performance

Part of the motivation for the Laguerre–Hermite pseudospectral approach is the ability to successfully run at low velocity-space resolution without uncontrolled approximations, since in the lowest-resolution limit the system corresponds to established gyrofluid models like the one of Beer & Hammett (1996), albeit without Beer & Hammett's collisionless closure schemes. Another motivation for the pseudospectral approach is its fit for GPU computing, since the spectral algorithm relies heavily on fast transform methods that are well-optimized on GPUs, and the memory requirements of the algorithm are low enough to fit a problem onto one or a few GPUs. We also make exclusive use of single-precision arithmetic, which is often sufficient for turbulence calculations (Maurer et al. 2020); there are no matrix inversions in our algorithm that would require higher precision, and we do not presently target simulations spanning both ion and electron scales that may stress the limits of single precision. The success of the benchmarks presented in the previous section is an additional indication that single-precision arithmetic is sufficient for a wide range of linear and nonlinear calculations.

In this section we examine the velocity-space convergence for the nonlinear CBC calculations with GX along with the performance of the calculations on one or several GPUs. Convergence and performance are complementary factors towards enabling first-principles transport calculations that are fast enough to be used for fusion reactor design along with wide-ranging parameter scans for physics discovery. In this section we will present results detailing convergence and performance for the nonlinear CBC benchmark cases from § 6.1.2.

Starting with the case with Boltzmann electrons, table 1 shows the results of several GX calculations, with coarsening velocity resolution as we progress down the table. For this problem the convergence of the Laguerre–Hermite basis is quite remarkable, allowing accurate results with resolution as coarse as $( N_\ell, N_m) = (4,6)$. This is still a factor of four more moments than used in the Beer & Hammett (1996) six-moment gyrofluid model, but it is quite coarse relative to typical gyrokinetic calculations with $O(10$–$100)$ velocity grid points. Thus, the flexibility of the Laguerre–Hermite representation to smoothly interpolate between gyrofluid and gyrokinetic resolution enables the full nonlinear CBC calculation to be run accurately in less than 4 min on a single GPU, a remarkable result.

Table 1. Velocity-space convergence and performance for the nonlinear CBC with Boltzmann electrons. Here $N_m$ is the number of Hermite modes and $N_\ell$ is the number of Laguerre modes. Other resolution parameters were $N_x=192,\ N_y=64,\ N_z=24$. Accurate ion heat flux calculations can be obtained with resolution as coarse as $( N_\ell, N_m)\geq (4,6)$ (above the double bar). Each simulation was run to $t = 1000 a/v_{\rm ti}$, and we report the total wallclock time in minutes, the time per timestep in seconds, and the average timestep size $\langle \Delta t\rangle$ (normalized to $a/v_{\rm ti}$), which changes with $N_\ell$ and $N_m$ due to linear stability constraints. The number of NVIDIA A100 GPUs used for each calculation is listed in the final column. The resolutions shown in figure 4 are marked with a $\star$.

When we look at the Laguerre and Hermite spectra for these cases shown in figure 8, where $W_{g,s}(\ell,m) = \int |G_{\ell,m}^s|^2\, {{\rm d} x}\, {{\rm d} y}\, {\rm d}z$ is the free energy in each moment, and $W_{g,s}(\ell ) = \sum _m W_{g,s}(\ell, m)$ and $W_{g,s}(m) = \sum _\ell W_{g,s}(\ell, m)$, we can see why the convergence is very good: even for the coarsest resolution, the spectra at the scales that contribute dominantly to the heat flux (small $\ell$ and $m$) are nearly identical. Note that hypercollisions are being used in these calculations, and this is helping to enhance convergence. Without hypercollisions (not shown), the results are only accurate with $( N_\ell, N_m) \geq (6,12)$. This means that hypercollisions enable a five-fold reduction in simulation (wallclock) time for this case.

Figure 8. Laguerre (a) and Hermite (b) free energy spectra, $W_g$, for the nonlinear CBC with a Boltzmann electron response with varying velocity-space resolution. All the spectra are nearly identical at the scales that contribute dominantly to the heat flux (small $\ell$ and $m$), which is consistent with the excellent convergence observed in table 1.

We perform a similar convergence study for the CBC with kinetic electrons, as shown in table 2. As above, $N_\ell =4$ is sufficient for accurate heat flux predictions in both ion and electron channels. However, Hermite convergence is not as strong here, with reasonable accuracy (within 15 % of the highest resolution case) requiring $N_m\gtrsim 16$. Examining the Laguerre and Hermite free energy spectra for these cases in figure 9, we see more structure in the Hermite spectra, especially for the electrons. The oscillatory nature of the electron spectrum, with higher amplitudes in even Hermite modes than odd, could be an indication that the toroidal drifts (which couple every other Hermite mode) are playing a strong role in the dynamics. This is consistent with the results of analytical and numerical convergence studies of the Laguerre–Hermite basis associated with toroidal drifts by Frei et al. (2022a, 2023). Despite this structure, reasonably accurate results can be achieved in 159 min on four GPUs using $( N_\ell, N_m)=(4,16)$. This is quite good for a nonlinear, electromagnetic gyrokinetic ITG simulation with real-mass-ratio kinetic electrons. Additional details about multi-GPU scaling are presented in § 7.1.

Table 2. Velocity-space convergence and performance for the nonlinear CBC with kinetic electrons. Here $N_m$ is the number of Hermite modes and $N_\ell$ is the number of Laguerre modes. Other resolution parameters were $N_x=192,\ N_y=64,\ N_z=24$. Reasonably accurate heat flux calculations (within 15 % of the highest resolution case) can be obtained with resolution as coarse as$( N_\ell, N_m)\geq (4,16)$ (above the double bar). Each simulation was run to $t = 600 a/v_{\rm ti}$, and we report the total wallclock time in minutes, the time per timestep in seconds, and the average timestep size $\langle \Delta t\rangle$ (normalized to $a/v_{\rm ti}$). The number of NVIDIA A100 GPUs used for each calculation is also shown. The resolutions used for the GX calculations shown in figure 5 are marked with a $\star$.

Figure 9. Laguerre (i) and Hermite (ii) spectra for ions (a) and electrons (b) for the nonlinear CBC with kinetic electrons with varying velocity-space resolution.

Note that using more Hermite modes not only adds to the size of the problem, but also makes the Courant–Friedrichs–Levy (CFL) constraint on the timestep more restrictive because the maximum parallel velocity $v_{\parallel \mathrm {max}}$ on the Hermite collocation grid increases with the number of Hermite modes. In future work we will try to alleviate this issue by constraining the maximum velocity on the collocation grid. However, a commonly used method of scaling the argument of the Hermite polynomials to reduce the extents of the collocation grid is ill-suited to our problem, since the resulting projection of a Maxwellian onto the scaled basis often diverges (Fok, Guo & Tang 2001). Instead, an approach similar to that used by Candy et al. (2016) to bound the energy grid in CGYRO could be used here.

7.1. GPU scaling studies

For single-GPU calculations, the GPU threading architecture effectively enables dynamic parallelism over the entire phase space. In figure 10, we show that for a nonlinear CBC calculation with Boltzmann electrons, the algorithm scales roughly linearly with the number of radial grid points, $N_x$. This is better than the $O(N\log N)$ cost that would be expected when the fast Fourier transform (FFT) in the pseudospectral algorithm dominates the cost. This is also better scaling than other gyrokinetic algorithms that require $O(N^2)$ operations due to matrix inversion.

Figure 10. Scaling of time per timestep (in seconds) versus the number of radial grid points, $N_x$, for nonlinear CBC with Boltzmann electrons on a single GPU. Ideal linear scaling is shown with the dashed line, while the dot–dashed line shows $N_x \log N_x$ scaling (the expected scaling when FFTs dominate the algorithm).

When considering multi-GPU parallelization, the cost of inter-GPU memory transfers can be quite large compared with floating point operations. Thus, it is important to design the scheme to minimize memory transfers and take advantage of the capability of modern GPUs to overlap communication with computation. Thus, in GX we have chosen a targeted multi-GPU parallelization scheme that currently parallelizes only species and Hermite modes across GPUs. In this scheme, communication is required to compute the total charge density and currents in the field equations, which takes the form of an all-reduce operation on a configuration-space-sized ($N_{k_x}\times N_{k_y}\times N_z$) object. Additionally, since the equation for the Hermite mode $m$ couples to modes $m-2$, $m-1$, $m+1$ and $m+2$, we use ‘halo’ (or ‘ghost’) modes on each GPU, similar to standard parallelization schemes for finite differencing. Halo exchange is thus the other dominant form of communication, requiring inter-GPU transfers of objects of size $N_{k_x}\times N_{k_y}\times N_z\times N_\ell \times 2$. We overlap the halo exchange with computation of most of the nonlinear terms (halo modes are only required for the nonlinear terms involving $A_\parallel$, and even for these the terms on the interior modes can be computed before the halo transfers have been completed).

In figure 11 we show a strong scaling study for nonlinear CBC cases with two kinetic species. In each case we fix resolution (shown in the legend) and plot the time per timestep (figure 11a) and scaling efficiency (figure 11b) as we increase the number of GPUs used for the calculation. These calculations have been performed on Perlmutter at NERSC, where each node contains four A100 GPUs. The three cases have the same configuration space resolution and number of kinetic species as the cases in table 2. For reference, the (16,64) case running on a single GPU consumes approximately 90 % of the available GPU memory on a 40GB A100 GPU. In each case, the scaling efficiency is above 75 % parallelizing across four GPUs. The efficiency also improves as the problem size increases, which is expected because there is more computing time relative to communication time.

Figure 11. Strong scaling study for the nonlinear CBC with kinetic electrons showing the time per timestep (a) and the scaling efficiency (b) as a function of number of GPUs used, with fixed resolution for each curve. In each case the ideal scaling is shown with a dashed black line. The Laguerre–Hermite resolution parameters are listed in the legend as $(N_\ell, N_m)$. The other resolution parameters are $N_x=192,\ N_y=64, N_z=24,\ N_\mathrm {species}=2$.

We also show in figure 12 weak scaling studies for the nonlinear CBC with two kinetic species and the same configuration space resolution as the cases in table 2. Here, we increase the number of Hermite modes proportionally to the number of GPUs used, so that $N_m=8 N_\mathrm {gpu}$ in each case. After parallelizing first over the two kinetic species, this results in 16 Hermite modes per GPU in all the calculations shown. We show the weak scaling for both $N_\ell =4$ (blue) and $N_\ell =8$ (yellow). The weak scaling efficiency, shown in (figure 12b), is ideal in both cases up to 32 GPUs ($N_m = 256$), despite the fact that using more than four GPUs requires multiple nodes on Perlmutter, and communication between nodes can be slower than communication within a node due to the details of the interconnect hardware.

Figure 12. Weak scaling study for the nonlinear CBC with kinetic electrons showing the time per timestep (a) and the scaling efficiency (b) as a function of number of GPUs used, with the number of Hermite modes $N_m$ scaling with the number of GPUs as $N_m=8N_\mathrm {gpu}$. We show results for both $N_\ell =4$ (blue) and $N_\ell =8$ (yellow). The ideal scaling is shown with dashed lines. The other resolution parameters are $N_x=192,\ N_y=64,\ N_z={24}, N_\mathrm {species}=2$.

8. Conclusion and future opportunities

In this work we have presented GX, a GPU-native gyrokinetic code focused on tokamak and stellarator design at reactor scales. We have described the numerical algorithms that we use to solve the electromagnetic $\delta f$ gyrokinetic system, in particular the discretization scheme that is pseudospectral in the entire five-dimensional phase space, leveraging the Laguerre–Hermite velocity-space formulation developed by Mandell et al. (2018). We have shown several linear and nonlinear benchmarks against established flux-tube gyrokinetic codes in both tokamak and stellarator geometry, verifying that GX correctly simulates gyrokinetic turbulence in the small $\rho _*$ limit. The combination of GPU acceleration and favourable convergence properties of the Laguerre–Hermite velocity-space basis for nonlinear problems enables useful turbulence simulations in minutes.

Additional work will also focus on improving the efficiency of simulations that include kinetic electrons. At present, the need to resolve the fast parallel electron motion results in an increase in simulation cost by a factor of roughly $2\sqrt {m_i/m_e}\approx 120$ compared with simulations with Boltzmann electrons (the factor of 2 is from the additional kinetic species, and the square root of the mass ratio comes from $v_{\rm te}/v_{\rm ti}$). We have also shown that kinetic electron cases may require more Hermite resolution, especially if resolving the trapped–passing boundary is essential. Exploration of different choices of velocity-space coordinates for the electrons or improvements to the basis functions used (for example, a method to bound the Hermite collocation points to constrain $v_{\parallel \mathrm {max}}$, as in Candy et al. (2016), which could alleviate the CFL timestep restriction) could increase efficiency. In addition, the reduced electron models of Beer & Hammett (1996), Snyder & Hammett (2001) and Abel & Cowley (2013) provide insight into how to eliminate the fast time scales associated with parallel electron motion. A combination of these approaches with implicit timestepping schemes could enable kinetic electron simulations with nearly the same efficiency as the Boltzmann electron simulations we have presented. Additionally, we are investigating machine-learning methods for subgrid models and closures (Barbour et al. 2021, 2022), as well as methods for dynamically adapting the number of modes in the system.

An essential target for GX is transport time scale macroscale profile evolution via coupling to a transport solver like Trinity (Barnes et al. 2009; Qian et al. 2022) or a steady-state profile prediction solver like TGYRO (Candy et al. 2009) or PORTALS (Rodriguez-Fernandez, Howard & Candy 2022). This will be described in a forthcoming paper. In this case, several GX flux-tube calculations are run in parallel at various radii, and the resulting turbulent fluxes are used to advance the transport equations for the equilibrium profiles on transport time scales. The efficiency of GX makes these simulations quite tractable in comparison with previous efforts. Further, since transport in a fusion reactor is usually quite stiff, small changes in profile gradients result in large changes in transport fluxes; conversely, this means that errors of order 10 % are quite tolerable since they will result in relatively negligible differences in the final equilibrium profiles. Thus, we can also leverage the flexibility of GX to run at quite low velocity resolution with only 10 %–20 % errors in the turbulent fluxes, as shown in the convergence studies.

Finally, since GX is specialized for reactor design, using GX as a turbulence model in the inner loop of optimization frameworks is of great interest. Kim et al. (2024) have leveraged the DESC (Dudt et al. 2022) optimization framework to demonstrate that stellarators can be optimized for turbulence using nonlinear GX calculations in the optimization loop. Additionally, coupling of GX to the SIMSOPT stellarator optimization framework (Landreman et al. 2021) is in progress. Following the work of Highcock et al. (2018), optimization of core equilibrium profiles (and hence fusion power) is also tractable with a fast multiscale transport model composed of GX coupled to a transport solver.