Fast equilibrium reconstruction by deep learning on EAST tokamak

Reconstructing magnetic configuration using magnetic measurements is a routine task of tokamak operation. There are many equilibrium solvers, e.g., EFITLao_1985 (1, 2, 3, 4, 5, 6), that can do this kind of reconstruction by solving the Grad-Shafranov equation under the constraint of magnetic measurements. In recent years, accumulation of data resulting from these reconstruction practices, along with the development of machine learning algorithms, software frameworks and computing power, have made it possible to train deep neural networks to provide reconstructions as accurate as those by EFIT. This has been demonstrated on KSTARKSTAR_2020 (7) and DIII-DDIII-D_2022 (8).

On the EAST tokamakWan_2017 (9), EFIT has been routinely used in tokamak operations for more than ten years and substantial equilibrium data have been accumulatedQian2009 (10, 11, 12, 13). In this paper, we report the results of magnetic reconstruction by a deep neural network trained on the magnetic measurements and EFIT reconstructed 2D magnetic poloidal flux.

There are two versions of EFIT used on EAST, one is for real-time control and one for off-line analysis. The former is often restricted to lower accuracy due to the time constraint of real time control, while the latter is of higher accuracy. In this work, we use the off-line EFIT data in training the neural network. The model trained this way have the higher accuracy as the off-line EFIT while still meets the time constraint of the real-time control.

There are many hyperparameters in a neural network that usually need to be set manually, such as number of hidden layers, units per layer, mini-batch size, learning rate, number of epochs of training. In recent years, there appear optimization libraries that can automatically set the values of these hyperparameters. In this work, we use the Optuna optimization frameworkoptuna_2019 (14) in setting the hyperparameters. The hyperparameters found this way turn out to be much better than our previously manually set ones in terms of the model accuracy. The size of the network architecture found by the automatic hyperparameter tuning turns out to be relatively small (with less than 2 million parameters). This small size allows for very fast equilibrium construction that can be easily deployed.

The input to the network is limited to only the magnetic measurements. The output of the network is the 2D poloidal magnetic flux function, $\Psi(R,Z)\equiv A_{\phi}R$, which is related to the poloidal magnetic field, $B_{R}$ and $B_{Z}$, by

$$B_{R}=-\frac{1}{R}\frac{\partial\Psi}{\partial Z},$$

(1)

$$B_{Z}=\frac{1}{R}\frac{\partial\Psi}{\partial R},$$

(2)

where $(R,\phi,Z)$ are the cylindrical coordinates. The 2D contours of $\Psi$ in $(R,Z)$ plane correspond to the magnetic surfaces. We compare the inner magnetic surfaces and the LCFSs with those given by EFIT, in order to evaluate the accuracy of $\Psi$ predicted by the network. We also calculate the normalized internal inductance, $l_{i}$, which is a quantity that is solely determined by $\Psi$ and thus can reflect how accuracy the predicted $\Psi$ is. The time evolution of the internal inductance in full discharges is compared with that provided by EFIT. All of the comparisons show good agreement, demonstrating the accuracy of the machine learning model, which has the high spatial resolution as the off-line EFIT while still meets the time constraint of real-time control.

The rest of this paper is organized as follows. Section II presents how the data are collected and normalized. Sec. III explains the structure of our neural network and how the hyperparameters are chosen by automatic optimization. In section IV, we test the predicting capability of the trained network. Section V discusses a small network used to predict some volume-averaged quantities, namely the plasma stored energy $W_{\operatorname{mhd}}$, normalized toroidal beta $\beta_{N}$, and edge safety factor $q_{95}$. A brief summary is given in section VI.

Figure 1a illustrates the poloidal locations of the magnetic measurements used as inputs to our model. A typical time evolution of some of the magnetic measurements from EAST discharge 113019 are plotted in Figure 1b-f.

The inputs (features) to the neutral network (NN) are 84 magnetic measurements: 35 poloidal magnetic flux ($\Psi_{\operatorname{FL}}$) values measured by flux loops, 34 equilibrium poloidal magnetic field (MP) values measured by magnetic probes, 14 poloidal field (PF) coil currents and 1 plasma current (I_p) measured by a Rogowski loop.

The outputs (targets) of the neutral network are the values of the poloidal magnetic flux $\Psi$ at $R\times Z=129\times 129=16641$ spatial locations. The $\Psi$ used in the training process is computed by off-line EFIT and are downloaded from EAST MDSplus server (mds.ipp.ac.cn). The input and the output signals are interpolated to the same time slices before they are fed to the NN.

The inputs and outputs are summarized in Table 1.

