Magnetohydrodynamic Equilibrium Reconstruction with Consistent Uncertainties ^†

Abstract

We report on progress towards a probabilistic framework for consistent uncertainty quantification and propagation in the analysis and numerical modeling of physics in magnetically confined plasmas in the stellarator configuration. A frequent starting point in this process is the calculation of a magnetohydrodynamic equilibrium from plasma profiles. Profiles, and thus the equilibrium, are typically reconstructed from experimental data. What sets equilibrium reconstruction apart from usual inverse problems is that profiles are given as functions over a magnetic flux derived from the magnetic field, rather than spatial coordinates. This makes it a fixed-point problem that is traditionally left inconsistent or solved iteratively in a least-squares sense. The aim here is progressing towards a straightforward and transparent process to quantify and propagate uncertainties and their correlations for function-valued fields and profiles in this setting. We propose a framework that utilizes a low-dimensional prior distribution of equilibria, constructed with principal component analysis. A surrogate of the forward model is trained to enable faster sampling.

1. Introduction

A frequent starting point for calculations in magnetic confinement fusion is the magnetohydrodynamic equilibrium. Ideal magnetohydrodynamics (MHD) [1] describes plasma as a single quasi-neutral fluid under the assumptions of infinite electrical conductivity, small ion gyroradius and negligible electron momentum. An MHD equilibrium is a solution to the static MHD equations (Equations (1)–(3)), describing the magnetic field ($\mathit{B}$) and its relation to the current density ($\mathit{J}$) and the pressure (p):

$$\begin{array}{cc}\hfill \mathit{J}\times \mathit{B}& =\nabla p,\hfill \end{array}$$

(1)

$$\begin{array}{cc}\hfill \nabla \times \mathit{B}& ={\mu }_0\mathit{J},\hfill \end{array}$$

(2)

$$\begin{array}{cc}\hfill \nabla ·\mathit{B}& =0.\hfill \end{array}$$

(3)

The solutions to Equations (1)–(3) are typically obtained in magnetic flux coordinates [2], curvilinear coordinates that greatly simplify calculations, for example, by aligning in such a way that $\mathit{B}$ appears straight. A large part of a typical MHD solution is, therefore, the mapping from flux coordinates to real-space coordinates. Inferring an MHD equilibrium from experimental diagnostics, a task known as equilibrium reconstruction, is of great interest for fusion experiments, to gain insight into unmeasurable or hard-to-measure quantities [3], such as flux surface geometry.

This task is routinely performed for tokamaks [4,5], where the toroidal symmetry further simplifies the MHD equations but proves to be more complicated for three dimensional (3D) systems, like stellarators. Typically, MHD solvers take the pressure and toroidal current profile as functions of a radial flux coordinate $\rho$ as input and solve for the optimal geometry of the magnetic flux surfaces. From the resulting equilibrium configuration, synthetic diagnostic signals can be calculated and compared to the corresponding real physical measurements. This makes 3D equilibrium reconstruction a fixed-point problem that is traditionally solved iteratively in a least-squares sense [3,6,7]. While the least-squares approach can result in suitable reconstructed parameters and uncertainty estimates [8], evaluations of the forward model are computationally demanding (a single equilibrium reconstruction can take up to several hours [9,10]), and common error estimates might fail to capture relevant uncertainties when high levels of noise are present for the equilibrium diagnostics [11]. Existing Bayesian frameworks [12] are also limited by the high computational costs of the forward model. Therefore, we aim to build a Bayesian framework for 3D equilibrium reconstruction that allows for fast sampling from the posterior distribution of equilibria while maintaining relevant physical constraints. In this work, we present, as a proof of concept, such a framework for a configuration of the Wendelstein 7-X stellarator (W7-X).

2. Methods

