EuroPED-NN: Uncertainty aware surrogate model

The EuroPED model is conceptually identical to the EPED model, which uses an ideal MHD stability description to predict the critical plasma pressure gradient, $\alpha$, as a function of the pedestal width, $\Delta$. This curve is then crossed against a semi-empirical $\Delta$-$\beta_{p,ped}$ relation reflecting the KBM constraint [14] to provide a prediction of the pedestal properties. The assumption of a two-species plasma in equilibrium is then used to estimate the top-pedestal electron pressure, $p_{e,ped}$, and further decompose that into the top-pedestal electron temperature, $T_{e,ped}$, by requiring the top-pedestal electron density, $n_{e,ped}$, as input.

Thus, for the EPED model, the required input parameters include:

• •

  plasma parameters:

  • –

    plasma current, $I_{p}$, in [MA];

  • –

    toroidal magnetic field, $B_{t}$, in [T];

  • –

    line-integrated effective charge, $Z_{\text{eff}}$;

  • –

    top-pedestal electron density, $n_{e,ped}$ , in [$\times 10^{19}$ m^-3];

  • –

    normalized plasma pressure, $\beta_{N}$;

• •

  and magnetic geometry parameters:

  • –

    triangularity, $\delta$;

  • –

    elongation, $\kappa$;

  • –

    minor radius, $a$, in [m];

  • –

    major radius, $R_{0}$, in [m].

In practice, one drawback of the EPED model is the necessity of providing $n_{e,ped}$, $\beta_{N}$ and $Z_{eff}$ as inputs, as this precludes the model from predictive exercises as these parameters are often unknown before performing an actual experiment.

To remedy this, the EuroPED model was developed by coupling an EPED-like pedestal parameter calculation workflow to the Bohm/gyro-Bohm (BgB) turbulent transport model [32] inside a 1D transport solver. This removes the need to explicitly define $\beta_{N}$ and $n_{e,ped}$, as the former is provided by the kinetic profiles resulting from the self-consistent transport calculation and the latter is defined either via a regression or a neutral penetration model provided within EuroPED. However, this model additionally requires two engineering parameters as input:

• •

  injected auxiliary heating power, $P_{tot}$;

• •

  separatrix electron density, $n_{e,sep}$.

This is seen as an improvement over the EPED model as $\beta_{N}$ or $n_{e,ped}$ are difficult to estimate prior to conducting an experiment. Although often the experiments are run in feedback mode where these parameters can actually be chosen. However, from the model point-of-view EPED is still not fully predictive. While the $n_{e,sep}$ is equally difficult to know prior to an experiment, the option to estimate it via the fuelling gas flow rate, $R_{\text{gas}}$, and a neutral penetration model is provided within EuroPED. This connection was excluded from this study as it was not used in the construction of the particular database forming the training set.

The EuroPED model then returns as outputs:

• •

  the top-pedestal electron density, $n_{e,ped}$;

• •

  the top-pedestal electron temperature, $T_{e,ped}$;

• •

  and pedestal width, $\Delta$.

However, due to the implementation details connecting the EPED-like workflow to the transport model, the actual radial locations of the top-pedestal quantities may not coincide in practice. This has some implications on the construction of an appropriate surrogate model, which is further discussed in Section 2.4.

As mentioned earlier, Bayesian neural networks are a specific class of NNs trained using Bayesian inference methods, involving an estimation of a posterior distribution over network weights given observed data. One of the difficulties with training BNNs is that they can be computationally expensive, particularly for large networks with many parameters. By providing a fast and efficient approximation to the posterior distribution, the Noise Contrastive Prior technique can speed up BNN training.

