2.1. Ray tracing/FP modelling of heating and current drive

Several approaches are possible to simulate the RF current drive profile in tokamak plasmas. Analytical and reduced numerical (e.g. adjoint method) models provide a reasonable estimate of the total driven current but lack key physics necessary to determine the current profile shape that can only be realised by solving the FP equation. Moving upward in complexity to ray tracing/FP modelling includes the necessary wave propagation and quasi-linear Landau damping physics for accurate current profile prediction in several frequency ranges (lower hybrid, electron cyclotron). Ray tracing makes use of the simplifications of geometrical optics and reduces the problem from solving Maxwell's equations in toroidal geometry (the ‘fullwave’ method) to the following set of ray equations:

(2.1a)\begin{gather} \frac{{\rm d}\boldsymbol{x}}{{\rm d}\tau}=\frac{\partial D_0}{\partial \boldsymbol{k}}, \end{gather}

(2.1b)\begin{gather} \frac{{\rm d}\omega}{{\rm d}\tau}=\frac{\partial D_0}{\partial t}, \end{gather}

(2.1c)\begin{gather} \frac{{\rm d}\boldsymbol{k}}{{\rm d}\tau}={-}\frac{\partial D_0}{\partial \boldsymbol{x}}, \end{gather}

(2.1d)\begin{gather} \frac{{\rm d}t}{{\rm d}\tau}={-}\frac{\partial D_0}{\partial \omega}, \end{gather}

where $D_0$ is the wave dispersion relation, $\omega$ is the wave frequency, $\boldsymbol {k}$ is the wavenumber, $\tau$ is the ray step parameter, $t$ is the time coordinate and $\boldsymbol {x}$ is the spatial coordinate. Equation (2.1b) is trivial under the condition that the time variation of $D_0$ is small on the timescale of the operating frequency $f_0$, which is well satisfied for waves in the gigahertz range.

Solving the FP equation is required for both fullwave and ray-tracing approaches to determine power absorption and current drive. The FP equation

(2.2)\begin{equation} \frac{{\rm d} f}{{\rm d} t} = RF + C, \end{equation}

describes the time evolution of the distribution function, $f$, in the presence of RF damping, $RF$, and collisions, $C$. The RF term

(2.3)\begin{equation} RF = \frac{\partial}{\partial \boldsymbol{v}} D_{QL}(\boldsymbol{v}) \frac{\partial}{\partial \boldsymbol{v}} f_{e0}, \end{equation}

depends on the distribution function and the quasi-linear diffusion operator, $D_{QL}$. The wave solvers provide $D_{QL}$ as an input to the FP solver, as shown in figure 1. The FP solver then feeds information on power absorption back to the wave equation solver and the process iterates until convergence. In the ray-tracing approach, the ray trajectories need not be recalculated at each time step of the FP solver, only the power absorption along the ray. A schematic of the workflow between wave equation solvers and FP solvers is shown in figure 1. In this study we use the GENRAY (Smirnov & Harvey Reference Smirnov and Harvey1995) ray-tracing code coupled to CQL3D (Harvey & McCoy Reference Harvey and McCoy1992; Petrov & Harvey Reference Petrov and Harvey2016) for solving the FP equation, although the approach may be expanded to include other RF simulation tools such as TORIC (Brambilla Reference Brambilla1996) for fullwave simulations.

Figure 1. Wave equation solvers (TORIC+Petra-M or GENRAY) couple to FP equation solvers (CQL3D) through the quasi-linear diffusion operator, $D_{QL}$. The FP solver perturbs the distribution function and couples information on quasi-linear RF power absorption, $P_{abs}$, back to the wave solver. This process iterates until the solutions converge. Rough per-iteration wall clock runtimes for each code are indicated in the figure. TORIC+Petra-M and GENRAY figures reproduced from Shiraiwa et al. (Reference Shiraiwa, Wright, Lee and Bonoli2017) and Wallace et al. (Reference Wallace, Parker, Bonoli, Hubbard, Hughes, LaBombard, Meneghini, Schmidt, Shiraiwa and Whyte2010), respectively.

