3.1. 1D UEDGE model

UEDGE (Rognlien et al. Reference Rognlien, Brown, Campbell, Kaiser, Knoll, McHugh, Porter, Rensink and Smith1994; Rognlien & Rensink Reference Rognlien and Rensink2002) is a finite-volume simulation code to capture plasma and neutral transport. UEDGE has been used extensively to model the edge region of present-day and future tokamak devices for almost three decades. It has been utilised to study the divertor detachment physics on many machines, such as Alcator C-Mod (Wising et al. Reference Wising, Krasheninnikov, Sigmar, Knoll, Rognlien, LaBombard, Lipschultz and McCracken1997), DIII-D (Porter et al. Reference Porter, Allen, Brown, Fenstermacher, Hill, Jong, Leonard, Nilson, Rensink and Rognlien1996), ITER (Wising et al. Reference Wising, Knoll, Krasheninnikov, Rognlien and Sigmar1996) and SPARC (Ballinger et al. Reference Ballinger, Kuang, Umansky, Brunner, Canik, Greenwald, Lore, LaBombard, Terry and Wigram2021). UEDGE solves the multi-species fluid transport model for plasma and neutral gas, including important atomic physics, such as line radiation from ionisation, excitation and recombination processes in realistic tokamak equilibria. With the fully implicit Newton–Krylov scheme and adaptive stepping to advance all model equations in time, UEDGE is able to quickly evolve the system to a steady state.

In this proof-of-principle study, highly flexible and efficient 1D UEDGE code is used to quickly generate a large training data set that covers a wide range of parameter space beyond normal tokamak operation. Unlike the commonly used 2D UEDGE model, which simulates the entire tokamak periphery, including closed flux surface region, scrape-off-layer (SOL) and private flux surface region, the 1D UEDGE in this study models only one flux tube in the low-field side SOL, starting from the outer mid-plane to the outer divertor target plate, as depicted in figure 2(a). Nevertheless, 1D UEDGE code solves the same set of equations as its 2D counterpart except radial transport and drifts. Plasmas are allowed to transport in the parallel direction that is along the magnetic field whereas neutrals transport poloidally along the flux tube. As a result, the essence of plasma detachment physics is retained despite the one-flux-tube assumption. In fact, simplified 1D geometry disentangles many secondary effects and is sometimes preferred to better quantify the role of different physics processes played in the detachment (Dudson et al. Reference Dudson, Allen, Body, Chapman, Lau, Townley, Moulton, Harrison and Lipschultz2019). In the 1D UEDGE set-up, spatial variation of the magnetic field $\boldsymbol {B}$ is retained, i.e. the flux expansion along a flux tube in the SOL is considered in the 1D UEDGE mesh. Power, or the energy influx, across the core boundary is injected at the outer midplane (upstream), and is equally partitioned between electrons and ions. Plasma particles are also assumed to be fueled at upstream. As plasma moves downside the flux tube, it recombines to neutrals which eventually are recycled at the divertor target plate. In this study, the particle and momentum recycling coefficients are set to be $0.99$.

Figure 2. (a) 2D (black) versus 1D (red) UEDGE simulation meshes and (b) sketch of 1D UEDGE simulation set-up.

Figure 3 shows the typical attached (in blue) and detached (in orange) plasma in 1D UEDGE simulations. Compared with attached plasma, detached plasma normally has a lower electron temperature $T_e$, a higher neutral gas density $n_g$ at the divertor target and a distinct detachment front, or radiation peak away from the divertor target. Not surprisingly, 1D UEDGE code is able to qualitatively reproduce the target ion saturation current density $J_\parallel ^{\text{sat}}$ rollover (e.g. figure 2 in Loarte et al. Reference Loarte, Monk, Martin-Solis, Campbell, Chankin, Clement, Davies, Ehrenberg, Erents and Guo1998) and electron temperature cliff (e.g. figure 5(a) in McLean et al. Reference McLean, Leonard, Makowski, Groth, Allen, Boedo, Bray, Briesemeister, Carlstrom and Eldon2015) phenomena observed in detachment experiments as shown in figure 4. We admit that more sophisticated models such as 2D UEDGE and SOLPS-ITER considering additional details in principle yield more trustworthy results; but the fundamental physics picture of plasma detachment, that plasma radiates enough power to reach a low-enough temperature (a few electronvolts) and to recombine into neutral gas, will not change. Rather than developing a more comprehensive data-driven model with training data from state-of-the-art numerical simulations, here we focus on exploring the data-driven approach for detachment prediction with a somewhat simplified physics-based model (1D UEDGE) that could cover lots of plasma discharge scenarios relatively cheap computationally.