Within this technique, the posterior distribution is approximated by a mixture of the prior distribution over the output and a noise distribution. The network is trained to distinguish between observed data and noise samples generated by the noise distribution during training. Essentially, it uses simulated noise to shift the hypothesis from the probability distribution of the neural network weights to a probability distribution in the data space, requiring less computational resources. In practice, this means that a random point from the joint input distribution is sampled per training epoch. This sampled point is passed forward through the current BNN configuration, and the resulting mean and epistemic uncertainty is compared to a user-defined prior output distribution via some distance metric. This significantly reduces the number of samples required per epoch compared to a typical Monte Carlo method by leveraging the large number of epochs required for training the model to effectively accumulate the necessary statistics, which is similar to the training advantages offered by mini-batching and stochastic gradient descent. Furthermore, the NCP method improves the detection of out-of-distribution (OOD) regions in data space by directly accounting for input variations within its epistemic uncertainty estimation. In summary, this method allows for an approximate calculation of the output distribution without the need to sample the weights to arrive at a reliable estimate of the uncertainty. This approach is presented in [29] for one input and one output problems, in this study a multidimensional approach to fit the needs of the plasma pedestal models is built.

As a brief overview, the BNN-NCP architecture consists of a number of fully connected feed-forward (i.e. dense) layers with a custom layer inserted just before the output layer, as shown in Figure 3. This custom layer consists of a variational layer and a dense layer with a “softplus" activation function, as shown in Figure 3, each with only a single neuron. This variational layer represents its weights and biases as joint probability distribution, also known as a kernel, which effectively allows them to be sampled to assemble its predictive distribution, $p\!\left(\mathbf{y}^{*}|\mathbf{x}^{*},\mathbf{x},\mathbf{y},\theta\right)$, where the ^∗ notation is used to differentiate a generic point from one which is explicitly within the training dataset. The “softplus" activation function on the dense layer enforces its output to be positive, allowing it to represent an independent uncertainty estimate. All other neurons in the network have a “leaky ReLU" activation function, which is effectively a piecewise function of two linear components, both with non-zero gradients and joined together at the origin.

Due to the extension of this technique to multiple output variables, the final dense layer before the custom layer is separated into dedicated blocks per output variable. This serves as an attempt to ease the training process by reducing the amount of cross-talk between the network predictions of each variable. It also introduces a degree of flexibility in the last hidden layer, allowing adjustments to suit the complexity of each variable, e.g. the most difficult variables to fit can have a large number of neurons to account for their dependencies. This results in a multi-input multi-output model, transforming input points, $\mathbf{x}=\left\{x_{i}\right\}$, to output points, $\mathbf{y}=\left\{y_{j}\right\}$, which can then be trained on a labelled training dataset, denoted as $D=\left\{\left(\mathbf{x},\mathbf{y}\right)_{k}\right\}$.

As a result of the BNN-NCP architecture, there are two modes of evaluating it on a given input point, $\mathbf{x}$:

1. 1.

   By sampling the probability distribution represented by the kernel in the variational layer, resulting in a stochastic output, $\mathbf{\hat{y}}$;

2. 2.

   By computing the joint statistical moments of the kernel in the variational layer, resulting in a deterministic mean output, $\mathbf{\mu}$, and its standard deviation, $\mathbf{\sigma}$.

Both of these evaluation modes return another deterministic output from the softplus neuron, $\mathbf{s}$. The predictive output value of the BNN-NCP model is taken as the mean value, i.e. $\mathbf{y}^{*}\equiv\mathbf{\mu}$, from the second evaluation method due to its deterministic nature, although the first evaluation method is still described here due to its importance for the training methodology. Thus, the entire BNN-NCP model can be mathematically represented as a black box function, $\left\{\mathbf{\mu},\mathbf{\sigma},\mathbf{s}\right\}^{*}=f\!\left(\mathbf{x}^{*}\right)$.

Then, based on these descriptions, the output standard deviation from taking the statistical moment of the kernel, $\mathbf{\sigma}$, was labelled as the epistemic uncertainty prediction and the output of the independent softplus neuron, $\mathbf{s}$, as the aleatoric uncertainty prediction. Both uncertainty outputs are taken to represent a value of one standard deviation when converting them from statistical moments back into probability distributions.

The cost function used in this study is similar to the one introduced in Ref. [29] except extended into the multidimensional approach. This results in expression provided in Equation (1). The terms are then summed over each output variable, represented by the subscript $j$.

