Uncertainty quantification in three-dimensional magnetohydrodynamic equilibrium reconstruction via surrogate-assisted Bayesian inference

In three-dimensional (3D) equilibrium, reconstruction defining parameters of an ideal magneto-hydrodynamic equilibrium are inferred from a set of plasma diagnostic measurements. For the reconstructed parameters, various forms of uncertainty estimates exist within common 3D reconstruction frameworks. These estimates often assume a Gaussian posterior distribution. The validity of this assumption is not obvious in such highly nonlinear inverse problems, and therefore the accuracy of the estimates cannot be guaranteed. In this work, we formulate the problem of 3D equilibrium reconstruction in a Bayesian sense and explore the posterior distribution of reconstruction parameters via Markov chain Monte Carlo (MCMC) sampling. The target reconstruction parameters, that is, shape and scaling factors of the pressure and toroidal current-density profiles as well as the total toroidal flux, are taken from a reduced subspace of Wendelstein 7-X equilibrium configurations. Since the corresponding forward model evaluations are computationally demanding, we replace the forward model via a polynomial chaos expansion surrogate. We


INTRODUCTION
The analysis of uncertainties and their propagation in physical models is a key task and concept that assists in establishing trust in predictions as well as the general validation and verification process. The well-established Bayesian framework often poses a natural way to tackle the problem of uncertainty quantification by casting the issue into a probabilistic setting that allows the incorporation of prior knowledge. The most essential formula in this context is given by Bayes' theorem which, written in terms of probability density functions (PDF), reads, In the context of solving an inverse problem, we denote with some model parameters that are of interest to us and with D the corresponding observations that form the input for the inference. For example, D can be a set of experimental measurements for which a physical forward model that depends on is known. Hence, following Equation (1) the posterior PDF p( |D) of given the observations D is determined by the likelihood p(D| ), the prior p( ) and the evidence p(D). While Equation (1) provides an elegant formalism for defining problems, one, more often than not, has to rely on numerical methods when solving for the full posterior or integral quantities like its moments. A commonly used method for approximating posterior distributions is through Markov chain Monte Carlo (MCMC) sampling. Within this approach, random samples are drawn from the posterior distribution. This usually requires thousands of sequential evaluations of the forward model. While MCMC can aid in providing elaborate knowledge about the posterior, it can also become numerically extremely demanding if model evaluations are costly. One way of circumventing such repeated prohibitive costs is by building a cheap-to-evaluate surrogate model for the forward model.
Various approaches that deal with or follow along the lines of building surrogate models have already been applied within the field of nuclear fusion plasma physics. For example, Merlo et al. [1] used neural networks to build a surrogate for the ideal magneto-hydrodynamic (MHD) equilibrium solver VMEC [2] for the Wendelstein 7-X stellarator. Coster et al. [3] utilised a polynomial chaos expansion (PCE) model for uncertainty quantification in a turbulence-transport workflow. Callaghan et al. [4] applied the function parametrisation technique [5] to the problem of equilibrium reconstruction in the W7-AS stellarator. Similarly, function parametrisation was applied by Sengupta et al. [6] to recover equilibrium parameters for Wendelstein 7-X configurations.
In this work, we employ a PCE surrogate model for a forward model that maps MHD equilibrium defining quantities ( ) to synthetic magnetic diagnostic signals of the Wendelstein 7-X stellarator. The surrogate is used together with Laplace's approximation or MCMC sampling to approximate the posterior of the parameters , given a set of simulated measurements. The posterior approximations obtained through MCMC sampling are then compared to Laplace's approximation, which assumes a Gaussian posterior.
The assumption of a Gaussian posterior is present in existing uncertainty estimates within common three-dimensional (3D) equilibrium reconstruction frameworks (Section 1.1). However, a posterior distribution is not guaranteed to be Gaussian for such highly non-linear inverse problems. Therefore, this work provides further novel insight into the validity and limitations of these uncertainty estimates for the relevant stellarator case of Wendelstein 7-X.
In Section 1.1 we recapitulate the common approaches to the equilibrium reconstruction process and describe the major steps taken in the evaluation of the forward model. In Section 1.2 the basic concept behind the PCE approach is presented. Section 2 describes the setup for the numerical calculations and discusses the results obtained through this model. Finally, in Section 3 a conclusion is given together with an outlook on possible future work.