Figure 3. Examples of (a) plasma density, (b) electron temperature, (c) neutral density and (d) radiation profile of attached (blue) and detached (orange) divertor plasma from 1D UEDGE simulations.

Figure 4. Reproductions of (a) $J_\parallel ^{\text{sat}}$ rollover and (b) $T_e$ cliff features in 1D UEDGE density scan simulations.

3.2. Data generation

Although there are more than a dozen independent geometry and plasma parameters that can be adjusted in the 1D UEDGE model, here we only vary the three most relevant parameters, namely, upstream density $n_{e,u}$, injection power ${P_\text {inj}}$ and impurity fraction $f_Z$. These three parameters $(n_{e,u},P_\text {inj},f_z$) are hereafter referred to as ‘model inputs’.

In this study, we use DIII-D tokamak geometry with a fixed divertor leg length ${L_\text {leg}}=0.2696\,{\rm m}$, defined as the poloidal distance between the $X$-point to the outer divertor target as shown in figure 2(b). We uniformly sample $n_{e,u}\in [1,7]\times 10^{19}\,{\rm m}^{-3}$, ${P_\text {inj}} \in [1,10]$ MW and $f_Z\in 0\unicode{x2013}10\,\%$. Because the first wall material of DIII-D tokamak is graphite, here the impurity species is set to be carbon. With 60 sample points for $n_{e,u}$ and $f_Z$, and 40 sample points for ${P_\text {inj}}$, 144 000 cases in total are launched. However, not all of these cases can reach a physically feasible steady-state solution within a limited simulation time. Hence, only 111 598 converged simulation cases with good power balance relation are accepted as the training data set. Similarly, an independent validation data set with a coarser sampling rate ($50\times 50\times 30$ and 57 655 accepted solutions) is generated. Unless stated otherwise, the neural networks presented in this paper are trained and validated based on the training data set, only the combined forward prediction model is validated with the validation data set.

To avoid manual intervention of hundreds of thousands simulations individually, all the UEDGE runs are launched and managed by Merlin (Peterson et al. Reference Peterson, Bay, Koning, Robinson, Semler, White, Anirudh, Athey, Bremer and Di Natale2022), a ML workflow framework designed for large-scale high-performance computing (HPC) environments. It takes 4 days with 200 CPUs to obtain about 170 000 converged cases. The data-generation time could be further shortened if more CPUs were utilised.

For the diagnostic description of plasma state in this 1D set-up, five synthetic measurements are taken across the simulation domain, including the upstream electron temperature $T_{e,u}$, ion saturation current density $J_\parallel ^{\text{sat}}$, electron density $n_{e,t}$ and temperature $T_{e,t}$ at divertor target, as well as the radiation profile $P_{{\rm rad}}$. As illustrated in figure 2(b), $T_{e,u}$ mimics the mid-plane Thompson scattering measurement, $J_\parallel ^{\text{sat}}$ mimics the Langmuir probe measurement, $n_{e,t}$ and $T_{e,t}$ come from either divertor Thompson scattering or Langmuir probe and $P_{{\rm rad}}$ is deduced from either bolometer or C-III emission. Note here $J_\parallel ^{\text{sat}}$, $T_{e,u}$, $n_{e,t}$ and $T_{e,t}$ are scalars, whereas $P_{{\rm rad}}$ is a 1D profile. These diagnostic measurements are collected for all 111 598 cases and then normalised for model training.