where $\gamma$ represents a user-defined weight coefficient used to adjust the associated loss term, $\theta$ represents the distribution of weights in the BNN, and the $\sim$ notation is used to indicate quantities related to the OOD region sampling. The sampled OOD input values, $\mathbf{\tilde{x}}$, were generated from a normal distribution about each dataset point, $\mathbf{x}$, according to a user-defined input noise term, $\mathbf{\Sigma_{x}}$. This input noise should generally be chosen as a large enough value to ensure the OOD sampling goes adequately beyond the training dataset input values, $\mathbf{x}$, but not so large that the OOD sampling would have difficulty resolving well the area just outside the training dataset.

The first term of the loss function corresponds to the negative log-likelihood (NLL) which represents a goodness-of-fit metric for the approximate predicted posterior distribution, selected to be:

By minimizing this term with respect to the weights and biases, $\theta$, the maximum likelihood value is obtained for the network prediction. As the aleatoric uncertainty prediction, $s_{j}$, is used as the weighting parameter for this metric, it also encourages the aleatoric uncertainty to grow where the network cannot match the specific data point where the combination of other data points and regularization discourage it, e.g. where data noise is present.

The second term represents a metric to anchor and regularize the predicted epistemic uncertainty with weight coefficients, $\gamma_{\text{epi},j}$, which are used to adjust its relative importance in the loss function per output variable. The last term represents the same anchoring and regularization metric for the predicted aleatoric uncertainty, with its own relative weight coefficients per output variable, $\gamma_{\text{alea},j}$.

Both uncertainty regularization terms use the Kullback-Leibler (KL) divergence metric, which effectively compares the similarity of two distributions. For the epistemic term, the approximate output distribution predicted by the network using an OOD input distribution, $q\!\left(\mu_{j}|\mathbf{\tilde{x}},\theta\right)$, was compared against an epistemic prior distribution, $p_{\text{prior,epi},j}$, selected to be:

This represents a normal distribution where the mean is the corresponding output value in the dataset, $y_{j}$, and the standard deviation, $\sigma_{y_{j}}$, is a user-defined parameter encapsulating any expected epistemic uncertainty already known to be associated with the data-generating system. This term is expected to generate a significant loss for OOD samples far away from the training point, effectively encouraging the epistemic uncertainty at those OOD samples to grow in order to minimize the overall loss. In order to save computational time during training, only one OOD point is sampled per training point per epoch, which relies on the stochastic gradient descent process to generalize this feature to the entire OOD domain and is one of the essential acceleration tricks provided by the NCP methodology. For the aleatoric term, a normal distribution with zero mean created from the aleatoric uncertainty prediction of the network, $\mathcal{N}(0,s_{j}^{2})$, was compared against an aleatoric prior distribution, selected to be:

Again, this represents another normal distribution where the standard deviation, $s_{y_{j}}$, is another user-defined parameter encapsulating any expected aleatoric uncertainty already known to be associated with the dataset output values. This is effectively meant to provide a competing pressure on the aleatoric uncertainty prediction, $s_{j}$, preventing it from growing to the point where a flat prediction with infinite aleatoric uncertainties becomes the most likely solution.

Basically, this way of treating the output results into the loss function is the main difference with respect to a regular neural network. For this matter, it is particularly meaningful to define several parameters that will help to adjust for the problem to solve. These user-defined values, $\mathbf{\Sigma_{x}}=\left\{\sigma_{x_{i}}\right\}$, $\mathbf{\Sigma_{y}}=\left\{\sigma_{y_{j}}\right\}$ and $\mathbf{s_{y}}=\left\{s_{y_{j}}\right\}$, are meaningful for the model as they are hyperparameters that introduce knowledge from the system and the data region to be modelled into the neural network. While in principle, this would allow extensions to the labelled training dataset, $D^{\prime}=\left\{\left(\mathbf{x},\mathbf{y},\mathbf{\Sigma_{x}},\mathbf{\Sigma_{y}},\mathbf{s_{y}}\right)_{k}\right\}$, for refining the uncertainty predictions, it was decided to leave these extra hyperparameters constant across the entire original training dataset, $D$, for this study. This is primarily for ease of demonstration for the BNN-NCP methodology, as it simultaneously reduces the burden of collecting the information necessary to refine these priors and the complexity needed to be captured by the training process. More details about the values used for these priors is available in Section 2.5.