For the proposed equilibrium reconstruction approach using the Bayesian formalism, two major components are required: a prior distribution of equilibria and a fast evaluation of the likelihood function, both of which are not readily accessible. A schematic depiction of the steps taken to obtain these major components is shown in Scheme 1. To define the prior distribution of equilibria, a physically motivated prior distribution over equilibrium-defining parameters is formulated (Section 2.1). Samples from this prior distribution are then propagated into the space of equilibria using the Variational Moments Equilibrium Code (VMEC) [13] (Section 2.2) and subsequently into the space of synthetic equilibrium diagnostic signals with the code DIAGNO [14]. Since the space of equilibria is very high-dimensional, dimensionality reduction is performed on the equilibria using principal component analysis (PCA) (Section 2.3). In this way, a low-dimensional distribution over equilibrium configurations, which enables the training of polynomial chaos expansion (PCE) [15] surrogate models, can be defined (Section 2.4). These models map low-dimensional parameters $c_i$ onto the space of synthetic equilibrium diagnostic signals and thus circumvent high computational costs arising from the DIAGNO forward model. Through the combination of the surrogates with the low-dimensional equilibrium prior, the low-dimensional posterior distribution of equilibrium parameters can be inferred using Markov Chain Monte Carlo (MCMC) sampling in a Bayesian formalism, based on a set of (synthetic) measurement signals (Section 3).

2.1. Prior Distribution of Current and Pressure Profiles

We define a prior distribution for pressure $p\left(s\right)$ and toroidal current $I\left(s\right)$ profiles. They are functions of normalized toroidal flux $s=\frac{\psi }{{\psi }_0}$ and define an MHD equilibrium together with the total toroidal flux (${\psi }_0$). Here, $\psi$ denotes the toroidal flux. Often, reconstruction methods utilize simple parameterized functions [3,7,8], for example, two-power profiles, for describing equilibrium profiles. In order not to restrict the method to only favorable profiles, a broader prior is desired. Following similar concepts already realized in the MINERVA framework [10,12,16], we propose the use of a Gaussian Process [17] (GP) as a function generator for the profile gradients. This allows us to more easily specify the length scales and smoothness constraints based on physics information. For the pressure profile, the correlation length scales are larger at the core of the plasma than at the boundary, which is encoded in the GP kernel used in the MINERVA framework for W7-X [10]. We further constrain the pressure profile to be monotonic and 0 at the plasma boundary and enforce these constraints with exponentiation and integration, as well as a renormalization. For the current profile GP, a smooth rational quadratic kernel [17] is used, since detailed information on the current profile shape can be hard to capture with magnetic diagnostics alone [11] and the toroidal current is expected to be comparatively low for standard W7-X equilibria. Again, exponentiation and integration ensure monotonicity, but in contrast to the pressure profile, the current is 0 at the magnetic axis, and the scaling factor can be positive or negative. Using exponentiation as a positivity transform for the profile gradients allows for steep gradients but discourages many local plateaus that are found when using the absolute value as the transform.

$$\begin{array}{cc}\tilde{p}\left(s\right)\sim G{P}_{\mathrm{Minerva}}\hfill & \hfill p\left(s\right)=\frac{p_0}{{\int }_0^{1}exp\left(\tilde{p}\right)d{s}^{\prime }}{\int }_s^{1}exp\left(\tilde{p}\right)d{s}^{\prime }\end{array}$$

(4)

$$\begin{array}{cc}\tilde{I}\left(s\right)\sim G{P}_{\mathrm{RQ}}\hfill & \hfill I\left(s\right)=\frac{I_0}{{\int }_0^{1}exp\left(\tilde{I}\right)d{s}^{\prime }}{\int }_1^{s}exp\left(\tilde{I}\right)d{s}^{\prime }\end{array}$$

(5)

Finally, the GP is approximated with a PCA to allow for the space-filling latin hypercube sampling of the parameter space. Six shape parameters for the pressure profile and three shape parameters for the current profile correspond to an explained variance ratio of $99\%$ each. Combined with a parameter ${\psi }_0$ for the total toroidal magnetic flux, this gives us 12 equilibrium-defining parameters, as detailed in

Table 1. The prior parameters and their distribution.