The data used in training, validation and testing process were downloaded from the EAST MDSplus server by using Python API, which scans a series of discharges and automatically skips discharges where necessary signals are missing. Specifically, we scan every 5 discharges among all the discharges spanning from #114000 to #117000, resulting in total 45,544 equilibria (time slices). These discharges are from experiments performed in one EAST campaign from June to July in 2022. This range is casually chosen with no particular criterion, except that we prefer recent discharges and avoid old discharges because locations of some magnetic probes were changed in previous campaigns. The auxiliary heating methods on EAST used in this campaign include neutral beam injection (50-70keV Deuterium beam), lower-hybrid waves (2.45GHz and 4.6GHz), electron cyclotron waves (140GHz), and ion cyclotron waves (25-70MHz). Typical values of total heating source power are between 4-10MW.

Figure 2 is the distribution of the EFIT equilibrium data in $(l_{i},\beta_{N})$ plane, where $l_{i}$ is the normalized internal induction and $\beta_{N}$ is the normalized plasma beta.

The collected data are split into three sets: training set (81%), validation set (9%), and testing set (10%), where training set is used in training the NN, validation set is used in monitoring potential overfitting and tuning hyperparameters, and testing set is used in testing the predicting capability of the trained model.

Figure 1b-f shows that there is a difference of up to six orders of magnitude in the values of the input signals. In order to eliminate scale differences among features, we use the min-max normalization method to normalize the input data. The general formula is given by

$$x^{\prime}=\frac{x-x_{\min}}{x_{\max}-x_{\min}}$$

(3)

where $x$ is the original value of the feature, $x^{\prime}$ is the normalized value, $x_{\min}$ and $x_{\max}$ are respectively the minimal and maximal value of a feature in the data sets excluding the testing set. The $x_{\min}$ and $x_{\max}$ obtained here are then used to normalize the input data in the testing set when doing prediction using the trained NN.

Figure 3 plots the time evolution of the normalized input signals corresponding to those in Fig. 1b-f.

The magnitude of the output ($\Psi$ in SI units) is near 1, so no normalization is applied to it.

An artificial neural network is a kind of computational network, which usually consists of multiple layers: input layer, one or more inner layers (known as hidden layers), and a layer of outputs. Each layer is made of units. Each unit in the computing layers (hidden and output layers) receives information and process the information using some linear transform (matrix multiplication) and some nonlinear transform (activation function).

A fully-connected feed-forward network showed in figure 4 is used here to predict the poloidal magnetic flux $\Psi$ based on the magnetic measurements. Here “fully-connected” means that each unit of a computing layer receives information from all the units of the previous layer. “Feed-forward” means that information move in only one direction (from the input layer to the hidden layers, and to the output layers), no cycles or loops, and no intra-layer connections.

Each unit (neuron or node) in the computing layers has trainable parameters, often called weights and biases. Denote the output of the $j^{\operatorname{th}}$ neuron in $l^{\operatorname{th}}$ layer by $a_{j}^{l}$, then a neural network model assumes that $a_{j}^{l}$ is related to the $a^{l-1}$ (output of the previous layer) via

$$a^{l}_{j}=\sigma\left(\sum_{k}w_{jk}^{l}a^{l-1}_{k}+b^{l}_{j}\right),$$

(4)

where $w^{l}_{jk}$ and $b_{j}^{l}$ are the weight and bias, the summation is over all neurons in the $(l-1)^{\operatorname{th}}$ layer, and $\sigma$ is a function called activation function. The weights and biases will be adjusted in the training process by gradient descent methods to reduce the loss (cost or error) function, which is defined in this work as

$$L(\mathbf{w},\mathbf{b})\equiv\frac{1}{2n}\sum_{i=1}^{n}\|\mathbf{y}_{i}-\hat{\mathbf{y}}_{i}\|^{2},$$

(5)

where $\mathbf{y}_{i}$ is the EFIT poloidal magnetic flux and $\hat{\mathbf{y}}_{i}$ is the NN output, and the summation is over all the samples in the training set. The loss function in Eq. (5) is the mean squared error (MSE). The loss function measures the derivation of the approximate solution away from the desired exact solution. So the goal of a learning algorithm is to find weights and biases that minimize the loss function. To minimize the loss function over $(\mathbf{w},\mathbf{b})$ using the gradient descent method, we need to compute the partial derivatives $\partial L/\partial w^{l}_{jk}$ and $\partial L/\partial b^{l}_{j}$, which can be efficiently computed by the well known back-propagating methodRumelhart1986 (15, 16). The back-propagating algorithm and the corresponding gradient descent method are the core algorithms in all deep learning software frameworks.

Besides the trainable parameters, there are various hyperparameters in a NN that usually need to be set manually, such as number of hidden layers, units per layer, activation function, NN optimizers, learning rate, batch size, number of epochs of training. In recent years, there appear automatic optimization libraries that can search for the best combination of hyperparameters. In this work, we use the Optuna optimization frameworkoptuna_2019 (14) in setting the hyperparameters. Optuna automates the hyperparameter optimization process by defining a search space of hyperparameters and exploring the space using efficient searching algorithms. The tree-structured Parzen estimator (TPE) algorithm is used in this work. This algorithm models the relationship between hyperparameters and their corresponding performance metrics and makes efficient decisions on which hyperparameters to try next.