Equilibrium reconstruction
Equilibrium reconstruction describes the process of inferring an ideal MHD equilibrium or its defining properties from a set of experimental measurements. The defining properties include, for example, the pressure-profile toroidal current density-profile and the rotational transform-profile given with respect to a flux surface label. The resulting reconstructed equilibria may give insight into unmeasurable or hard-to-measure quantities [7] and often form the basis for subsequent analyses. [8,9] Hence, suitable error estimates for the reconstructed parameters are of interest. Two core elements of the forward model for the reconstruction process are the actual equilibrium solver and the mapping of an equilibrium solution to corresponding experimental diagnostics. A common approach, given some experimental diagnostic measurements, is to minimise the least squared distance between some fit function incorporating the synthetic diagnostics and fit targets incorporating the actual measurements, Here, denotes the vector of model parameters that are to be reconstructed, N D the total number of targets, y i the i th fit target with the corresponding weight 1 i and f i ( ) the evaluation of the fit function at point . One usually defines i as the standard deviation of the measurement noise. Minimising such an expression is equivalent to finding the maximum posterior (MAP) estimate in the Bayesian sense when both prior and likelihood follow a normal distribution and the covariance matrices appearing in the likelihood and the prior are assumed to be diagonal. This assumption implies that the measurements and the parameters of the prior are uncorrelated. [10] This fundamental principle has been employed in various equilibrium reconstruction codes like the well-established EFIT code, [11] IDE, [12] which is based on CLISTE, [13] or NICE [14] for the axisymmetric scenario as well as within V3FIT, [15] the STELLOPT [7,16] and MINERVA framework [17,18] in the 3D case. While the details concerning the exact formulation of targets and fit function may vary throughout the codes, the fundamental idea of matching synthetic to real diagnostics stays the same. Fitting the targets requires an iterative adaptation of the input parameters to the forward model and its consecutive evaluation. Depending on the prior assumptions made on the model the number of free parameters can become rather high, for example, 63 free parameters were used by Lazerson et al. [7] The choice of these free parameters is crucial since it can introduce a significant bias in the reconstructed result. Naturally, this applies as well to the equilibrium solver, as it not only depends on the free parameters but also on internal assumptions. One prominent example is the assumption of nested closed flux surfaces in equilibrium codes such as VMEC which can therefore not treat equilibria with magnetic islands and chaotic regions. [15] While uncertainties associated with these modelling choices too are of great interest, we will focus solely on the propagation of uncertainties originating from the measurement signals in this work. The existing reconstruction frameworks usually provide some form of error estimates for the reconstructed parameters. Within the IDA [19] approach at ASDEX Upgrade a variety of uncertainty estimate methods, including MCMC and Gaussian process regression, have been utilised by Fischer et al. [20] Dikovsky et al. [21] performed Bayesian posterior sampling via Hamiltonian Monte Carlo for reconstruction tasks at the C-2 W experiment. [22] Lazerson et al. [7] used the asymptotic standard prediction error, the standard error and the asymptotic standard prediction error and Cianciosa et al. [10] propagated uncertainties using the best linear unbiased estimate in the 3D case. These latter estimates are based on derivatives that are obtained during the reconstruction process and assume a posterior that follows a Gaussian distribution centred at the parameter position resulting from the optimisation.

Polynomial chaos expansion
In PCE surrogate modelling, a computational model is approximated by a set of polynomials that are orthonormal with respect to the PDF of random input variables. [23,24] Consider a set of N input variables̃=

{̃1
, · · · ,̃N } for a deterministic mappingỹ = (̃) that are affected by uncertainty. Due to this uncertaintỹis described by a random vector Λ with values in a set Ω and a prescribed probability density function p Λ . Hence, the model response will be a random variable as well, denoted by Y = (Λ). If the random variables L are independent, one can expand the response Y on an orthogonal polynomial basis, [24][25][26][27] where the c i denote the unknown expansion coefficients and i the multivariate polynomials. In practice, one has to truncate the series at some given polynomial degree d so that, with the number of coefficients given by, The basis polynomials are ideally chosen in an optimal sense depending on the type of random variables, usually picked from the Askey scheme [28] for common distribution functions. For less common distributions, one might construct the optimal basis via the Gram-Schmidt procedure. [24] Optimality is given if, where ij is the Kronecker delta and Z j is a normalizing constant. [24] Vector-valued responses Y are approximated component-wise. If the number of required expansions becomes infeasible, one might utilise dimensionality reduction procedures such as principal component analysis (PCA) as suggested by Blatman and Sudret. [29] The calculation of the unknown coefficients can be performed in several different ways, broadly categorised into intrusive and non-intrusive approaches. [26] In this work, we will focus solely on the non-intrusive regression approach. Given a set of M realisations , with M typically chosen in the range between 2 ⋅ D and 3 ⋅ D, [23] and the corresponding model evaluations Y e = {ỹ (1) , · · · ,ỹ (M) } , one can write the solution for the estimation of the coefficients via the least-squares regression,ĉ Here, . [23,25] For high dimensional Λ, the required number of model evaluations quickly becomes infeasible if all the basis functions are retained. There exists a broad family of approaches to circumvent this curse of dimensionality by selecting a sparse subset of the basis functions to build the expansion. A great overview of these sparse PCE methods is given by Lüthen et al. [23] As with any regression method, one has to validate the performance of the employed model. In this work, we choose the leave-one-out error metric as the performance measure. The leave-one-out error can be efficiently evaluated for a PCE via, [25] Here, h i denotes the i-th diagonal element of the matrix ( T ) −1 T . This corresponds to an M-fold cross-validation and allows efficient use of the training data to access an error estimate. Weighting loo by the empirical variance of the response samples defines the relative leave-one-out error, with  (Y e ) = 1 . As mentioned by Blatman and Sudret, [25] the study of Molinaro et al. [30] indicates that the leave-one-out technique performs well in terms of bias and mean-squared error.