Both GENRAY and CQL3D require information about the plasma density, temperature and ion composition, as well as the magnetic equilibrium as inputs. In addition, CQL3D requires the toroidal DC electric field, whereas GENRAY requires the wave parameters (power, $n_{\|}$ spectrum, antenna position) as initial conditions. The outputs of GENRAY, namely the ray trajectories and quasi-linear diffusion operator, are of indirect rather than direct interest from the perspective of plasma control and integrated modelling. The outputs of CQL3D (wave power deposition and current profiles) are of direct interest and will be the prediction targets of our ML models.

2.2. Using ML to accelerate computations

Modern computational resources have advanced numerical simulation of physical systems to an increasing range of spatial and temporal scales. However, even with high-performance computing techniques, it is a challenge for real-time demand response in feedback control. Emerging data-driven methods have revolutionised how researchers model, predict and eventually control these complex systems in a diverse range of fields, including but not limited to automation, climate, combustion, fluids, high-energy physics, plasma science and plasmonics (Malkiel et al. Reference Malkiel, Mrejen, Nagler, Arieli, Wolf and Suchowski2018; Duraisamy, Iaccarino & Xiao Reference Duraisamy, Iaccarino and Xiao2019; Wilkes et al. Reference Wilkes, Emmett, Beltran, Woodward and Carling2020; Yellapantula et al. Reference Yellapantula, de Frahan, King, Day, Grout, H. and A.2020; Abbate, Conlin & Kolemen Reference Abbate, Conlin and Kolemen2021; Bai & Peng Reference Bai and Peng2021; Hatfield et al. Reference Hatfield, Gaffney, Anderson, Ali, Antonelli, Başeğmez du Pree, Citrin, Fajardo, Knapp and Kettle2021; Maschler & Weyrich Reference Maschler and Weyrich2021; Brunton & Kutz Reference Brunton and Kutz2022). A variety of ML techniques have been used such as deep neural networks (DNNs) to fit simulation data. The resulting surrogate models have been used in larger frameworks. For example, global climate models use ML models to represent kernels drawn from observations (Reichstein et al. Reference Reichstein, Camps-Valls, Stevens, Jung, Denzler, Carvalhais and Prabhat2019). A common thread is the use of numerical models and codes to generate training data to produce surrogate models with a low mean squared error (MSE). This is the approach of this paper.

2.3. The landscape of ML in plasma physics and fusion