During this study, the dataset is derived from a specialized collection JET-ILW plasma discharges [33] and their corresponding EuroPED predictions [16]. While the dataset includes several plasma quantities, this study focused on the input and output variables needed to predict the appropriate pedestal characteristics within a given plasma scenario accroding to EuroPED, as discussed in Section 2.1.

However, in order to improve the generality of the model via the introduction of dimensionless variables, the inverse aspect ratio, $\epsilon=a/R_{0}$, was used instead of the minor radius, $a$, and major radius, $R_{0}$. This process was also applied to the plasma current, $I_{p}$, replacing it by a safety-factor-like parameter, $\mu$, expressed as:

where $I_{p}$ is expressed in units of A, $B_{t}$ in T and $a$ in m. While it is uncertain whether pedestal phenomena follow any specific dimensionless scaling, it was attempted to improve its extrapolation capability to other tokamaks.

However, due to the particular application of the pedestal predictions inside the 1D transport solver within EuroPED, the exact top-pedestal temperature and density outputs of the ideal MHD calculation are no longer representative of the plasma state described by the profiles. Thus, the variables provided in this dataset used as the target outputs were:

• •

  electron density at $\psi_{\text{pol}}=\psi_{0}$, $n_{e}\!\left(\psi_{0}\right)$;

• •

  electron temperature at $\psi_{\text{pol}}=\psi_{0}$, $T_{e}\!\left(\psi_{0}\right)$;

• •

  and pedestal width, $\Delta$

where $\psi_{\text{pol}}$ is the poloidal magnetic flux, which is used as a geometry-independent radial coordinate, and $\psi_{0}\equiv 0.94$ for the purposes of this study. This specific flux coordinate value was chosen due to its statistical similarity to the $\beta_{\text{pol},ped}$ computed from EuroPED.

The pedestal dataset is composed of 1317 entries. However, 257 of them were removed in the training and testing because some of the variables of interest were not available. During the training of the model, the dataset was split into 90% for training and 10% for testing.

In Figure 4, the input data (train and test) used in EuroPED-NN is shown as histograms. The distribution of all the input variables as well as the range of values can be observed.

As observed in the loss function, Equation (1), several user-defined parameters and hyperparameters related with both the epistemic and aleatoric uncertainty predictions need to be determined. These values to the training routine such that the method to the particular problem. Therefore, additional considerations were required to provide reasonable targets for these values in training the EuroPED-NN.

The width of the OOD region for input sampling, $\mathbf{\Sigma_{x}}$, was chosen to span a constant proportion of the standard deviation within the dataset of the corresponding input quantity. These specific proportionality values are detailed Appendix A.

Regarding the target epistemic uncertainty, $\mathbf{\Sigma_{y}}$ in Equation (3), a small constant value was chosen across the entire dataset. So that, when comparing, using KL divergence, in the training loop the target distribution and the output distribution with OOD input, the OOD output distribution will be forced to become wider to match the thin target distribution.

Regarding the target aleatoric uncertainty, $\mathbf{s_{y}}$ in Equation (4), another small constant value was chosen across the entire dataset. This represents the prior knowledge that the EuroPED model is deterministic and thus, no data error should be present within the output values of the training dataset. The specific values used in this study are available in Appendix A.

All of these selections have been the result of different attempts and the personal reasoning of the authors regarding the nature of the model. Further details of these sensitivities are provided in Section 3, but the main reason is due to the fact that the epistemic and aleatoric uncertainties has been designated their own loss terms, each with their own weight coefficients, $\gamma_{epi}$ and $\gamma_{alea}$, in Equation (1).

Once these parameters informing the BNN-NCP priors have been decided, a rudimentary configuration optimization exercise yielded an architecture with a common dense layer with 20 neurons and specialized dense layers with 8, 8, and 10 neurons for $n_{e}\!\left(\psi_{0}\right)$, $T_{e}\!\left(\psi_{0}\right)$, and $\Delta$ outputs, respectively. Then, the BNN was trained using the Keras package [34] in Tensorflow [35] (version 2.15.0) with 2000 training epochs and a mini-batch size of 53, taking around 1.7 CPU-hours (10 minutes in 10 cores $\approx 100$ minutes). The Adam optimizer [36] was used with a learning rate of 0.001 and a exponential decay rate of 0.1.