Symbol	Distribution	Description
${\psi }_0$	Uniform in $\left[-2.5\mathrm{Wb},-1.6\mathrm{Wb}\right]$	Total toroidal magnetic flux
$p_0$	Uniform in $\left[0,200\mathrm{kPa}\right]$	Pressure at the magnetic axis
${\mathit{p}}_{\mathrm{shape}}$	${\mathcal{GP}}_{\mathrm{Minerva}}\left(\begin{array}{cc}{\sigma }_f& =1\hfill \\ {\lambda }_{\mathrm{core}}& =0.2\hfill \\ {\lambda }_{\mathrm{edge}}& =0.1\hfill \\ s_0& =0.9\hfill \\ s_w& =0.1\hfill \end{array}\right)$	Pressure profile shape factors reduced to 6 principal components
$I_0$	Uniform in $\left[-10\mathrm{kA},10\mathrm{kA}\right]$	Total toroidal current
${\mathit{I}}_{\mathrm{shape}}$	${\mathcal{GP}}_{\mathrm{RQ}}\left(\begin{array}{cc}\sigma & =2\hfill \\ \lambda & =0.2\hfill \\ \alpha & =2\hfill \end{array}\right)$	Current profile shape factors reduced to 3 principal components

. Several realizations of the equilibrium profile shapes obtained with the described procedure are visualized in Figure 1. From the prior defined in Table 1, 4000 samples are drawn and propagated into the space of MHD equilibria using VMEC.

2.2. VMEC

VMEC solves Equations (1)–(3) using a double-Fourier basis for the inverse coordinate representation $\mathit{x}\left(\rho ,\theta ,\zeta \right)=\left(R,Z,\varphi \right)$. Here, $\rho =\sqrt{s}$ denotes the magnetic flux surfaces, and $\theta$ and $\zeta$ are poloidal and toroidal angle-like flux coordinates. VMEC uses an energy functional in place of the local formulation, which is guaranteed to converge to an equilibrium in the absence of magnetic islands. This is not guaranteed for 3D geometries like stellarators [2], but VMEC can provide a suitable approximation of magnetic flux surfaces. In contrast to many other equilibrium codes, VMEC does not require a fixed plasma boundary but can calculate equilibrium flux surfaces, including the last closed flux surface, given a coil current configuration $\left[i_1...i_5\right]$, as well as ${\psi }_0$, $p\left(s\right)$ and $I\left(s\right)$.

A VMEC equilibrium can be fully described by the Fourier series for cylindrical coordinates $R\left(\rho ,\theta ,\zeta \right)$, $Z\left(\rho ,\theta ,\zeta \right)$ and potential $\lambda \left(\rho ,\theta ,\zeta \right)$, rotational transform profile $\iota \left(\rho \right)=\frac{d\chi }{d\psi }$ and total toroidal flux ${\psi }_0=\psi \left(\rho =1\right)$. Here, $\chi$ denotes the poloidal flux, and the cylindrical angle ($\varphi$) is a periodic continuation of the toroidal angle ($\zeta$) with the number of field periods ($N_{FP}$). The straight-field-line angle (${\theta }^{★}=\theta +\lambda \left(\rho ,\zeta ,\theta \right)$) can be calculated from $\lambda$. With stellarator symmetry, the Fourier series for R only contains cosine components, while Z and $\lambda$ only contain sine components. For the radial direction, VMEC uses a finite difference scheme and a linear spacing in the normalized toroidal flux $s={\rho }^2=\frac{\psi }{{\psi }_0}$.

The pressure profile ($p\left(\rho \right)$) can be calculated with a weak formulation of the radial force balance [13] under the assumption that Equation (1) is fulfilled:

$$p\left(\rho \right)={\int }_{\rho }^1-\frac{1}{{\mu }_0〈\mathcal{J}〉}\frac{d\psi }{d\rho }\left(\frac{d}{d\rho }〈B_{\zeta }〉+\iota \frac{d}{d\rho }〈B_{\theta }〉\right)d\rho$$

(6)

where $〈Q〉=\frac{1}{4{\pi }^2}\int \int Qd\theta d\zeta$ denotes an average over a flux surface and $\mathcal{J}$ is the determinant of the Jacobian.