Prior model training, all cases are labelled as either ‘detached’ or ‘attached’ because there is no ‘partial detached’ state in the 1D system. Here the choice of detachment criterion is straightforward. As shown in figure 5, electron temperature at the divertor target $T_{e,t}$ has two distinct distributions around $T_{e,t}=2.1\,{\rm eV}$. The gap at 2.1 eV is caused by the temperature cliff phenomenon shown in figure 4(b). Therefore, 56 047 cases with $T_{e,t}<2.1\,{\rm eV}$ are labelled ‘detached’, whereas the remaining 55 551 cases are labelled as ‘attached’ cases. As shown in table 1, nearly half of the cases in the validation data set are labelled ‘detached’, similar to the percentage of ‘detached’ cases in the training set. In DIII-D experiment, the temperature cliff occurs around 5 eV, this apparent discrepancy is mainly due to our simplified 1D set-up. The temperature cliff in our simulations is consistent with the sonic transition point where the ion parallel velocity reaches sonic moving away from the target. A thorough investigation of temperature cliff in one dimension is beyond the scope of this paper and will be presented in a separated paper by Zhao et al. (Reference Zhao, Rognlien, Zhu, Meyer, Xu, Dudson and Li2022). We would like to point out that labelling cases is primarily to better quantify the models’ accuracy in § 5; it does not affect the training process at all. In addition, the lower $T_{e,t}\sim 0.22\,{\rm eV}$ limit shown in figure 5 is likely related to the atomic physics. When $T_e<1\,{\rm eV}$, the ionisation rate decreases whereas the recombination rate increases dramatically. In this temperature range, electron collisionality is very high such that three-body recombination becomes dominant. This process moves plasma potential energy to electron thermal energy, stopping $T_e$ decreasing further.

Figure 5. Histogram and cumulative distribution of electron temperature at divertor target $T_{e,t}$ from the training data set (111 598 cases in total).

Table 1. Numbers of detached and attached cases and their percentages in the training and validation data sets.

4.1. Latent space identification using AEs

The basic idea behind latent space identification (also referred to as dimensionality reduction or data compression) is straightforward. Given an input data point, the goal of an AE is to reduce it to fewer dimensions such that the reduced representation allows the original input to be reconstructed as faithfully as possible. An AE comprises two parts: (1) an encoder that learns to meaningfully encode the data into a set of values, a compact representation of the input (i.e. latent coordinates); and (2) a decoder that attempts to reconstruct the original inputs from this latent space encoding. By forcing the reconstruction of data, an AE attempts to learn a suitable latent space that captures all necessary degrees of freedom while discarding trivial variations, noise and redundant correlations in the data.

In this work, we utilise two types of AE to create a suitable LSR in two steps. First, we develop a (vanilla) AE to identify an appropriate level of reduction, i.e. the dimensionality of the latent space and a suitable neural network architecture. Next, we turn this AE into a special type of AE, called the $\beta$-VAE, to create the final LSR.

4.1.1. AE design to identify the dimensionality of a suitable latent space