The application of ML techniques to fusion energy and plasma physics dates to the late 20th century, with initial applications of neural networks to the topics of tokamak plasma control (Bishop et al. Reference Bishop, Cox, Haynes, Roach, Smith, Todd, Trotman and J.G.1992, Reference Bishop, Haynes, Smith, Todd and Trotman1995a; Wróblewski Reference Wróblewski1997), disruption prediction (Hernandez et al. Reference Hernandez, Vannucci, Tajima, Lin, Horton and McCool1996; Wroblewski, Jahns & Leuer Reference Wroblewski, Jahns and Leuer1997), energy confinement prediction (Allen & Bishop Reference Allen and Bishop1992) and data analysis (Lister & Schnurrenberger Reference Lister and Schnurrenberger1991; Bishop et al. Reference Bishop, Strachan, O'Rourke, Maddison and Thomas1993; Morabito & Versaci Reference Morabito and Versaci1997), among other topics. These early applications showed the promise of ML techniques in fusion research. Considerable advances in both computing hardware and ML techniques after the turn of the century resulted in renewed interest in ML applications for fusion in recent years.

ML continues to be an extremely effective tool in some of the early adopted topics, particularly disruption prediction (Rea & Granetz Reference Rea and Granetz2018; Rea et al. Reference Rea, Montes, Pau, Granetz and Sauter2020; Churchill, Tobias & Zhu Reference Churchill, Tobias and Zhu2020a; Churchill et al. Reference Churchill, Choi, Kube, Chang, Klasky, J., B., A., O., S. and T.2020b). A more recent application of ML to fusion research is in the area of accelerating computationally intensive simulation models. Researchers developed fast surrogate models for plasma transport (Meneghini et al. Reference Meneghini, Snoep, Lyons, McClenaghan, Imai, Grierson, Smith, Staebler, Snyder and Candy2020; van de Plassche et al. Reference van de Plassche, Citrin, Bourdelle, Camenen, Casson, Dagnelie, Felici, Ho and Van Mulders2020) and neutral beam heating/current drive (Boyer, Kaye & Erickson Reference Boyer, Kaye and Erickson2019; Morosohk, Boyer & Schuster Reference Morosohk, Boyer and Schuster2021), used surrogate model frameworks for validation of plasma transport codes (Rodriguez-Fernandez et al. Reference Rodriguez-Fernandez, White, Creely, Greenwald, Howard, Sciortino and Wright2018) and also employed ML techniques to reduce computation time for codes such as XGC (Miller et al. Reference Miller, Churchill, Dener, Chang, Munson and Hager2021). The work of Boyer and Morosohk, in which they create a fast neural-network=based surrogate model for NUBEAM, is particularly relevant to the work presented in this paper. The work presented here represents the first efforts to create a ML-based surrogate model for RF heating and current drive.

3.1. Source data

Training ML models requires a database of simulations to draw training, validation and testing data. In the ML context, ‘validation’ refers to the process of using a holdout set of data to provide an unbiased evaluation of the model while tuning the hyperparameters. To this end, the ${\rm \pi}$Scope (Shiraiwa et al. Reference Shiraiwa, Greenwald, Stillerman, Wallace, London and Thomas2018) workflow engine generates a set of $16\,384$ GENRAY (Smirnov & Harvey Reference Smirnov and Harvey1995)/CQL3D (Harvey & McCoy Reference Harvey and McCoy1992) input files and submits five of these simulations at a time to run in parallel on the MIT Engaging cluster at the Massachusetts Green High Performance Computing Center. The wall-clock time required to generate the entire database was approximately two weeks. As each GENRAY/CQL3D simulation completes, ${\rm \pi}$Scope incorporates the relevant input and output quantities into a NETCDF database. Because a single NETCDF file containing all simulation data would be too large for efficient data read/write, ${\rm \pi}$Scope chunks the NETCDF files with $100$ simulations per file.

The GENRAY/CQL3D input files contain hundreds of parameters covering both the plasma and the numerical methods used in the simulation. The vast majority of these parameters do not vary in the simulation database, with nine zero-dimensional parameters describing the plasma and lower hybrid system varying across the ranges listed in table 1. The simulation database consists of LHCD simulations for the EAST tokamak (Wu et al. Reference Wu2007). The LHCD wave frequency is 4.6 GHz, and the antenna size/location match that of the LHCD antenna on EAST (Liu et al. Reference Liu, Ding, Li, Wan, Shan, Wang, Liu, Zhao, Li and Li2015). The kinetic profiles are analytic and defined by a central electron temperature/density ($n_{e0}$, $T_{e0}$) and edge density/temperature ($n_{eb}$, $T_{eb}$) according to the equations

(3.1)\begin{equation} n_e = (n_{e0}-n_{eb})\times(1-\rho^2) + n_{eb} \end{equation}

and

(3.2)\begin{equation} T_e = (T_{e0}-T_{eb})\times(1-\rho^2) + T_{eb}, \end{equation}

where $\rho$ is the normalised minor radius. In the GENRAY/CQL3D database $n_{e0}$ and $T_{e0}$ vary whereas $n_{eb}$ and $T_{eb}$ remain constant at $7.5 \times 10^{18}\,{\rm m}^{-3}$ and 0.1 keV, respectively.

Table 1. Parameters varied in the GENRAY/CQL3D database.

The magnetic equilibrium is an analytic model based on the work described in (Guazzotto & Freidberg Reference Guazzotto and Freidberg2021). The major radius of the plasma, plasma current and toroidal magnetic field vary over typical ranges for the EAST tokamak in this database. However, other equilibrium parameters such as inverse aspect ratio (0.24), triangularity (0.5), elongation (1.8) and magnetic configuration (double null) remain constant for all simulations in the database.

For sampling the parameter space, we chose to use Latin hypercube sampling (LHS) (McKay, Conover & Whiteman Reference McKay, Conover and Whiteman1976). LHS covers a large dimension parameter space with far fewer sample points than a uniform grid and guarantees better uniformity across all columns/rows of the multi-dimensional parameter space than random sampling methods. Given the utility of this sampling method, ${\rm \pi}$Scope now incorporates LHS functionality in parametric scans of input variables.

The parameter ranges in table 1 cover or exceed typical values for each parameter seen on the EAST tokamak, however not all combinations of parameters are relevant. In particular, a significant number of the simulations occur in the ‘slide-away’ regime where the collisional drag on fast electrons is insufficient to counteract the accelerating potential of the toroidal DC electric field:

(3.3)\begin{equation} m_e v \nu_{{\rm coll}}(v) < q E_{{\rm crit}}, \end{equation}

where $m_e$ is the electron mass, $v$ is the fast electron velocity, $\nu _{{\rm coll}}(v)$ is the fast election collision frequency, $q$ is the electron charge and $E_{{\rm crit}}$ is the critical toroidal DC electric field. Using

(3.4)\begin{equation} \nu_{{\rm coll}}(v) = \frac{q^4 n_e (\ln\varLambda)}{4{\rm \pi}\epsilon_0 m_e^2 v^3}, \end{equation}

results in

(3.5)\begin{equation} E_{{\rm crit}} = \frac{q^3 n_e (\ln \varLambda)}{4 {\rm \pi}\epsilon_0 m_e v^2} \end{equation}

which is the familiar expression for the Dreicer field (Dreicer Reference Dreicer1959). Taking $v = c/n_{\|}$ as the upper limit of the quasi-linear plateau for slide-away electrons (as opposed to $v=c$ for true runaway electrons (Connor & Hastie Reference Connor and Hastie1975)) yields

(3.6)\begin{equation} E_{{\rm crit}} = 7.9\times10^{{-}22} n_e n_{\|}^2 \end{equation}

with $E_{{\rm crit}}$ measured in ${\rm V}\,{\rm m}^{-1}$. The slide-away regime can be modelled accurately with CQL3D by increasing the resolution and range of the electron energy grid, however this requires additional computational resources and the slide-away regime is generally avoided in LHCD experiments. A threshold of $E_{{\rm DC}}/E_{{\rm crit}} > 0.9$ delineates the slide-away regime to be filtered from the dataset, resulting in the final dataset used to design the ML models as described in the following.

3.2. ML models

This section describes the three different ML models designed to predict current profiles and power deposition. Each of these ML models carry different characteristics which may affect prediction performance in different ways. Our goal is to evaluate multiple ML models as surrogates for the traditional more expensive type of computation described in the previous section.

3.2.3. Gaussian processes regression

The Gaussian processes regression (GPR) model is a probabilistic supervised ML framework that has been widely used for regression and classification tasks. This model can make predictions incorporating prior knowledge and provides uncertainty measures over predictions (Murphy Reference Murphy2013). Similar to other ML models, GPR is classified as a non-parametric model (Liu et al. Reference Liu, Ong, Shen and Cai2020), i.e. when conducting regressions using GPR, the complexity or flexibility of the model is not limited by the number of parameters. More specifically, the model does not assume a specific form for the function that maps the input to the output data, which would be the case for a linear regression model for example.

Formally, a Gaussian process mode is a probability distribution over possible functions that fit a set of points. The regression function modelled by a multivariate Gaussian is given as

(3.7)\begin{equation} P(\boldsymbol{f}|\boldsymbol{X}) = \mathcal{N}(\boldsymbol{f}|\boldsymbol{\mu},\boldsymbol{K}), \end{equation}

where $\boldsymbol {X} = [\boldsymbol {x}_1,\ldots,\boldsymbol {x}_n]$, $\boldsymbol {f} = [f(\boldsymbol {x}_1),\ldots,f(\boldsymbol {x}_n)]$, $\boldsymbol {\mu } = [m(\boldsymbol {x}_1),\ldots,m(\boldsymbol {x}_n)]$ and $K_{ij} = k(\boldsymbol {x}_i,\boldsymbol {x}_j)$. Here $\boldsymbol {X}$ are the observed data points, $m$ represents the mean function and $k$ represents a positive-definite kernel function. With no observation, the mean function is default to be $m(\boldsymbol {X})=0$ given that the data are often normalised to a zero mean. The Gaussian processes model is a distribution over functions whose shape is defined by $\boldsymbol {K}$. If points $\boldsymbol {x}_i$ and $\boldsymbol {x}_j$ are considered to be similar by the kernel, function outputs of the two points, $f(\boldsymbol {x}_i)$ and $f(\boldsymbol {x}_j)$, are expected to be similar. In the regression process, given the observed data and a mean function $\boldsymbol {f}$ estimated by these observed points, predictions can be made at new points $\boldsymbol {X}_*$ as $\boldsymbol {f}(\boldsymbol {X}_*)$. Based on the definitions, one can directly derive the joint distribution of $\boldsymbol {f}$ and $\boldsymbol {f}_*$ based on $\boldsymbol {K}$ and, more importantly, the conditional distribution $P(\boldsymbol {f}_*|\boldsymbol {f},\boldsymbol {X},\boldsymbol {X}_*)$ over $\boldsymbol {f}_*$, which is needed for the regression process. These derivations are explained in detail in earlier work (Rasmussen & Williams Reference Rasmussen and Williams2005).

The hyperparameters of the GPR model consist of the kernel and its lengthscale and variance. In this work, we explore six different kernels: Exponential, Matern12, Matern32, Matern52, SquaredExponential and Rational Quadratic. Each of these kernels is added together with a Linear kernel. Generally speaking, the lengthscale controls the smoothness of the learnt functions and the variance controls the scale of values of the learnt functions. Given that the input dimension for our problem is nine, the dimension of the lengthscale is also nine. To optimise these hyperparameters, we perform a grid search using the interval $[0.25,2.1]$ for each lengthscale value for each input dimension, and the interval $[0.1,2.2]$ for the variance of the kernels. This search is performed for all six different kernels. We present the final results for this optimisation process in § 4.

3.3. Design and implementation

To implement the three proposed surrogate models we use Python and its supporting libraries. The MLP model is implemented using the Pytorch deep learning framework. We use the Adam solver (Kingma & Ba Reference Kingma and Ba2014), an extension of stochastic gradient descent in the implementation of optimising the weights of the MLP, which is computationally efficient and requires less tuning for hyperparameters.

As for the random forest model, we use scikit-learn (Pedregosa et al. Reference Pedregosa, Varoquaux, Gramfort, Michel, Thirion, Grisel, Blondel, Prettenhofer, Weiss and Dubourg2011), more specifically sklearn.ensemble.RandomForestRegressor, which provides the functions fit and predict for training and prediction, respectively.

Finally, to implement the GPR model, we use an existing Python library called GPflow (Matthews et al. Reference Matthews, van der Wilk, Nickson, Fujii, Boukouvalas, León-Villagrá, Ghahramani and Hensman2017), which is based on GPy  (Sheffield Machine Learning Software 2012). Although these two libraries are very similar to each other, GPflow leverages TensorFlow for faster computation and gradient calculation, and focuses on variational inference and Markov Chain Monte Carlo methods. One special advantage about this library pertinent to our work is its recent extension to inter-domain approximations and multi-output priors (van der Wilk et al. Reference van der Wilk, Dutordoir, John, Artemev, Adam and Hensman2020), which is the case of our dataset. As mentioned previously, one essential component of a GPR model is the kernel function. GPflow provides implementations for the most well-known kernels used in the literature, including Exponential, Linear, Matern family, and others. The kernel selection is included in the hyperparameter optimisation process explained in § 3. For the optimisation process, we use the model gpflow.models.GPR also implemented in GPflow, along with gpflow.optimizers.Scipy for function minimisation.