2.3. Dimensionality Reduction

With $S=99$ radial points, $M=25$ poloidal modes and $N=12$ toroidal modes, a VMEC equilibrium can be described using $S·M·N$ parameters for R, Z and $\lambda$, respectively; S parameters for $\iota$; and a single parameter for ${\psi }_0$. With a total of 85,636 non-zero parameters, this space is difficult to handle with Bayesian methods. To reduce the dimensionality, principal component analysis (PCA) was applied on the VMEC equilibria.

To compensate for the incomparable units of the five quantities, they were all rescaled such that the total variance of each group of parameters was 1. For the total explained variance ratio of $99.9\%$, it was found that 11 principal components were sufficient to represent the equilibrium prior. The comparatively low dimensionality now facilitates Bayesian inference directly on the equilibrium parameters and also significantly reduces the number of data needed to represent the posterior distribution or a collection of samples.

For each of the five quantities $R,Z,\lambda ,\iota$ and ${\psi }_0$, the reconstruction error due to dimensionality reduction was evaluated using the $L^2$ metric (Figure 2). The relative error stayed below $1\%$ for all quantities except $\lambda$. This lead to lower quality of the direction of $\mathit{B}$ and derived quantities like p. The distribution of the low-dimensional parameters ($c_i$) was approximated with kernel density estimation (KDE) and assumed to be independent. This assumption held well enough and marginally increased the width of the total prior.

In Figure 3, the pressure profile is shown for a validation sample drawn from the prior distribution, together with a pressure profile calculated with Equation (6), and the same calculation applied after the validation sample was pushed through the dimensionality reduction. Indeed, the pressure profile can be calculated from the five quantities $R,Z,\lambda ,\iota$ and ${\psi }_0$, which determine $\mathit{J}$ and $\mathit{B}$, with only a small deviation near the magnetic axis, where VMEC is known to be inaccurate. However, dimensionality reduction significantly reduces the accuracy of this calculation.

2.4. Synthetic Diagnostics Surrogate Model

We trained polynomial chaos expansion (PCE) surrogate models that map from equilibrium coefficients $c_i$ to magnetic equilibrium diagnostics of W7-X. The aim of these surrogate models was to circumvent prohibitively high computational costs arising from DIAGNO forward model evaluations, which occur during sampling. The magnetic equilibrium diagnostics considered were saddle coils, diamagnetic loops and segmented Rogowski coils [18]. PCEs were built using the linear regression methods provided in the chaospy [19] Python package. The polynomials used in the expansions were chosen through hyperbolic truncation. In this truncation scheme, all multivariate polynomials, where the associated degrees of the univariate polynomials ${\alpha }_i$ fulfill [20]

$$d\ge {\left(\sum _i{\alpha }_i^h\right)}^{1/h},$$

(7)

are retained. Here, d denotes the maximum polynomial degree occurring in the expansion, and $h\in (0,1]$ determines the number of retained polynomials. While several other methods for choosing the basis polynomials in a sparse sense exist [21], we found that the hyperbolic truncation scheme with $d=6$ and $h=0.8$ worked well in the considered case. To estimate the performance of the resulting PCEs, the relative leave-one-out error ${\epsilon }_{loo}^R$ [22] was calculated for each surrogate model. The resulting ${\epsilon }_{loo}^R$ for these PCEs are visualized in Figure 4, grouped in terms of the synthetic diagnostic signal that is being approximated. Note that ${\epsilon }_{loo}^R$ is given relative to the variance present in the training data.

3. Results

With the previously discussed prior on the low-dimensional parameters ($c_i$) and the synthetic diagnostic surrogate, we tried to reconstruct a W7-X standard configuration equilibrium from a set of corresponding synthetic magnetic equilibrium diagnostic signals.