Given a set of diagnostic measurements ($\boldsymbol {y}$), we develop an AE to identify the compressed LSR ($\boldsymbol {z}$) of the plasma state. We use $\boldsymbol {y}'$ to denote the reconstructed input from the AE. In other words, the encoder maps $\boldsymbol {y} \to \boldsymbol {z}$ and the decoder maps $\boldsymbol {z} \to \boldsymbol {y}'$. In our case, $\boldsymbol {y}$ (and, hence, also $\boldsymbol {y}'$) has 34 values: it comprises 4 scalar values ($J_\parallel ^{\text{sat}}$, $T_{e,u}$, $N_{t}$ and $T_{e,t}$) and a 30-element 1D array ($P_\text {rad}$). However, the different values (scalars versus the array) exhibit different ranges of values. Therefore, we utilise a scaling factor $\alpha _i$ for each of the 34 input dimensions to prevent overfitting on the 30 values for the array. In particular, we choose $\alpha _i = 0.1$ for values corresponding to the array elements and $\alpha _i = 1.0$ for the scalars, effectively asking the model to consider each array element 10$\times$ less important than the scalars. Formally, our AE is trained to minimise the loss function

(4.1)\begin{equation} \mathcal{L}(\boldsymbol{y},\boldsymbol{y}') = \frac{1}{N}\sum_i \alpha_i|y_i - y'_i|^{n}, \end{equation}

where $N=34$ is the total number of elements in the input diagnostic measurement data $\boldsymbol {y}$. For $n=1$, the loss function is the mean absolute error (MAE) or $L1$ norm, and for $n=2$, it becomes the mean square error (MSE) or $L2$ norm.

As our input data are simply a collection of values, our AE design uses a series of fully connected neural network layers. Each such layer provides dense connections between input and output values through a linear transformation (a matrix multiplication with an additive bias) followed by a nonlinear activation. Given a chosen input/output dimensionality and an activation function, the training of the model then learns the appropriate matrix and bias to perform this transformation. We experimented with different depths of the neural network (i.e. the number of such fully connected layers) and the corresponding dimensionalities. All models were trained using the Tensorflow (Abadi et al. Reference Abadi, Agarwal, Barham, Brevdo, Chen, Citro, Corrado, Davis, Dean and Devin2015) and Keras (Chollet et al. Reference Chollet and Others2015) frameworks with the Adam optimiser (Kingma & Ba Reference Kingma and Ba2014). A dataset of 111 598 samples was used for developing the AE, with a randomly selected 80 % subset of these data for training, and the remaining 20 % for validating the model (i.e. assess its accuracy).

The final AE model, which provided the best reconstruction accuracy (see table 2), contains three fully connected layers that progressively reduce the data dimensionality as $34 \to 18 \to 10 \to 6$. Each layer is followed by the sigmoid activation function. The output of the last layer is the desired LSR; in other words, a 6-dimensional (6D) latent space is adequate to compress a 34-element diagnostic data in our case so as to recover the input with sufficient accuracy.

Table 2. To choose a suitable dimensionality, $D_{\boldsymbol {z}}$, of the latent space, we study the reconstruction errors posed by different dimensionalities. Given the significant jump in the error when going from six to five, $D_{\boldsymbol {z}}$ was found to be the appropriate choice.

4.1.2. VAE design to identify a suitable latent space

The AE set-up described here is capable of generating compressed encoding of the input data. However, such AEs provide no control of the distribution of the data in the learned LSR. Typically, this is not a big issue if the AE is used to encode/decide the data, i.e. use the AE (both encoder and decoder) as one model. However, it is not suitable for our application as our goal is to utilise only the decoder and in conjunction with a forward, surrogate model. Specifically, in our case it is important to provide smoothness guarantees on the LSR, otherwise even small prediction errors from the forward model could be amplified significantly by the decoder. To regularise the distribution in the latent space and to improve decoder's generative capability, we use a specific type of AE, called $\beta$-VAE (Higgins et al. Reference Higgins, Matthey, Pal, Burgess, Glorot, Botvinick, Mohamed and Lerchner2016). The key difference between the (vanilla) AE and the VAE is that the VAE forces the distribution of $\boldsymbol {z}$ to a given reference distribution, typically a multivariate normal distribution, offering some important mathematical guarantees, such as smoothness. To ensure that $\boldsymbol {z}$ mimics the reference distribution, a regularisation term is added to the loss function to measure how different the two distributions are. Specifically, the loss function from (4.1) is modified to

(4.2)\begin{equation} \mathcal{L}(\boldsymbol{y},\boldsymbol{y}') = \frac{1}{N}\sum_i \alpha_i|y_i - y'_i|^{n} + \beta D_{{\rm KL}}(f(\boldsymbol{z})\,||\,f_0), \end{equation}

where $D_{{\rm KL}}(f(\boldsymbol {z})\,||\,f_0)$ is the Kullback–Leibler divergence, which measures how the distribution $f(\boldsymbol {z})$ differs from the reference distribution $f_0$, and a hyperparameter $\beta$ is used to balance the decoder reconstruction accuracy (the first term in the loss function) and the orthogonality of latent space coordinates (the second term). As is common for VAEs, our reference distribution $f_0$ is a multivariate normal distribution, $\mathcal {N}(0,1)$. In many applications, $\beta$ is set to be larger than unity for better separation of independent latent space coordinate; in our case, forcing the LSR distribution function to match the normal distribution is not necessary. Instead, a small $\beta =10^{-9}$ is found to be suitable. The VAE used in this study has the same network architecture and training parameters as the AE described previously.

4.2. Predicting simulation outputs using a surrogate model

Given the LSR, $\boldsymbol {z}$, generated by the VAE described previously, we next develop another ML model that can predict $\boldsymbol {z}$ based on the physics model inputs, $\boldsymbol {x}$ (i.e. global discharge parameters). Due to a small numbers of model inputs (i.e. three varying parameters in 1D UEDGE model) and outputs (i.e. 6D latent space), here a simple yet robust MLP is chosen to be our forward model. This MLP has four fully connected layers with 24 neurons on each layer, uses the rectified linear (ReLu) activation function, and MSE as the loss function. An Adam optimiser with a learning rate $0.001$ is used for training the model.

5.1. Latent space identification using AEs

Through the AE described previously, we identify the dimensionality of a suitable latent space, $D_{\boldsymbol {z}}$, to be six. Table 2 indicates that with both $L1$ or $L2$ norm, the validation loss (in normalised units) increases as $D_{\boldsymbol {z}}$ decreases. The $L2$ norm appears to have a critical turning point at $D_{\boldsymbol {z}}=6$; the validation loss increases substantially when $D_{\boldsymbol {z}}$ is set to be less than six, whereas for $D_{\boldsymbol {z}}>6$, the improvement is less profound.

Next, we study the performance of the VAE using figure 7 with UEDGE data $f_\text {UEDGE}$ on the horizontal axis and the residual between VAE-reproduced data and UEDGE data $\varepsilon =f_{\beta \text {-VAE}}-f_\text {UEDGE}$ on the vertical axis. If a perfect reconstruction can be achieved, all data points would lie exactly on the horizontal $\varepsilon =0$ lines. In our case, two of the four scalar diagnostic measurements, the electron density at divertor $n_{e,t}$ and the electron upstream temperature $T_{e,u}$ can be reproduced with excellent accuracy, whereas the other two scalar measurements, the ion saturation current, $J_\parallel ^{\text{sat}}$, and the electron temperature at divertor, $T_{e,t}$, also show good quality reconstruction, with a slight performance degradation at high values of the respective diagnostics. Nevertheless, the $R^{2}$ scores (or, coefficient of determination that measures the goodness-of-fit between the two variables $f=\{f_i\},g=\{g_i\}$ defined as $R^{2}=1-\sum _i(f_i-g_i)^{2}/\sum _i f_i^{2}$ with $i$ being the index of the variable; hence, $R^{2}=1$ indicates perfect correlation whereas $R^{2}=0$ means no correlation) for these four scalars all exceed $0.997$, evidencing a close to perfect replication of input data from trained $\beta$-VAE statistical-wise. To quantify the reconstruction quality of the radiation profile, $P_\text {rad}$, we choose two metrics: peak amplitude and location. The peak radiation amplitude prediction by the VAE is consistent the UEDGE data with a slightly larger variance ($R^{2}\simeq 0.994<0.997$). It appears to be challenging to reproduce the peak radiation or detachment front location with $R^{2}_\text {detached}\approx 0.7$ and $R^{2}_\text {attached}\approx 0$. However, due to the discrete nature of the simulation mesh, radiation peak locations are discretised as well so that the $R^{2}$ score here could be misleading, especially for the attached cases. In fact, all 55 551 attached cases whose peaking radiation location at the divertor target are correctly reproduced by $\beta$-VAE but they are represented by what appears to be ‘a single’ orange dot in figure 7( f).

Figure 7. Performance of $\beta$-VAE in terms of absolute residual error for (a) $J_\parallel ^{\text{sat}}$, (b) $n_{e,t}$, (c) $T_{e,u}$, (d) $T_{e,t}$, (e) $P_\text {rad}^{\text{max}}$ and ( f) peak radiation or detachment front location. Here, the $x$-axis is UEDGE (true) value and the residual error is defined as $\varepsilon =f_{\beta \text {-VAE}}-f_\text {UEDGE}$ for the quantity $f$. All predictions show excellent correlation with the true values, as quantified with almost-perfect $R^{2}$ scores. The distribution of data points are shown in adjoining panels to overcome the issue of overplotting in the scatter plots; these indicate that for most cases, relatively low and unbiased error is obtained, whereas higher errors are associated with relatively smaller number of samples.

One plasma quantity that we are especially interested in for detachment prediction is the electron temperature at the divertor target, $T_{e,t}$, as this is arguably the most important and direct indicator of detachment. As shown in figure 8(a), VAE is able to reproduce $T_{e,t}$ near the detachment onset value (${\sim }2.1\,{\rm eV}$ for this study) very well. Figure 8(b) shows that the VAE is able to also accurately retain the ion exit velocity, $v_{\parallel i,e}=J_\parallel ^{\text{sat}}/(e n_{e,t})$, which is an inherent quantity that was not trained directly. A discrepancy between the VAE's predictions and UEDGE results appears for $J_\parallel ^{\text{sat}}/(en_{e,t}) > 55\,{\rm km}\,{\rm s}^{-1}$, which corresponds to attached plasma cases with $T_{e,t}\gg 2.1\,{\rm eV}$. However, this type of discrepancy exists only for less than $2\,\%$ of the total cases: the small number of $T_{e,t}\gg 2.1\,{\rm eV}$ samples may cause the relatively large prediction error for these cases. This result indicates that our trained VAE captures the correct physics constraints between different plasma variables (i.e. AE inputs and outputs) for at least the majority of the cases in our dataset.

Figure 8. Residual plots for (a) $T_{e,t}$ near detachment transition and (b) $J_\parallel ^{\text{sat}}/n_{e,t}$. Detached cases are in blue whereas attached cases are in orange. Note here that the $y$-axis of the distribution plots is in logarithmic scale.

Finally, we show that $\beta$-VAE encoded LSRs are indeed apart in latent space based on plasma states (e.g. detached or attached). Figure 9 illustrates the distinctive clusters of LSRs for detached and attached cases in latent space using the $t$-distributed stochastic neighbour embedding ($t$-SNE) method (Van der Maaten & Hinton Reference Van der Maaten and Hinton2008). Similarly, figure 10 has the distributions of all six latent coordinates, again separated by attachment/detachment label. The result shows that despite overlaps between detached and attached plasmas in some latent variables (e.g. latent variables 1, 4 and 5), the two scenarios are broadly well separated in the other latent variables (e.g. latent variables 2 and 6), which is a nice feature to have for the upcoming forward detachment prediction.

Figure 9. $t$-SNE visualisation of a subset of 10 000 randomly picked LSR samples in the training data set in which detached (blue) and attached (orange) cases are clearly separated.

Figure 10. Performance of the MLP model in terms of absolute residual error for all six latent variables. Here the $x$-axis is $\beta$-VAE encoded value, whereas the residual is defined as $\varepsilon = f_\text {MLP}-f_{\beta \text {-VAE}}$ for quantity $f$. As with earlier results, an almost-perfect $R^{2}$ score indicates exceptional prediction quality.

5.2. Predicting simulation outputs using a surrogate model

Here, we study the performance of our surrogate model, an MLP, after 10 000 epochs. In particular, we compare the LSR generated by the VAE, $\boldsymbol {z}$, and that predicted by the surrogate model, $\boldsymbol {z_f}$. Figure 10 depicts plots the six coordinates within the two quantities ($\boldsymbol {z}$ vs. $\boldsymbol {z_f}$) as individual scatter plots and shows high degree of correlation. As shown in the figure, the $R^{2}$ values for all six coordinates exceeds $0.99$, indicating exceptional performance and that our surrogate model is capable of predicting the LSR with almost full accuracy.

5.3. Forward prediction

With both $\beta$-VAE and forward model trained, we are now able to predict diagnostic measurement of a plasma state from the model inputs by combining the forward model and decoder (figure 6c). To ensure that this model is evaluated properly, a separated validation data set consists 57 655 cases generated independently following the similar data-generation process described in § 3.2.

The performance of forward detachment prediction model is illustrated in figure 11 where all 57 655 cases are evaluated with the ML model then validated with the UEDGE simulation results. As expected, the overall accuracy degrades marginally compared with the performance of $\beta$-VAE (e.g. figure 7) in this sequential MLP and decoder architecture.

Figure 11. Performance of combined forward prediction model in terms of residuals for (a) $J_\parallel ^{\text{sat}}$, (b) $n_{e,t}$, (c) $T_{e,u}$, (d) $T_{e,t}$, (e) $P_\text {rad}^{\text{max}}$ and ( f) peak radiation or detachment front location. Here the $x$-axis is UEDGE (true) value, whereas the residual is defined as $f_\text {residual}=f_{{\rm ML}}-f_\text {UEDGE}$ for quantity $f$. The distribution of $f$ and the $R^{2}$ scores between $f_{{\rm ML}}$ and $f_\text {UEDGE}$ for detached (blue) and attached (orange) cases are also provided.

Figure 12 and table 3 elucidate the accuracy of this combined forward detachment prediction model, quantified with error statistic analysis. We examine model accuracy for both detached and attached cases, and find that our model does perform slightly different for different plasma states. For attached plasma, it gives better results when predicting certain plasma quantities, such as ion saturation current and divertor target electron density, whereas it is less satisfactory for divertor target electron temperature prediction. However, no significant quantitative difference is found except the detachment front location prediction. Out of four scalar measurements, upstream temperature $T_{e,u}$ is the most accurately predicted diagnostic quantity with mean and standard derivation of relative error $\mu =-0.37\,\%$ and $\sigma =0.83$ for the entire validation data set. The other three scalar measurements are also well predicted with $|\mu |<3\,\%$ and $\sigma <4\,\%$ for all the relative errors. The peak radiation strength prediction appears to have a fairly large uncertainty (${\sim }10\,\%$). This is possibly due to the shape peak structure (i.e. lack of resolution) for detached plasmas and boundary effect (i.e. maximum value resides at the last point) for attached plasmas as shown in figure 13(b). Fortunately, the more useful information, the peak radiation location (equivalent to the detachment front location $L_{{\rm DF}}$ to some extent) prediction is remarkably accurate. Here $L_{{\rm DF}}$ is correctly predicted to be at the divertor target plate for nearly all attached plasmas, and has a small uncertainty ($\sigma =0.60\,{\rm cm}$) for detached plasmas. Note that the error distributions of four scalar measurements appear to follow Gaussian distribution, whereas for the peak radiation amplitude and location, the error distribution are no longer Gaussian : the standard derivation $\sigma$ is skewed by rare (i.e. probability ${\sim }0.1\,\%$) extreme cases.

Figure 12. Probability distribution functions (PDFs) of relative error of (a) $J_\parallel ^{\text{sat}}$, (b) $n_{e,t}$, (c) $T_{e,u}$, (d) $T_{e,t}$, (e) $P_\text {rad}^{\text{max}}$ and ( f) absolution error of peak radiation or detachment front location. Detached cases are in blue whereas attached cases are in orange. Dashed lines are shifted Gaussian distribution functions with $\mu$ and $\sigma$ from table 3 correspondingly.

Figure 13. (a) Electron temperature at divertor target $T_{e,t}$ prediction versus UEDGE simulation result and (b) examples of radiation profile or detachment front prediction. True (UEDGE) detached cases are in blue whereas attached cases are in orange.

Table 3. Mean ($\mu$) and standard deviation ($\sigma$) of the relative or absolute error of forward prediction model.

We are particularly interested in $T_{e,t}$ prediction as this is the primary indicator of detachment. If we use the same detachment criterion (e.g. $T_{e,t}=2.1\,{\rm eV}$), there are 35 out of 28 468 detached cases (i.e. blue triangles at the top left quadrant of figure 13a) misclassified as ‘attached’. Similarly, 25 out of 29 187 attached cases (orange diamonds at the bottom right quadrant of figure 13a) are mislabelled ‘detached’. Even counting in the other marginal cases, the misclassification rate is lower than $0.2\,\%$. These misclassified cases appears to congregate near the detachment onset point (i.e. all these cases have $T_{e,t}\in (1,4)\,{\rm eV}$). However, no common feature can be identified in terms of the control parameters $\boldsymbol {x}$.

In addition to accuracy, speed is another important metric for a real-world application. Table 4 summarises our predictive model's performance in terms of the wall-clock time required to predict diagnostic measurements based on controlled inputs. Even without any optimisation, this model takes about 36 ms to carry out a prediction for one case: about 10 000 times faster than 1D UEDGE simulations which normally require a few minutes to find a converged solution. It is already within the minimum requirement for detachment control (${\sim }100\,{\rm ms}$). The model efficiency also increases when predicting multiple cases (for $n<100\,000$).

Table 4. Wall-clock time versus number of cases.

6.1. Basic 2PM

The basic 2PM is derived to evaluate $T_{e,u}$, $T_{e,t}$ and $n_{e,t}$ for given $n_{e,u}$, parallel heat flux $q_\parallel$ and flux tube length $L$ based on particle, pressure and power balance. Because of its simplicity, basic 2PM has been implemented in the tokamak detachment control a priori to estimate the degree of detachment (Eldon et al. Reference Eldon, Anand, Bak, Barr, Hahn, Jeong, Kim, Lee, Leonard and Sammuli2022). Because 2PM does not account for momentum and power loss between the upstream and downstream points, we therefore benchmark UEDGE, the forward detachment model and the 2PM on a case with zero impurity to remove the impurity radiation. Figure 14 displays upstream density scan of $J_\parallel ^{\text{sat}}$ and $T_{e,t}$ for the three models. Clearly, the data-driven model prediction and the UEDGE simulation results are in good agreement whereas the 2PM gives a quite different result. This one-order-of-magnitude discrepancy between our data-driven model/UEDGE and 2PM is due to the more comprehensive physics in our model/UEDGE. For instance, in this comparison the upstream ion temperature is much higher than the electron temperature in the upstream region in the UEDGE simulations due to slow ion parallel thermal conduction whereas the electron and ion temperatures are assumed to be the same in the basic 2PM. In addition, in the basic 2PM, the parallel heat flux is assumed to be conductive using the Spitzer–Härm formula, whereas in 1D UEDGE simulations, the parallel heat flux is assumed to be local flux-limited thermal transport. UEDGE is only able to recover 2PM result by further simplifying simulation (e.g. using slab configuration and turn off non-ideal terms) and manually adjusting boundary conditions to match 2PM assumptions, such as $T_e=T_i$ upstream.

Figure 14. Upstream density scan of (a) $J_\parallel ^{\text{sat}}$ and (b) $T_{e,u}$ for UEDGE (blue), ML data-driven model (orange) and basic 2PM (black dashed) with $P_\text {inj}=1.71\,{\rm MW}$ and $f_Z=0$.

6.2. 2PMF

To address the volumetric momentum and power losses, and the magnetic geometry effects, the basic 2PM has been extended to semi-analytical 2PMF with pre-fitted power and momentum loss coefficients $f_\text {cooling}$, $f_\text {mom-loss}$ (Stangeby Reference Stangeby2018). Benchmark between our data-driven model, UEDGE simulation and 2PMF are also performed. Two fitting curves (from Stangeby Reference Stangeby2018) are proposed for both power and momentum loss coefficients in 2PMF:

Stangeby fitting formula 1,

(6.1)\begin{gather} 1-f_\text{cooling} = [1-\exp({-T_{e,t}/2.4})]^{1.9} \end{gather}

(6.2)\begin{gather}1-f_\text{mom-loss} = [1-\exp({-T_{e,t}/0.8})]^{2.1}; \end{gather}

Stangeby fitting formula 2,

(6.3)\begin{gather} 1-f_\text{cooling} = 0.9[1-\exp({-T_{e,t}/6})]^{1.7} \end{gather}

(6.4)\begin{gather}1-f_\text{mom-loss} = 1.3[1-\exp({-T_{e,t}/1.8})]^{1.6}; \end{gather}

where UEDGE observed $T_{e,t}$ is used to estimate $f_\text {cooling}$ and $f_\text {mom-loss}$. Figure 15 displays upstream density scan of $J_\parallel ^{\text{sat}}$ and $T_{e,t}$ for the three models. The results from 2PMF compare with UEDGE simulation results better than those from the basic 2PM. All three models show two common features: (1) the ion saturation current density rollover and (2) temperature cliff at the onset of detachment. Once again, our data-driven model prediction matches UEDGE simulation results very well. Although there is no qualitative disagreement between our data-driven model/UEDGE and 2PMF, quantitatively the predictions between data-driven model/UEDGE and 2PMF can be off by an order of magnitude. This is because 2PMF prediction depends heavily on the choice of fitting curves for the power and momentum loss coefficients which may simplify the nonlinear dynamics setting the divertor plasma conditions. As illustrated in figure 15, fitting formula 1 does a better overall job than fitting formula 2 for this test case as $T_{e,t}$ predicted by fitting formula 2 is nearly an order of magnitude lower than UEDGE result when divertor plasma is detached ($n_{e,u}>2.4\times 10^{19}\,{\rm m}^{-3}$).

Figure 15. Upstream density scan of (a) $J_\parallel ^{\text{sat}}$ and (b) $T_{e,u}$ for UEDGE (blue), ML data-driven model (orange) and 2PMF (black lines) with $P_\text {inj}=2.66\,{\rm MW}$ and $f_Z=0.02$.