PCE surrogate
We build a PCE surrogate model for a set of N D magnetic diagnostics of the Wendelstein 7-X stellarator, to enable a framework where comparisons between different approaches for uncertainty estimations of reconstructed equilibrium parameters are easily accessible. The basic idea behind this approach is to circumvent the prohibitive computational costs arising from the evaluation of the forward model. For this analysis, the forward model comprises a free boundary evaluation of VMEC to obtain the ideal MHD equilibrium and the code DIAGNO [31] to calculate the signal of the diagnostics due to this equilibrium. For the magnetic diagnostics, a set of diamagnetic flux loops, saddle coils and segmented Rogowski coils 1 is used. To calculate different equilibria, the radial pressure profile P( ), the radial toroidal current-density profile I ′ ( ) = dI d and the total toroidal flux were varied in the input to VMEC. Here, ∈ [0, 1] denotes a magnetic flux surface label. For the two radial profiles we chose a functional representation, the two-power profile, which in our case depends on two parameters: While in general more sophisticated profile representations can be desirable, we choose the two-power profile as it provides a suitable trade-off between dimensionality, for example, two free parameters per profile, and flexibility in the profile shape. This is especially important, due to classical PCE's curse of dimensionality affliction. The specific parametrisation of the profiles was constructed around a reference VMEC calculation with the same hyper-parameters as the ones used in this work. This motivated the fifth power in Equations (10) and (11), which effectively limits the magnitude of the profiles in the vicinity of the last closed flux surface. Therefore, it contributes to the stability and accuracy of the forward model evaluations in this region. Hence, our forward model depends on five parameters: the pressure and toroidal current scaling factors P 0 and I 0 , withÎ 0 ∝ I 0 so that I 0 = I( = 1), the corresponding profile shape parameters P 1 and I 1 and the total toroidal flux . These form the set of free equilibrium-defining parameters = {P 0 , P 1 , I 0 , I 1 , }. All other input parameters for VMEC and subsequently DIAGNO, such as, for example, the coil current ratios, are kept fixed throughout this analysis. For the generation of the diagnostic signal data set used in the PCE regression, we placed independent uniform priors (denoted by ) on the parameters of interest and sampled from the joint five-dimensional distribution via Latin hypercube sampling. The individual priors are given by, In total 1000 samples were drawn and propagated through the VMEC/DIAGNO forward model. As some parameter combinations can still result in equilibria that are not of interest, that is, exceed relevant physical boundaries, we filtered the results with regard to the minor radius position and the total plasma volume. Both, the physical limits on the scaling factors and the toroidal flux as well as the volume and minor radius restriction follow along the lines of the procedures employed by Merlo et al. [1] and Sengupta et al. [6] Within the generated data set, given the reduced complexity of the setup, only the plasma volume of some calculations exceeded relevant upper boundaries. Therefore, those equilibria with plasma volumes above 38 m 2 are discarded, resulting in a total of 878 evaluations utilised in the PCE regression. The PCE for the different diagnostic signals itself was performed utilising the Chaospy [33] Python package. Prior to performing the regression, the input parameters, and respectively their corresponding prior distributions, were scaled to lie within the interval [0, 1]. The final degree of the polynomial expansion was chosen by evaluating loo for all degree values that still would satisfy M ≥ 2D and picking the one with the smallest average error over the different diagnostics. This did yield final PCEs with polynomials up to the fifth order for each diagnostic signal. In Figure 1 the corresponding relative loo are shown. With the given data set and the chosen PCE degree, we obtain a maximum loo that would correspond to approximately 0.3% of the empirical variance of the response samples. 1 The (segmented) Rogowski coils at W7-X are assembled from series-connected individual coil rods (see Endler et al. [32] for details). In this work, the responses are computed individually for these rods to explicitly exclude the crossover connections between the rods from the computation. Comparison with experiment would require to sum up the contributions from all rods of a Rogowski coil segment in order to compare to experimental measurements.