In order to balance the relative importance of the various loss terms in Equation (1), a rudimentary optimization was performed. This resulted in the uncertainty weight coefficients of $\gamma_{\text{epi}}=0.1$ and $\gamma_{\text{alea}}=0.1$ for all outputs, $j$, and individual NLL weight coefficients of $\gamma_{\text{nll}}=\left\{7,7,2\right\}$ for $\Delta$, $T_{e}\!\left(\psi_{0}\right)$, and $n_{e}\!\left(\psi_{0}\right)$, respectively.

To test the performance of the model, the EuroPED-NN predictions are first compared against the data points within the training set itself. Figure 8 shows the accuracy for the output variables pedestal width, $\Delta$, electron temperature, $T_{e}\!\left(\psi_{0}\right)$, and electron density, $n_{e}\!\left(\psi_{0}\right)$, respectively. While they are all reasonably accurate, it can be visually inferred that $\Delta$ returns a worse match with the EuroPED database. This may indicate a higher degree of complexity behind the determination of this particular output, as it was found that it already required more neurons in the last BNN layer to achieve this fit quality. Overall, it can be said that the accuracy obtained through the training of the neural network is remarkably good, in spite of the typical difficulties of training multi-output regression networks.

However, another important point of investigating the BNN architecture revolves around its potential to differentiate interpolation regions from extrapolation regions via its uncertainty predictions [29]. This study examines this capability using the epistemic uncertainty prediction, where a generally lower uncertainty is expected in the interpolation regions. Figure 11 shows the relative predicted epistemic and aleatoric uncertainty for $\Delta$, $T_{e,ped}$, and $n_{e,ped}$ using randomly generated input data, as a function of the Euclidean distance of the input data to the centroid of the training input dataset. This distance is calculated after re-scaling these random samples with the same relation used to re-scale the input data distributions to have a mean of 0 and a standard deviation of 1. This is done to avoid introducing any bias on the distance metric due to the orders of magnitude differences spanned by the absolute values of the various input variables. It is also noted that some of the random input data points used may not resemble realistic operational plasma conditions, though this is not a requirement for demonstrating the characteristics of the uncertainty predictions under extrapolation. The training data centroid was chosen as the reference point as it should best represent the region which has the highest data density. Thus, as the distance increases, the epistemic error should increase as well because the neural network would have less information to train its predictive capability. 4000 input entries were randomly generated using two uniform distributions: the first one over the central 60 % of the dataset and the second over the central 98 %. This has been done intentionally to highlight both the points inside the training convex hull and those far from it.

Indeed, Figure 11 demonstrates that when the distance increases, the epistemic uncertainty also increases notably. The random points that fall inside the convex hull of the training data region are distinguished in these plots. In the epistemic case, the uncertainty is expected to be low inside the hull and growing as the points move further out of it. In the aleatoric case, the uncertainty should be higher inside the hull and decrease as the point move further out of it. Indeed, this is roughly observed in Figure 11. As the neural network catches the variance of the training input data through the aleatoric uncertainty, going far from training data means low aleatoric uncertainty as determined by the target $s_{j}$ values chosen in this study. On the other hand, epistemic uncertainty reflects the precision of the neural network regarding the stochastic variance of the weights and biases with respect the training data, then out of the data region this precision should be much lower, due to the lack of information. Although the general behavior of the uncertainty is well captured in Figure 11, a 1D toy example can is provided in Section 3.2 where it is more intuitively depicted while simultaneously verifying the model via a physical scan.

Additionally, in Appendix B a study over the input uncertainties effect on the output result of EuroPED-NN has been done. It can be observed that in most cases the aleatoric uncertainty accounts for the input uncertainty.

In order to verify the model based on existing physical intuition about pedestal behaviours, this section compares the EuroPED-NN predictions against known empirical relations.

Data-driven methods, in particular neural networks, have been applied to predict plasma pedestal quantities in the past. To put EuroPED-NN in its context a comparison will be carried out with EPED1-NN [21] and PENN [27] models, to know which are their main attributes and where EuroPED-NN stands in relation to them.