A set of equilibrium-defining parameters from the distribution discussed in Section 2.1 were propagated using the VMEC/DIAGNO forward model, providing a known ground truth. The resulting synthetic magnetic equilibrium diagnostic signals were then used together with a Gaussian likelihood function and the low-dimensional equilibrium prior to estimate the posterior distribution of $c_i$. For the likelihood function, the synthetic diagnostics were assumed to be uncorrelated with a $5\%$ error. Two methods for estimating the posterior were used: Laplace’s approximation and MCMC sampling. In both methods the evaluation of the synthetic magnetic equilibrium diagnostics was performed with the PCE surrogate models. For the MCMC method, samples were drawn using NUTS [23]. Figure 5 depicts one-dimensional projections of the posterior estimates for the parameters $c_i$ with respect to the associated prior distribution. Seen with respect to the prior width, there is a high uncertainty in the posterior distribution of parameters $c_4$ to $c_{10}$. However, they contribute little to the explained variance of the equilibrium PCA. To visualize how this estimated uncertainty appears in the physically relevant space, the maximum a posteriori estimate (MAP) and several samples were transformed back into the space of VMEC equilibria. The flux surface geometry, total toroidal flux and rotational transform profile are visualized in Figure 6, Figure 7 and Figure 8, respectively. In addition to the reconstructed equilibria, the ground-truth evaluation of VMEC is visualized. One can observe that the MAP estimate is comparatively close to the ground-truth values and that for all three quantities, the ground-truth value is covered within the area spanned by the posterior samples.

With Equation (6), the radial pressure profile of the posterior samples can be estimated (Figure 9). Significant differences between the MAP estimate and the ground-truth profile can be observed. However, the ground-truth profile is still covered within the space spanned by the posterior samples, and the magnetic diagnostics used for this reconstruction are rather insensitive to the pressure profile. The systematic discrepancy close to the magnetic axis can be attributed to the use of VMEC, which is known to have problems in this region. Additionally, the profile shapes expose nonphysical values near the last closed flux surface, as well as non-monotonicity, which was not present in the pressure profile prior. One reason for this, as shown above in Figure 3, is that PCA dimensionality reduction affects the derived pressure profile more than magnetic quantities, like the rotational transform profile (Figure 2).

4. Discussion

The results presented in Section 3 show that the presented framework is able to reconstruct an equilibrium configuration drawn from the prior distribution. That is, the ground truth lies within the region of high probability of the estimated posterior distributions, and the MAP estimate closely matches the true flux surface geometry and rotational transform profile. While the MAP estimate and posterior mean are found to be equivalent, Laplace’s approximation is not sufficient to model the posterior. Quantities that are derived from the magnetic quantities $R,Z,\lambda ,\iota$ and ${\psi }_0$, such as pressure profile $p\left(\mathit{J},\mathit{B}\right)$ are not reconstructed well, partially due to dimensionality reduction, but also because the considered magnetic diagnostics are not very sensitive to the pressure profile. While the derived pressure profiles fulfill the weak force balance by construction, they lack the physically motivated constraints that are encoded into the profile prior distribution in Section 2.1. In particular, they can violate the positivity constraint. However, the true profile does lie within the posterior distribution, and a pressure-sensitive diagnostic could be used to determine it, particularly as the mapping between real-space and flux coordinates is reconstructed well.

The presented framework offers a flexible prior distribution of profiles but directly reconstructs the flux surface geometry from magnetic diagnostics, allowing samples of the flux surface geometry to be drawn from the posterior. The equilibrium prior has to be constructed once for each device configuration, e.g., coil currents, but it offers reusability from this point onward and can be well described with only a few principal components. This not only allows Bayesian methods to be applied effectively but also significantly reduces the number of data needed to describe a posterior distribution or collection of samples. The PCE surrogate model similarly reduces the computational cost of 3D equilibrium reconstruction and offers fast access to uncertainties.

With the presented methods, we provide a proof of concept for a fast Bayesian equilibrium reconstruction framework. The accurate reconstruction of the derived quantities and thus the physical consistency of the posterior samples warrant further investigation. This could be achieved with different dimensionality reduction methods, for example, using a variational autoencoder with additional physics constraints, or different surrogate models, as well as more accurate MHD solvers.