F I G U R E 1
Relative leave-one-out error ( R loo ) of the polynomial chaos expansions (PCE) for the different diagnostic signals. The error estimates are measured relative to the variance present in the training data set. The error estimates are grouped in terms of the corresponding diagnostic signal that is approximated via PCE (saddle coils-red crosses, diamagnetic loops-blue diamonds, segmented Rogowski coils-black dots).

F I G U R E 2
First-order Sobol indices obtained from the polynomial chaos expansions for the different magnetic diagnostics. Each row of the matrices belongs to one diagnostic signal, with the contribution to the corresponding variance split between the five reconstruction parameters. Thus, the values in each row add up to a total of one. The higher the value, the more the parameter of that column contributes to the signal variance of the diagnostic of the value's row.

Sensitivity analysis
One advantageous feature of PCE is the straightforward calculation of Sobol indices from the expansion coefficients. Through the first-order Sobol indices, one can gain insight into the sensitivity of the response quantity with respect to single input parameters. Higher-order indices account for the mixed influence of combinations of parameters. [26] Figure 2 depicts the first-order Sobol indices for the magnetic diagnostics with respect to the input parameters of the PCE. It can be seen that the combination of diagnostic signals in general is sensitive to changes in the pressure profile parameters. As expected, sensitivity to is present in the diamagnetic loops and I 0 mainly influences the segmented Rogowski coil signals. Solely the shape parameter I 1 appears to have little impact on the signal response and therefore, one can expect high uncertainties in I 1 when using these diagnostics for reconstruction. Figure 3 provides a visualisation of the second-order Sobol indices calculated from the PCEs. It shows that from the mixed influences, the one of the pressure profile parameters is the most dominant. This hints at a relevant correlation between these two. Also, a minor mixed influence between and P 0 as well as and P 1 can be observed for all considered types of magnetic diagnostics.

F I G U R E 3
Second-order Sobol indices obtained from the polynomial chaos expansions for the different magnetic diagnostics. For each diagnostic, a 5 × 5 matrix containing the second-order indices is calculated. The x-y-plane in the visualisation depicts the entries to these matrices, whereas the z-axis represents the index value that is the mixed influence of the parameters determined by the x-y position. Each graph depicts all matrices belonging to the signal type defined in the title. The dashed black lines are intended to act as visual guidance for projecting elevated points back onto the x-y-plane.

Equilibrium reconstruction methods
With the created PCE surrogate, the aim is to reconstruct the parameters for a given noisy set of magnetic diagnos- To ensure that the underlying ground truth values for lie within the model under consideration, one predefined set of parameters t is chosen and propagated into the space of diagnostic signals using the VMEC/DIAGNO forward model. The obtained signals D are perturbed via independent Gaussian noise, chosen as a fraction of the standard deviation of the training data for the PCEs ( i = ⋅ i train and ∈ [0.01, 0.2]) so that, and Thus, a numerically simulated measurement is obtained which forms the basis for inferring the ground truth value again. Hence, the original t that created the simulated measurement signal is known; however, one does not know how the added uncertainty, due to the noise, will translate back into the space of reconstruction parameters during the reconstruction process. In other words, the quantity of interest for this analysis is the joint posterior distribution of P 0 , P 1 , I 0 , I 1 and . In particular, we compare two approaches for approximating this posterior, in both of which we replace the forward model with the PCE-surrogate: In the first approach, we approximate the posterior trough MCMC sampling using Bayes' theorem Equation (1), a Gaussian likelihood for the simulated measurements D M and the prior Equation (12).
In the second approach, we estimate the posterior via Laplace's approximation (p post L ). To do so, we find the MAP estimate, again using the chosen prior and the Gaussian likelihood function, via optimisation and calculate the Hessian matrix at that position with respect to changes in the reconstruction parameters. Due to the choice of the prior and likelihood, the MAP estimate reduces to the 2 optimisation Equation (2) within the region of considered parameters. The necessary derivatives are in our case available through the PCE models. This method is in close analogy to existing error estimates for the 3D reconstruction mentioned in Section 1.1. Both approaches have been realised utilising the Python PyMC3 [34] library. For the MCMC sampling, we apply the sequential Monte Carlo method readily available within PyMC3.