Optuna automate the selection of the best hyperparameter combination. After multiple experiments, we have found that the model accuracy is not sensitive to the number of hidden layers, activation functions, and optimizers (an example showing the relative importance of these hyperparameters is given in Fig. 5). Therefore, these hyperparameters are fixed in the fine tuning step, in order to improve the speed of the model selection process, and explore more hyperparameter regimes to which the model may be sensitive. For other hyperparameters, we use Optuna framework to find the optimal combination of hyperparameters. Relative importance of these hyperparameters are shown in Fig. 6. The above results indicate that learning rate is the dominant factor that determines the model performance.

The final values of the hyperparameters used in the model are shown in Table 2.

The network are constructed and trained using Keras & TensorFlow2 chollet2015keras (17, 18), which is a broadly adopted open source deep learning framework in industry and research community. Figure. 7 plots loss function values as a function of the training epochs. The loss function is also evaluated on the validation set, which serves as a monitor for the possible overfitting. The validation loss follows the same trend as the training loss, indicating no overfitting.

Besides the $l_{i}$ discussed above, there are some other global parameters that can be constructed from the magnetic measurements, namely the plasma stored energy $W_{\operatorname{mhd}}$, normalized plasma beta $\beta_{N}$, and edge safety factor $q_{95}$. These parameters depend on information beyond the poloidal magnetic flux, namely the toroidal magnetic field and plasma pressure. Therefore they can not be fully determined by using only the poloidal magnetic flux predicted from the above network.  Following Ref. DIII-D_2022 (8), we construct a new NN for predicting these parameters (called NN2 in the following; the previous one will be called NN1), where the network has only 3 output values, namely $W_{\operatorname{mhd}}$, $\beta_{N}$, $q_{95}$. The input to NN2 includes a new signal, the current in the toroidal field (TF) coils, which determines the toroidal field. (In the NN1, this signal is not included because it has negligible effect on the prediction of the poloidal magnetic flux.) The NN2 has only one hidden layer consisting of 16 units, and uses the sigmoid as the activation function for both the hidden and output layers. The input and output signals of NN2 are normalized by using the same min-max scalar as used for NN1.

The training data consist of about $1/4$ randomly selected part of the data used for NN1. We found that using larger dataset makes this small network prone to overfitting. The testing set consists of 1000 time slices. Figure 20 plots the NN2 predictions against the EFIT values for the testing set. The results indicate that the NN2 predictions are in reasonable agreement with the EFIT values for all the 3 parameters. The NN2 predictions of $q_{95}$ are a little worse than those of the other two parameters, judging from the values of $r$ and $R^{2}$.

To evaluate the accuracy of NN2 prediction for a full discharge, we arbitrarily chosen a discharge and compare the time evolution of $W_{\operatorname{mhd}}$, $\beta_{N}$, and $q_{95}$ between the NN2 predictions and EFIT values. The results are shown in Fig. 21, which shows good agreement between the network predictions and EFIT values.

In this work, we train a multiply-layer neural network on the magnetic measurements (input) and EFIT poloidal magnetic flux (output) on EAST tokamak. The prediction capability of the network is examined by comparing the reconstructed magnetic surfaces, last closed flux surfaces, plasma current, and normalized internal inductance with those of EFIT. The neural network shows good agreement with EFIT for the data unseen in the training process.

In constructing the neural network, we use automatic optimization in searching for the best hyperparameters of the model. The hyperparameters found this way turn out to be better than our previously manually set hyperparameters in terms of the model accuracy.

Based on the model’s good prediction capability and efficiency in terms of computational time (about 0.5ms per equilibrium on a desktop computer, using an 11th Gen Intel(R) Core(TM) [email protected] CPU with a single thread), it looks promising to apply the neural network to real-time magnetic configuration control. The above computational time does not include the time used for tracing boundary/internal magnetic surfaces, and other related calculations to obtain $l_{i}$. These computations (not optimized in this work) seem too inefficient to be used in real-time control. The purpose of computing $l_{i}$ for the NN1 model is to evaluate the accuracy of the predicted $\Psi$. To predict these volume-integrated parameters, one usually uses an additional small network, as we did in Sec. V, which is efficient enough for real-time control because the network size is usually very small.

This work is limited to magnetic measurements. We plan to add more diagnostics relating to the inner safety factor profiles and pressure profiles into the model, in order to construct more realistic equilibria. This will rely on the kinetic EFIT output. We are accumulating these kind of training data.

The data that supports the findings of this study are available from the corresponding author upon reasonable request.

The authors thank Ting Lan, Tonghui Shi, Chengguang Wan, Zhengping Luo, Yao Huang, Guoqiang Li, and Jingping Qian for useful discussions. This work was supported by Comprehensive Research Facility for Fusion Technology Program of China under Contract No. 2018-000052-73-01-001228, by Users with Excellence Program of Hefei Science Center CAS under Grant No. 2021HSC-UE017, and by the National Natural Science Foundation of China under Grant No. 11575251.