Regarding EPED1-NN, the model was developed to replicate EPED pedestal model [14]. The inputs of the model are the inputs of EPED: $n_{e,ped}$, $Z_{eff}$, $\beta_{N}$, $I_{p}$, $B_{t}$, $a$, $R$, $\kappa$ and $\delta$. The outputs are pedestal pressure $p_{ped}$ and pedestal width. EPED1-NN is built using a regular feed-forward neural network, and its accuracy achieves values around $R^{2}=0.987$ and $R^{2}=0.964$, for the respective outputs, when compared with EPED predictions as described in [21].

On the other hand PENN model is intended to replicate experimental data from JET-ILW pedestal database [33]. The inputs of the model are $Z_{eff}$, $\beta_{N}$, $I_{p}$, $B_{t}$, $a$, $\kappa$, $\delta_{up}$, $\delta_{low}$, $\kappa$, NBI Power, Total Power, plasma volume and $q_{95}$. Where $\delta_{up}$, $\delta_{low}$ and $q_{95}$ stands for upper triangularity, lower triangularity and q value at $\psi=0.95$ respectively. PENN is also build using a regular feed-forward neural network and the output variables are the pedestal density $n_{e,ped}$ and temperature $T_{e,ped}$, for which the accuracy in test data when compared with experimental data is $R^{2}=0.91$ and $R^{2}=0.93$, respectively.

Therefore, we can compare both models characteristic with EuroPED-NN, where the set of inputs (explained in the Section 2.1) differ from both EPED1-NN and PENN, then the usage cases are different, as happens between EPED and EuroPED. Depending on the available data one model can be more suitable than the other. The accuracy is reasonably good in all of them, including EuroPED-NN with $R^{2}_{\Delta,test}=0.846$, $R^{2}_{T_{e},test}=0.874$, and $R^{2}_{n_{e},test}=0.970$ (see Figure 8). However, this accuracy is quoted against the respective datasets of each model, meaning that the insights that can be extracted from this comparison are limited. Regarding uncertainty analysis, only EuroPED-NN is an uncertainty-aware model inherently providing its corresponding model and data uncertainty estimates. Regarding the model applicability, it is known that neither EPED nor EuroPED includes the effect of plasma resistivity. Thus, the empirical model PENN could be a better option when predicting the properties of a more resistive plasma. However, to explore the physics and behavior of the pedestal models in a fast way, both EPED1-NN and EuroPED-NN are also applicable.

Finally, an experimental pedestal database from ASDEX-Upgrade (AUG) tokamak [38, 39] with 50 shots, and its subsequent EuroPED predictions were available. This data is used to compare the EuroPED-NN predictions with the EuroPED evaluations. As the AUG database parameters lie outside the hull of JET parameters, the exercise constitutes an extrapolation of the model. This will test the adaptability of EuroPED-NN out of JET database values. The distribution of input values from the 50 AUG shots can be compared in Figure 17 with the JET database. The biggest difference between both of them are found in $\mu$, $\epsilon$, $P_{tot}$ and triangularity $\delta$ variables.

In the Figure 21 the EuroPED-NN predictions performance in AUG data is displayed for the three output variables. The points are split between those whose input points lie inside the central 99% of JET database range, and those who doesn’t. This central 99% boundary was chosen to effectively exclude outliers from artificially expanding the range limits. It can be appreciated how most of the points are outside this range, so it is therefore an extrapolating exercise. For $T_{e}(\psi_{0})$ and $n_{e}(\psi_{0})$ the match with the EuroPED predictions is reasonable with $R^{2}=0.645$ and $R^{2}=0.641$, respectively. However, in the case of the pedestal width, the extrapolation does not look to have succeeded. The reason for this is that the AUG predictions were run using a KBM constant of 0.09 (instead of the standard 0.076) as AUG has found that the standard prediction gives too narrow values. For this reason, the pedestal width variable cannot be extrapolated, while the rest variables perform much better.
The fact that EuroPED-NN performance in AUG data is reasonable good in the pedestal height quantities, even though the model is extrapolating, might mean that EuroPED-NN is adequately catching the trend between the input data and the pedestal quantities, rather than just adapting to the training data. However, to increase its accuracy and generalization power, future work training EuroPED-NN with data from other machines is encouraged.