Equilibrium reconstruction results
We performed the described reconstruction procedures (Section 2.3) for different levels of measurement noise. Figure 4 depicts a summary of the estimated posterior distribution obtained through MCMC (blue) in comparison to Laplace's approximation (red). The assumed noise level for this reconstruction was set to be = 0.15 ⋅ train for each diagnostic signal independently and only one set of noisy measured signals was considered. The measurement accuracy for the equilibrium diagnostics at Wendelstein 7-X is typically specified in terms of a minimum and maximum flux change (Δ min , Δ max ) and is of the order of Δ min Δ max ∼ 10 −3 . [32] This level of accuracy is comparable to noise levels below = 0.04 ⋅ train for most equilibrium diagnostics considered. However, contributions from systematic uncertainties in measurements, for example, due to the sensitivity of the measurement electronics to temperature drifts, [35] and modelling, for example, due to the limited resolution of the forward model, usually raise the level of uncertainty. Andreeva et al. [36] assumed a 5% uncertainty of the measured magnetic diagnostic signals for the equilibrium reconstruction of  and p post MCMC is present below a signal noise threshold of about = 0.12 ⋅ train . However, for very low noise levels it is plausible that the influence of the error introduced via the PCE becomes relevant as ground truth values may no longer be covered within high-density regions of the posterior if this uncertainty is not accounted for. For noise levels above approximately = 0.14 ⋅ train , the posteriors start to significantly differ from each other. This might be associated with the increasingly prominent non-linear correlation between P 0 and P 1 as well as the fact, that less information about I 1 can be retrieved. That is, I 1 's prior's influence prevails.
Removing diagnostic signals, and therefore information, from the reconstruction procedure translates into a reduced region of similarity between p post L and p post MCMC . This is shown via the orange diamond-shaped points in Figure 7. Here, the number of diagnostic signals from the segmented Rogowski coils was reduced to 256 signals from the original 1167. This causes an increased level of uncertainty due to the missing information that translates into effects such as an emphasised curvature in the projected posterior between P 0 and P 1 as well as longer tails on the projected distribution of I 1 . Hence, non-Gaussian behaviour of the posterior for is already observable at approximately = 0.10 ⋅ train .
The similarity between p post L and p post MCMC for the full set of diagnostics below = 0.14 ⋅ train provides further validation for the use of the existing error estimates (Section 1.1) in these signal-noise regions. Nevertheless, we argue that care

F I G U R E 7
Kullback-Leibler divergence between Laplace's approximation for the posterior distribution and the posterior estimated via Markov chain Monte Carlo sampling for different levels of signal noise with respect to the standard deviation of the training data set train . The blue crosses correspond to the posterior distributions obtained using the full set of magnetic diagnostics, whereas the orange diamonds correspond to those obtained using the reduced set. must be taken when utilising these error estimates in regimes above this threshold, as they might fail to capture relevant features of the posterior.

CONCLUSION
We successfully constructed a PCE surrogate model for the mapping of five equilibrium-defining parameters to a set of magnetic diagnostic signals for the Wendelstein 7-X stellarator. For this reduced-complexity parameter space, we replaced the costly forward model in the framework of 3D equilibrium reconstruction with the PCE surrogate model. Therefore, we were able to estimate the posterior of the equilibrium-defining parameters, in a Bayesian setting, from a set of noisy synthetic diagnostic signals, which were created from the known ground truth. To this end, MCMC sampling was utilised. The posterior distributions obtained from this inference procedure were then compared to Laplace's approximation of the same posterior, which is similar to commonly used error estimates in 3D equilibrium reconstruction. The similarity between these two posterior estimates for different levels of noise on the diagnostic signals was then measured in terms of the Kullback-Leibler divergence. This provided novel insight into the region of validity of approaches such as Laplace's approximation for the estimation of uncertainties within the considered model. We found that for this model the posterior distribution of the reconstruction parameters can be rather well approximated for a significant range of the considered diagnostic signal noise levels using Laplace's approximation. At high noise levels, the posterior cannot be described with a Gaussian distribution and, therefore, Laplace's approximation fails to capture some of the uncertainty present in the reconstructed parameters. The specific ranges of validity are dependent on the underlying model. This includes the profile parametrisation as well as the types and the number of equilibrium diagnostics. However, the presented approach enabled us to perform the comparison between Laplace's approximation and MCMC sampling in a controlled and computationally trackable environment. Hence, this analysis contributes to further validate existing error estimation approaches and showcases some of the limitations of such methods. To utilize the presented methodology for equilibrium reconstruction using experimental diagnostic signals would require an extension to higher complexity scenarios. Such scenarios could include different radial profile representations and additional free parameters. This as well as the analysis of the influence of different model uncertainties on the reconstruction process might raise topics for future work.