3.1. Magnetohydrodynamic equilibria

An MHD equilibrium is determined by the boundary shape together with two functions of the toroidal flux, typically the pressure and enclosed toroidal current (Kruskal & Kulsrud Reference Kruskal and Kulsrud1958). It is not obvious how to best sample this space, particularly the space of boundary shapes. Boundaries with self-intersections or other such pathologies should be avoided. Also, a compromise must be struck between preserving similarity of the shapes to ‘real’ devices (built experiments or theoretical configurations designed through serious optimisation) while also allowing for new possible geometries.

To balance these considerations, we assemble a set of 23577 equilibria drawn from three groups. The first group has heliotron-like rotating ellipse shapes, in which parameters of the shape (number of field periods, aspect ratio, elongation, axis torsion, beta) are sampled randomly. Equilibria in the second group are taken from the QUASR database of quasi-axisymmetric (QA) and quasi-helically symmetric (QH) configurations (Giuliani Reference Giuliani2024; Giuliani, Rodríguez & Spivak Reference Giuliani, Rodríguez and Spivak2024). We use both the original QUASR configurations, which are all vacuum fields, and also generate new configurations by adding pressure while keeping the plasma boundary shape fixed. For the third group, random boundary shapes are generated by sampling Fourier modes that have been fit to a dataset of previous stellarator shapes. The combined set of configurations includes values of aspect ratio ranging from 2.9 to 10, volume-averaged beta from 0 % to 5 %, and number of field periods from 2 to 8. Thus, the set of equilibria is diverse, and includes both omnigenous and non-omnigenous geometries. All equilibria have the same minor radius and same boundary toroidal flux, resulting in the same normalising length and normalising field strength, so the gyro-Bohm normalisations are identical. Examples of equilibria from the three classes are shown in figures 2–4. More details of the procedures for generating equilibria are given in Appendix A.

Figure 2.

Examples of the rotating-ellipse equilibria included in the dataset. Two-dimensional plots show the cross-sections at which the toroidal angle is 0, $1/4$ , $1/2$ and $3/4$ of a field period. Three-dimensional images show each configuration from two angles, with colour indicating $|B|$ (red = high, blue = low) and field lines in black. Left columns, configurations in which the boundaries are centred on a circle. Right columns, configurations in which the boundaries are centred on a curve with torsion.

Figure 3.

Examples of the QUASR quasi-axisymmetric and quasi-helically symmetric equilibria included in the dataset. Two-dimensional plots show the cross-sections at which the toroidal angle is 0, $1/4$ , $1/2$ and $3/4$ of a field period. Three-dimensional images show each configuration from two angles, with colour indicating $|B|$ (red = high, blue = low) and field lines in black.

Figure 4.

Examples of the equilibria generated with random boundary Fourier modes. Two-dimensional plots show the cross-sections at which the toroidal angle is 0, $1/4$ , $1/2$ and $3/4$ of a field period. Three-dimensional images show each configuration from two angles, with colour indicating $|B|$ (red = high, blue = low) and field lines in black.

3.2. Gyrokinetic simulations

In each equilibrium, we extract at least four flux tubes, more in the QUASR vacuum configurations so as to include more omnigenous geometries. The radial location $\rho =\sqrt {s}$ for each flux tube is sampled randomly from $[0, 1]$ , where $s=\psi /\psi _a$ is the normalised toroidal flux. All flux tubes are stellarator-symmetric, with at least one flux tube in every configuration centred on each of the points $(\theta ,\phi )=(0,0)$ , $(\pi ,0)$ , $(0,\pi /n_{\mathrm{fp}})$ and $(\pi ,\pi /n_{\mathrm{fp}})$ . Nothing in the procedure limits the method to stellarator-symmetric equilibria or flux tubes. The total number of flux tube geometries in the dataset is 100 705. The length of the flux tubes in real space was chosen to be the same for every configuration, 75 times the minor radius, which for aspect ratio 6 corresponds to roughly two full toroidal transits. Due to the variation in $\iota$ , the number of poloidal transits varied between configurations. By making all flux tubes have the same physical length, distances in the $z$ coordinate can be compared meaningfully between flux tubes, and for each integer $j$ , the $j$ th Fourier series coefficients in all flux tubes correspond to the same physical $k_{||}$ . These properties may be advantageous for extracting patterns in the dataset.

Nonlinear electrostatic simulations with adiabatic electrons were then carried out with the GX code (Mandell et al. Reference Mandell, Dorland, Abel, Gaur, Kim, Martin and Qian2024). We run GX twice for each flux tube, resulting in two datasets each of size $N=100\,705$ . In the first dataset, the temperature and density gradients were fixed to the same values for all flux tubes. In the second dataset, these gradients were chosen randomly for each flux tube. Using the latter dataset, models can potentially find interactions between geometry and gradients in their effect on the heat flux; for example, the dependence on geometry may differ close to the critical gradient compared with far above the critical gradient. However, we also assembled the dataset with fixed gradients since it enables several of the analyses in § 5 and allows us to focus cleanly on the effects of geometry. For the fixed-gradient dataset, we take $a/L_{Ti} = 3$ and $a/L_n=0.9$ , where $a/L_{Ti} = -(a/T_i) \,{d}T_i/{d}x$ and $a/L_n=-(a/n)\, {d}n/{d}x$ are the normalised gradients of ion temperature and density, reflecting typical measurements at the $s=0.5$ surface of W7-X (Beurskens et al. Reference Beurskens2021; Lunsford et al. Reference Lunsford2021; Zhang et al. Reference Zhang2023). For the varying-gradient dataset, each simulation has $a/L_{Ti}$ and $a/L_n$ sampled randomly within plausible experimental ranges.

Additional details of the turbulence simulations are given in Appendix B.

3.3. Configurations with very low or high heat flux

The distribution of heat fluxes for the dataset is shown in figure 5. It is striking that even for the dataset with fixed gradients, $Q$ varies by over four orders of magnitude over the data. This variation is evidently due purely to the geometry.

Figure 5.

Distribution of heat fluxes for the fixed-gradient and varied-gradient datasets. In the latter, 30 % of the simulations were stable, with $Q \approx 0$ . Simulations with $Q\lt 0.1$ , considered stable, are included in the leftmost bar.

Figure 6.

Some trends are apparent between flux tubes with very low or high heat flux. The columns show six flux tubes from the $n_{fp}=3$ equilibria with random boundary Fourier modes, the first three stable, the last three with very high $Q$ at the same gradients. The top six rows are the inputs to the gyrokinetic-quasineutrality system ( $B^{-3}\boldsymbol{B}\times \boldsymbol{\nabla} B\boldsymbol{\cdot} \boldsymbol{\nabla} y$ is omitted for simplicity since it is similar to $B^{-2}\boldsymbol{B}\times \boldsymbol{\kappa }\boldsymbol{\cdot} \boldsymbol{\nabla} y$ ), while the bottom row shows the contribution to the total heat flux versus $z$ .

It is interesting to examine configurations with particularly low or particularly high heat flux. Several such flux tubes are displayed in figure 6. The columns represent six flux tubes from the dataset. In the fixed gradient dataset with $a/L_{T}=3$ and $a/L_n=0.9$ , the first three tubes are stable, while the other three have $Q \gt 500$ , even though the gradients are identical. These flux tubes are all taken from $n_{fp}=3$ equilibria from the group with random boundary Fourier modes. The first six of the seven rows show the geometric inputs to the gyrokinetic-quasineutrality system. To simplify the figure, $B^{-3}\boldsymbol{B}\times \boldsymbol{\nabla} B\boldsymbol{\cdot} \boldsymbol{\nabla} y$ is not shown since it is similar to $B^{-2}\boldsymbol{B}\times \boldsymbol{\kappa }\boldsymbol{\cdot} \boldsymbol{\nabla} y$ . The bottom row shows the contribution to the heat flux $Q$ as a function of $z$ . For each of the seven rows, the vertical scales are the same for each column to allow comparison.

Several patterns are apparent. The configurations with highest heat flux have regions with very large $|\boldsymbol{\nabla} x|^2$ . The high-heat-flux tubes also have larger magnitudes of $B^{-3}\boldsymbol{B}\times \boldsymbol{\nabla} B\boldsymbol{\cdot} \boldsymbol{\nabla} x$ . The stable tubes have more negative values of $B^{-2}\boldsymbol{B}\times \boldsymbol{\kappa }\boldsymbol{\cdot} \boldsymbol{\nabla} y$ , meaning mostly good curvature. All of these patterns foreshadow findings in § 5.

For the geometry-gradient pairings with very large heat flux, where the gyro-Bohm-normalised $Q$ is not small compared with $1 / {\rho _*}$ , with $\rho _*$ the ratio of typical gyroradius to macroscopic scale, the ordering used to derive the gyrokinetic model may no longer be accurate. Therefore, these simulation runs may not accurately describe real physical scenarios. Nonetheless, it is still reasonable to consider the $\delta f$ -gyrokinetic model as an abstract mathematical function, a map from seven functions (the geometry) and two real numbers (the gradients) to a real number ( $Q$ ), and to develop approximations of this function. Moreover, since surrogates will be more accurate at interpolation than extrapolation, it is desirable to include at least a few of these extreme data points in the dataset.

3.4. Weighting of errors for regression

When fitting regression models to the heat flux data and assessing their accuracy, some care is required to measure and weigh errors appropriately. The first issue is that the heat flux can span several orders of magnitude. If a model makes an absolute prediction error of 1 in gyro-Bohm units, this error is significant if the true heat flux is 0.5, but it is a relatively minor error if the true heat flux is 100. Moreover, the standard deviation (over time) of the heat flux within one simulation tends to scale approximately linearly with the mean heat flux, so high- $Q$ cases have higher uncertainty than low- $Q$ cases. For these reasons, it is natural to measure errors in a relative rather than absolute sense. Weighing errors in this relative sense can be achieved by performing regression on $\ln Q$ rather than $Q$ directly. In particular, models should be scored based on the mean squared error or coefficient of determination $R^2$ evaluated using $\ln Q$ . Thus, in the standard definition, $R^2 = 1-\left [\sum _j (y_j-f_j)^2\right ]/\left [\sum _j (y_j - \bar {y})^2\right ]$ , where $y_j$ are the true target values, $\bar {y}$ is their mean and $f_j$ are the predicted values, one would take $y_j$ and $f_j$ to be the true and predicted values of $\ln Q$ rather than of $Q$ itself.

However, an additional important consideration is that some simulations are stable. In a stable simulation, the heat flux is typically slightly positive but ${\ll} 1$ at the end of the computation due to decay of the initial perturbation and the exact value of $Q$ is not very meaningful. There is not a steady state to average over and the final $Q$ is influenced by the magnitude of the initial condition. If the heat flux from a simulation is $Q=10^{-12}$ and a regression model predicts it to be $10^{-4}$ , the model is accurately predicting that the system is stable (or nearly so), but scoring the model based on the error in $\ln Q$ would treat the prediction error to be large. In some stable simulations where $|Q|$ has very small magnitude, $Q$ may even be slightly negative, in which case, $\ln Q$ cannot even be evaluated.

There are several possible approaches to account for these considerations. One approach is to consider a classification problem for stability versus instability separately from the regression problem. In this case, the regression only needs to be evaluated if the classifier has first predicted instability, i.e. in the unstable region of parameter space, so the regression can be fit using only the subset of the data in which $Q$ exceeds some small threshold. Here, we adopt this approach with a threshold of 0.1 and our experience is that results are insensitive to the exact choice of threshold. With this approach, the training set does not include any points with $Q \lt 0.1$ or negative $Q$ , so it is straightforward to use $\ln Q$ as the regression target. The full dataset is still used for training the classifier. If the classifier is interpretable, it directly provides insight into the factors that determine the critical gradient.

Several other approaches are possible. As reported by van de Plassche et al. (Reference van de Plassche2020), a regression could be fit by minimising a custom objective that penalises errors in $Q$ (or $\ln Q$ ) only in the unstable region, while in the stable region, only penalising positive values of $Q$ , to prevent spurious prediction of instability. However, not all regression methods allow for a custom objective, e.g. nearest-neighbours. Or, regression could be performed on the full dataset with a standard objective after all values of $Q$ below a threshold are replaced by the threshold. While straightforward to implement, this method introduces non-smoothness in the predicted $Q$ ; this is not a problem for decision tree models, but it may introduce inaccuracies with smooth models like neural networks. A fourth option is to perform regression on the full dataset using the target quantity

(3.1) \begin{equation} \hat {Q} = \ln (1 + Q). \end{equation}

For $Q \gg 1$ , the quantity $\hat {Q}$ behaves like $\ln Q$ , while for $Q \ll 1$ , $\hat {Q}$ varies approximately linearly with $Q$ . Therefore, for a regression model that is fit by minimising the mean squared error in $\hat {Q}$ , or $R^2$ computed from $\hat {Q}$ , prediction errors will effectively be based on relative error in $Q$ for large $Q$ , but based on absolute error for small $Q$ , as desired. We have tried the first, third and fourth of these methods, and our experience is that they give very similar results. For the rest of the discussion, we use the first method.

4.1. Cyclic invariant neural network

Our goal is to develop a surrogate model based on a neural network that maintains invariance to cyclic permutations inherent to the gyrokinetic system. An analogy exists with computer vision research, where translation invariance is important as well: if a cat is present in an image, it should be recognised as a cat even if its location within a bitmap is shifted. A standard approach to achieving approximate translation-invariance in computer vision is through convolutional (as opposed to fully connected) neural networks. This type of neural network contains convolutional layers, in which the spatial data are convolved with a kernel, which is a translation-equivariant operation. Between convolutional layers, pooling layers then reduce the spatial resolution, typically by a factor of two, and a sufficient number of alternating convolutional and pooling layers converts the spatial data to approximately translation-invariant features. While computer vision applications use two-dimensional convolution and pooling operations, for the turbulence application here, we will use one-dimensional convolution and pooling.

During training, we further reinforce the translation-invariance by augmenting the dataset with a sufficient number of sequences that are randomly permuted (cyclically) in $z$ , allowing the network to learn patterns that are independent of any particular alignment with respect to $z$ . This strategy prevents the model from developing biases towards specific alignments, ultimately improving generalisation and stability.

Another key aspect of our surrogate model development is the use of ensemble learning, a powerful machine learning technique that aggregates predictions from multiple base models to improve overall performance. Ensemble methods have proven to be highly effective in reducing prediction variance, mitigating model biases and improving generalisation, making them well suited for applications involving complex, high-dimensional data such as ours. By leveraging an ensemble of multiple models, we aim to construct a predictive framework that is not only more accurate, but also more stable when encountering unseen data. The core advantage of ensemble learning is its ability to balance trade-offs between different models – some models may be more sensitive to specific features, while others may capture different patterns. By combining them, we create a more comprehensive representation of the underlying physical system.

To implement this strategy, we design our base model to be highly flexible, allowing variations in neuron count, layer depth, kernel shape, activation functions, etc. also known as hyperparameters. This flexibility is essential for tuning the model to achieve optimal predictive performance across different configurations. However, selecting the best architecture requires careful exploration of the hyperparameter space. To efficiently search for the best hyperparameters, we employ DeepHyper (Balaprakash et al. Reference Balaprakash, Salim, Uram, Vishwanath and Wild2018; Egelé et al. Reference Egelé, Guyon, Vishwanath and Balaprakash2023), an advanced hyperparameter optimisation tool designed for deep learning. DeepHyper allows us to automate the search for high-performing configurations, reducing manual tuning efforts and accelerating model development. We will discuss the details of this optimisation process in subsequent sections.

With these considerations in mind, we design our surrogate neural network using a structured approach that consists of three main components, shown in figure 7.

1. (i) Feature extraction with convolutional layers. The first component comprises $n$ consecutive layers of neural network blocks, each containing one-dimensional convolutional layers (Conv1D). These layers are responsible for extracting spatial features from the input sequences while preserving cyclic invariance. Each convolutional layer processes $x$ data channels and outputs $y$ number of transformed channels, which are subsequently normalised using batch normalisation to stabilise training and improve convergence. A max-pooling layer follows each block to downsample the extracted features, reducing computational complexity while retaining the most relevant information.

2. (ii) Global average pooling (GAP) for dimensionality reduction. To further condense the extracted feature maps, we integrate GAP. Unlike traditional pooling methods that operate on small local regions, GAP computes the average activation value for each feature map across all spatial dimensions. This not only reduces the number of parameters, but also ensures that the model focuses on high-level feature representations rather than localised details, making it particularly effective for cyclic data.

3. (iii) Fully connected layers for prediction. The final stage consists of $m$ fully connected linear layers that process the compact feature representation and map it to the target variable: predicted heat flux averages. These layers refine the extracted features and produce the final output by capturing complex, nonlinear relationships in the data. Since the temperature and density gradient are not functions of $z$ , these two inputs bypass the convolutional layers and feed directly to the first fully connected layer.

Table 1.

List of hyperparameters and their search ranges explored using DeepHyper.

Figure 7.

Surrogate model architecture to learn heat flux averages using a structured neural network. It consists of three main components: (a) feature extraction with Conv1D; (b) global average pooling; and (c) fully connected layers.

4.2. Hyper parameter optimisations and ensemble learning

To enhance the predictive performance of our surrogate model, we employ ensemble learning, which aggregates predictions from multiple base models. By combining different models, we create a more robust and comprehensive representation of the underlying physical system, improving accuracy and generalisation.

In our model, we carefully parametrise several key hyperparameters that influence performance. These include the kernel size of the 1-D convolution layers, the number of output channels in Conv1D layers, the number of neurons in fully connected layers, batch size, learning rate and scheduler patience for dynamically adjusting the learning rate. A summary of these hyperparameters is provided in table 1. We fix the number of Conv1D layers at five, so the associated pooling layers reduce the number of spatial points from 96 to $96/2^5=3$ , and include two fully connected layers. While these constraints simplify the search space, the total number of possible hyperparameter configurations remains extremely large, exceeding one hundred quadrillion ( $10^{17}$ ). Exploring this vast space through brute-force methods is computationally infeasible.

To efficiently navigate this large hyperparameter space, we leverage DeepHyper (Balaprakash et al. Reference Balaprakash, Salim, Uram, Vishwanath and Wild2018; Egelé et al. Reference Egelé, Guyon, Vishwanath and Balaprakash2023), an advanced hyperparameter optimisation tool. DeepHyper systematically finds the best set of hyperparameters that maximise model accuracy using different search strategies including Bayesian optimisation, evolutionary strategies and random search. Bayesian optimisation uses a probabilistic approach to refining hyperparameters, while evolutionary strategies optimise parameters through genetic algorithms. Random search serves as a baseline method. DeepHyper is designed for parallel execution in distributed computing environments such as clusters and supercomputers, making it effective for hyperparameter optimisation tasks at scale. In DeepHyper, Bayesian optimisation (BO) that we adopt follows a single-manager, multiple-worker scheme to explore and exploit the search space. The manager maintains a probabilistic surrogate model (e.g. Random Forest) to predict the performance of configurations and select promising candidates based on an acquisition function (e.g. expected improvement or upper confidence bound). The workers asynchronously evaluate these candidates in parallel, returning results to update the surrogate model iteratively. To improve robustness and generalisation, DeepHyper generates an ensemble of models by aggregating the top-performing configurations identified during the search.

In this study, we employ DeepHyper to explore the hyperparameter space of our base model, aiming to identify and select the top-N best models for ensemble learning. DeepHyper systematically searches the hyperparameter space to optimise model performance. To guide this optimisation, we use the coefficient of determination $R^2$ , as the objective function, ensuring the selected models achieve the highest predictive accuracy. The base neural network model is implemented using PyTorch.

Figure 8 presents the results of our hyperparameter search. Each point in the figure represents a model explored during the optimisation process. The x-axis denotes the time of completion, while the y-axis shows the model’s performance measured by the $R^2$ score. Over a period of approximately nine hours, we used 64 GPUs to conduct this search. DeepHyper evaluated 443 different models with varying hyperparameter configurations, systematically assessing their performance. After completing the search, we selected the top 100 models with the highest $R^2$ scores to form an ensemble, ensuring improved predictive accuracy and generalisation. The bottom plot presents a histogram of model sizes, represented by the number of parameters on the x-axis, for all 443 models explored by DeepHyper and the top 100 high-performing selected models, illustrating the distribution of model sizes before and after selection.

Figure 9 illustrates the prediction performance of this ensemble, which consists of the 100 highest-performing models from the DeepHyper optimisation, on the varied-gradient data. Each red dot in the figure represents the mean prediction of the ensemble at each target value in the test dataset, while the vertical bars indicate the range of $\pm 1\sigma$ (standard deviation) of predictions across the 100 models. We evaluated the ensemble using a total of 9785 held-out test samples, achieving an overall $R^2$ score of 0.989, computed by treating the ensemble mean as the prediction. Additionally, the neural networks exhibited efficient inference speed, taking a total 1031.0 s on a single NVIDIA A100 GPU to process 9785 test samples across the 100 models. To make a single prediction of the heat flux, this time corresponds to an average of 0.001 s for a single CNN and 0.1 s for the ensemble. These times represent a factor of ${\sim}4000$ speed-up for evaluating the ensemble compared with one of the gyrokinetic simulations used to generate the training data; for evaluating a single CNN, the speed-up factor over a gyrokinetic simulation is $4\times 10^5$ . These results demonstrate the effectiveness of our ensemble-based approach in providing accurate and stable predictions across the dataset while maintaining competitive inference performance.

Figure 8.

Results of the DeepHyper hyperparameter search. Each point represents a model with a unique hyperparameter configuration explored by DeepHyper, plotted against its completion time (x-axis) and performance score (y-axis). The search was conducted using 64 GPUs over approximately nine hours, evaluating a total of 443 models. The top 100 highest-performing models, selected for the final ensemble, are highlighted in red. The bottom plot shows the histogram of model sizes (number of parameters on the x-axis) for all 443 models explored by DeepHyper and the top 100 selected models.

Figure 9.

Prediction performance of an ensemble of the top 100 models selected from DeepHyper for the varied-gradient dataset. After exploring 443 models by DeepHyper, the top 100 were chosen based on their performance and evaluated against a test dataset of 9785 samples. Each dot represents the mean prediction of the ensemble, while the vertical bars indicate $\pm$ 1 standard deviation. The ensemble achieved an overall $R^2$ score of 0.989, demonstrating strong predictive accuracy and stability.

5.1. Feature extraction

Since the heat flux from the gyrokinetic model is invariant under periodic translation in $z$ , as discussed in § 2, a regression or classification model should preserve this invariance. However, this invariance would generally be broken if we use the raw features on grid points as inputs to regression or classification methods (other than CNNs) directly. For example, a regression could, in principle, find that the heat flux depends specifically on $|\boldsymbol{\nabla} x|^2$ at $z=0.3$ , which is not a translation-invariant quantity. Instead, the inputs to the regression or classification should be translation-invariant features that are derived from the raw $z$ -dependent geometric features. To ensure this invariance in the extracted features, we consider features that are the composition of equivariant operations with translation-invariant reductions. To explain this method, let $P$ denote the set of periodic real-valued functions on the interval $[0, 1)$ , representing the $z$ -domain. Each of the seven original raw features can be considered as an element of the set $P$ . An equivariant operation $E$ is a function from $P\to P$ , such that translating the input and then applying $E$ gives the same result as first applying $E$ and then translating. More formally, let $f$ be any element of $P$ and let $g:P\to P$ indicate a periodic shift by some distance $\varDelta$ , so $g(f)(x)=f(x + \varDelta \mod 1)$ . Then the equivariance of $E$ means that $E(g(f)) = g(E(f))$ for all $f$ and $\varDelta$ . A reduction is a map from $P \to \mathbb{R}$ . A reduction $A$ is translation-invariant if $A(f(x + \varDelta \mod 1)) = A(f(x))$ for all $x$ and any shift $\varDelta$ .

For the results here, we consider 23 possible translation-invariant reductions: max; min; range (max $-$ min); mean; median; mean of the square; variance; skewness; $L_1$ norm; quantiles 0.1, 0.25, 0.75 and 0.9; count above −2, −1, 0, 1 and 2; the absolute amplitudes of the three longest-wavelength Fourier modes; the expected $k_{||}$ ; and the largest amplitude $k_{||}$ . The last two quantities are defined as follows: letting $\sum _{j=-N}^N \hat {f}_j \exp (i k_j z)$ , denote the Fourier transform of a raw feature $f(z)$ , the expected $k_{||}$ is $\sum _{j=1}^N k_j |\hat {f}_j| / \sum _{j=1}^N |\hat {f}_j|$ , while the largest amplitude $k_{||}$ is $k_j$ , where $j$ is the solution of $\arg \max _j |\hat {f}_j|$ . These last two reductions were included due to the importance of parallel length scales in critical-balance models of turbulence (Barnes, Parra & Schekochihin Reference Barnes, Parra and Schekochihin2011). We indicate the set of these 23 reductions by $R$ .

For the equivariant operations, we consider combinations of unary operations and products. The unary operations, each an equivariant map $P \to P$ , include 11 operations on an input function $f$ : the identity $f$ , $|f|$ , ${d}f/{d}z$ , Heaviside $(f)$ , Heaviside $(-f)$ , ReLU $(f) = \max (f,0)$ , ReLU $(-f)$ , $1/f$ , $f^2$ , $fB$ and $f / B$ . The last two operations, multiplying or dividing by the field strength $B$ , were included since the coordinate Jacobian is $1/B$ .

We can now state the full set of features that are considered. Let $F$ denote a set of eight features, given by the original seven raw features together with the local shear $S = ({d}/{d}z)(\boldsymbol{\nabla} x \boldsymbol{\cdot} \boldsymbol{\nabla} y / \boldsymbol{\nabla} x \boldsymbol{\cdot} \boldsymbol{\nabla} x)$ . The local shear (Greene & Johnson Reference Greene and Johnson1968) is added here since it is plausible that it could play a role in determining the turbulence intensity. Then, the set of features $U(F)$ is obtained by applying each possible unary operation to each of the eight elements of $F$ . We let $C(U(F))$ denote the set of all possible pairwise products of functions in $U(F)$ , supplemented with $U(F)$ . Then, $U(C(U(F)))$ indicates all possible unary operations applied to all elements of $C(U(F))$ . The full set of features is finally $R(U(C(U(F))))$ , obtained by applying all of the reductions, for a total of just over 1 million extracted features. This set includes a few features that are duplicated, improper due to division by 0, or that include unnecessary operations (e.g. ReLU of a positive-definite function), but a suitably large fraction are distinct and finite. Of course, larger feature sets could be considered by including more operations and combinations, but the set here is suitably rich to find accurate regressions of the data.

5.2. Spearman correlation

One approach to measuring the potential importance of an isolated feature in a regression problem is to compute its Spearman correlation with the target quantity. Spearman correlation is defined by the Pearson correlation between the sorted rank of the target with the sorted rank of the feature. The absolute magnitude of the Spearman correlation has the appealing properties of being fast to compute and invariant to any monotonic nonlinear function, e.g. for any sequence of points $\{x_j\}$ , the Spearman correlation of $\{x_j\}$ with $\{\exp (x_j)\}$ is 1. We can evaluate the Spearman correlation between each feature from the previous subsection and the heat flux. Note that no regression or classification model is used. The most highly correlated features for the fixed-gradient dataset are listed in table 2. All these top features include the factor $\Theta (\boldsymbol{B}\times \boldsymbol{\kappa }\boldsymbol{\cdot} \boldsymbol{\nabla} y)$ , where $\Theta$ is the Heaviside function. In our sign convention, this quantity is 1 in regions of bad curvature. The most correlated features all have a similar form, indicating the heat flux is highest when the flux surface compression $|\boldsymbol{\nabla} x|$ is large in regions of bad curvature. For the features listed, the sign of the correlations is positive, so greater flux surface compression is associated with a higher heat flux, as expected physically. The features also include an inverse weighting with $B$ , which perhaps reflects the Jacobian $\propto 1/B$ . The exact powers of $|\boldsymbol{\nabla} x|$ and $B$ vary among the top features, but a consistent pattern is evident.

Table 2.

Geometric features from § 5.1 with highest absolute magnitude of Spearman correlation to the nonlinear heat flux at fixed temperature and density gradient. Here, $\Theta$ denotes the Heaviside function.

As mentioned in § 1, these features with highest Spearman correlation to the heat flux are consistent with ideas by Mynick et al. (Reference Mynick, Pomphrey and Xanthopoulos2010), Xanthopoulos et al. (Reference Xanthopoulos, Mynick, Helander, Turkin, Plunk, Jenko, Görler, Told, Bird and Proll2014), Stroteich et al. (Reference Stroteich, Xanthopoulos, Plunk and Schneider2022) and Goodman et al. (Reference Goodman, Xanthopoulos, Plunk, Smith, Nührenberg, Beidler, Henneberg, Roberg-Clark, Drevlak and Helander2024). In these earlier works, stellarator configurations were sought with smaller $|\boldsymbol{\nabla} x|$ in regions of bad curvature, motivated by the following physical intuition. At fixed ${\rm d}T/{\rm d}x$ , reducing $|\boldsymbol{\nabla} x|$ reduces the real-space temperature gradient $|\boldsymbol{\nabla} T| = ({\rm d}T/{\rm d}x)|\boldsymbol{\nabla} x|$ . Reducing $|\boldsymbol{\nabla} T|$ reduces the source of free energy for instabilities and the associated turbulence. This is particularly true if done in the regions where instabilities and turbulence are localised due to bad curvature. Positive correlation between $|\boldsymbol{\nabla} x|$ and quasilinear heat flux estimates was also noted by Jorge et al. (Reference Jorge, Dorland, Kim, Landreman, Mandell, Merlo and Qian2024).

Kendall’s $\tau$ is another correlation coefficient with many similarities to Spearman’s correlation. For all results in this paper involving Spearman correlation, essentially the same findings are obtained in comparable computational time if Kendall correlation is used instead, though numerical values of Kendall correlation are smaller.

Spearman and Kendall correlation cannot account for dependence of the target on multiple features that may interact, so we now proceed to more sophisticated methods.

5.3. Sequential feature selection

In FSFS (Efroymson Reference Efroymson1960; Whitney Reference Whitney1971), regression or classification is first performed using one feature at a time. In other words, if there are $n$ features, we fit $n$ distinct models, each with one feature. Any regression or classification method can be used. The feature that yields the best model fit to the data is selected to progress to the next step. Then, the data are fit using $n-1$ independent models that each use two features, the one selected in the first step plus all possible second features. Of these $n-1$ models, the one with closest fit is selected to progress to the next step and so forth. Thus, FSFS results in a parsimonious set of features, which are effectively ranked from most to least important by the order in which they are selected.

For the regression and classification models to use in FSFS, we primarily rely on the gradient boosted decision tree package XGBoost (Chen & Guestrin Reference Chen and Guestrin2016). This choice is due to its speed with datasets of this size. Speed is a priority since $O(10^6)$ independent models must be fit at each step. Like other decision tree methods, XGBoost fits a piecewise-constant function to the data, with the location and number of discontinuities chosen to balance accuracy of fitting the data against complexity of the surrogate. At each step of FSFS, we use an average score from five-fold cross-validation. When applying FSFS to the varied-gradient dataset, the temperature and density gradients are also included among the features that can be selected. Other scalar features such as $\iota$ , ${\rm d}\iota /{\rm d}x$ , ${\rm d}p/{\rm d}x$ and $n_{fp}$ can be included, but these features are not selected by the procedure, which makes sense because they do not appear in the gyrokinetic equation and hence have only a more indirect relation to the heat flux.

Figure 10 shows the results of FSFS for regression on both the fixed-gradient and varied-gradient datasets, as well as for stable versus unstable classification on the varied-gradient dataset. For regression, we assess the performance using the coefficient of determination $R^2$ . For classification, we choose features in FSFS using the log-loss measure (also known as cross-entropy; lower is better), while reporting also the accuracy and ROC-AUC scores (receiver operating characteristic area under the curve; higher is better) in figure 10(b). It can be seen that the prediction accuracy rapidly improves with the first few features, then levels off. With three features, the heat flux for the varied-gradient dataset can be predicted with $R^2=0.88$ ; adding nine more features results in a marginal improvement to $R^2=0.92$ . Qualitatively similar results are obtained using other regression or classification models. As an example, figure 10 also shows FSFS results for 10-nearest-neighbours regression; the behaviour is similar to the XGBoost regression but with slightly lower $R^2$ .

Figure 10.

Scores for regression (left) and classification (right) in forward sequential feature selection, showing improvement as the first few features are added. Each point shows the mean score on held-out data using five-fold cross-validation.

Another view on FSFS is provided by figure 11. Each panel shows, for a given number of features, how the predictions with that feature subset compare with the actual heat flux. In each panel, the regression was fit using 80 % of the data, and the figure shows the performance on the held-out 20 % of the data. A perfect prediction would correspond to all the points lying on the grey dotted line. As more features are added, the regression models more accurately predict the heat flux. There is a trade-off in that the models with more features are more complicated and harder to interpret due to possible interactions between the features.

Figure 11.

Accuracy of regression models improves to a point as more features are included from sequential feature selection. Each panel shows the performance of the XGBoost regression on 20 % held-out test data from the varied-gradient dataset.

For each of the FSFS curves in figure 10, the first five selected features are listed in table 3. In the table, absFFTCoeff1 indicates the absolute magnitude of the longest-wavelength (but not constant) Fourier mode in $z$ . In the varied-gradient dataset, for both regression and classification, the first selected feature is the temperature gradient and the second selected feature is the density gradient. These findings are consistent with expectations from linear theory that the temperature gradient is the primary drive for the instability and the density gradient can significantly affect stability. In both the varied-gradient and fixed-gradient cases, the first geometric feature selected again reflects the flux surface compression in regions of bad curvature. While the top geometric feature for the classifier involves unfavourable $\boldsymbol{\nabla} B$ drift rather than curvature drift, the feature $\mathrm{mean}(\Theta (\boldsymbol{B}\times \boldsymbol{\kappa }\boldsymbol{\cdot} \boldsymbol{\nabla} y)|\boldsymbol{\nabla} x|^2/B)$ has almost an identical score, log-loss $=$ 0.123 as opposed to 0.122, so there is little distinction between the two drifts in this case. Hence, the results are quite consistent with the Spearman correlation analysis of § 5.2, and support the physical intuition from Mynick et al. (Reference Mynick, Pomphrey and Xanthopoulos2010), Xanthopoulos et al. (Reference Xanthopoulos, Mynick, Helander, Turkin, Plunk, Jenko, Görler, Told, Bird and Proll2014), Stroteich et al. (Reference Stroteich, Xanthopoulos, Plunk and Schneider2022) and Goodman et al. (Reference Goodman, Xanthopoulos, Plunk, Smith, Nührenberg, Beidler, Henneberg, Roberg-Clark, Drevlak and Helander2024).

Table 3.

First five features from § 5.1 selected with forward sequential feature selection. Results are shown both for classification of stability vs instability, and for regression on the logarithm of the heat flux $Q$ . Results are also shown for both the gradient-boosted decision tree package XGBoost and for 10-nearest-neighbors (10NN). Here, $\Theta$ denotes the Heaviside function.

For each of the four tables within table 3, the second geometric feature again includes the flux surface compression $|\boldsymbol{\nabla} x|$ , but this time involving the radial rather than binormal component of the $\boldsymbol{\nabla} B$ drift, $\boldsymbol{B}\times \boldsymbol{\nabla} B\boldsymbol{\cdot} \boldsymbol{\nabla} x$ . Note that there is no distinction between the $\boldsymbol{\nabla} B$ drift and curvature drift in the radial direction: $\boldsymbol{B}\times \boldsymbol{\nabla} B \boldsymbol{\cdot} \boldsymbol{\nabla} x = B\boldsymbol{B}\times \boldsymbol{\kappa }\boldsymbol{\cdot} \boldsymbol{\nabla} x$ for any MHD equilibrium at any $\beta$ . So, the quantity $\boldsymbol{B}\times \boldsymbol{\nabla} B\boldsymbol{\cdot} \boldsymbol{\nabla} x$ appearing in the second geometric features is the geodesic curvature (times $B|\boldsymbol{B}\times \boldsymbol{\nabla} x|$ ). As mentioned in § 1, the geodesic curvature has been discussed recently as a correlate of turbulence (Xanthopoulos et al. Reference Xanthopoulos, Mischchenko, Helander, Sugama and Watanabe2011; Nakata & Matsuoka Reference Nakata and Matsuoka2022), motivated by the fact that geodesic curvature plays a prominent role in the gyrokinetic equation for zonal flow modes (Rosenbluth & Hinton Reference Rosenbluth and Hinton1998). More work would be needed to confirm that the high significance of geodesic curvature in our analysis is due to the effect of zonal flows on the turbulence, but the connection is at least suggestive.

At each step of FSFS, there are typically many features which are variations on a theme and which would give nearly the same score, similar to table 2. For instance, table 4 shows the top five features in steps three and four of FSFS for regressions on the varied-gradient dataset. At step 3, the top features are all based on the flux surface compression $|\boldsymbol{\nabla} x|$ in regions of bad curvature, weighted by various powers of the Jacobian. At step 4, the top features all involve the radial magnetic drift, weighted by the flux surface compression and powers of the Jacobian. Since the $R^2$ score is so similar among the various options at each step, we should not ascribe too much importance to the details that vary, such as the power of $B$ at step 3. However, aspects that are consistent among the top features, such as the consistent appearance of the radial magnetic drift (i.e. geodesic curvature) at step 4, are more likely to be physically meaningful.

Table 4.

Top-scoring features from steps 3–4 of forward sequential feature selection, for regression on the heat flux using the varied-gradient dataset with XGBoost. At each step, there are many features which are variations on a theme that have nearly identical $R^2$ score. Here, $\Theta$ denotes the Heaviside function.

Other regression and classification models besides XGBoost can achieve a reasonable fit to the data. Figure 12 shows a comparison of several regression methods for the heat flux, all using the same set of 12 features with the varied-gradient dataset ( $a/L_T$ , $a/L_n$ and the top 10 geometric features from FSFS). In addition to XGBoost and 10-nearest-neighbours (10NN), performance is also shown for three other models. One is the gradient-boosted decision tree package LightGBM (Ke et al. Reference Ke, Meng, Finley, Wang, Chen, Ma, Ye and Liu2017). Another is random forest regression. Lastly is linear regression after the features are transformed with the Yeo–Johnson power transform (Yeo & Johnson Reference Yeo and Johnson2000). Models other than XGBoost and LightGBM use the implementations from scikit-learn (Pedregosa et al. Reference Pedregosa2011). All models use the default hyperparameters from the relevant python package and the performance shown is the average score from five-fold cross-validation. While the prediction accuracy is highest for the decision-tree methods, it is meaningfully high for the other methods as well.

Figure 12.

Comparison of regression methods for the heat flux in the varied-gradient dataset. In all cases, the feature set is $a/L_T$ , $a/L_n$ and the top 10 geometric features from FSFS with XGBoost.

5.4. Shapley values

Another way to measure the importance of the features in regression and classification models is via Shapley values. A concept originating from game theory, Shapley values are a fair way to divide value among the input features, based on how the model performance degrades when that feature is removed, and considering all subsets of features (Shapley Reference Shapley1953; Lipovetsky & Conklin Reference Lipovetsky and Conklin2001). Here, we specifically use the Shapley Additive exPlanation (SHAP) formulation by Lundberg & Lee (Reference Lundberg, Lee, (ed., Luxburg, Bengio, Wallach, Fergus, Vishwanathan and Garnett)2017). Shapley values naturally provide the average sign of correlation between each feature and the target. For calculations here, we use the SHAP python package (Lundberg & Lee Reference Lundberg, Lee, (ed., Luxburg, Bengio, Wallach, Fergus, Vishwanathan and Garnett)2017; Lundberg et al. Reference Lundberg, Erion, Chen, DeGrave, Prutkin, Nair, Katz, Himmelfarb, Bansal and Lee2020), which allows for efficient computations with decision tree models.

Figure 13.

Distribution of Shapley values for regression with the varied-gradient dataset, using an XGBoost fit with top 12 features from FSFS. Features are listed in decreasing importance as measured by mean magnitude of the Shapley values.

Figure 13 shows the distribution of Shapley values for regression on the varied-gradient dataset, using the first 12 features from FSFS. Within each row, corresponding to one of the features, a histogram of the Shapley values is shown. The horizontal dimension gives the contribution to $\ln Q$ from that feature. Wider distributions therefore indicate more important features. The features are sorted in decreasing order of the average absolute Shapley value, i.e. decreasing importance. Each point is coloured by the value of the feature for that flux tube, scaled to the range of that feature over the dataset. So, if the distribution is purple/dark on the left and yellow/light on the right, larger values of that feature increase the predicted $Q$ . Conversely if the distribution is yellow/light on the left and purple/dark on the right, increasing values of the feature decrease $Q$ .

The order of feature importance as measured by Shapley values in figure 13 is not exactly the same as the order in which features were selected in FSFS, but the first three features are the same, and the pattern is broadly similar. Unsurprisingly, the most important feature is $a/L_T$ , which is found to increase $Q$ . The next most important feature is $a/L_n$ , which decreases $Q$ , as expected from linear theory. Consistent with the findings from Spearman correlation and FSFS, the most important geometric feature is an average of $|\boldsymbol{\nabla} x|$ in regions of bad curvature. Here, the next most important feature, mean $(|\boldsymbol{\nabla} x|^4/B^6)$ , also reflects the average flux surface compression, but now independent of bad curvature. As the distributions for these two features are blue on the left and yellow on the right, both of these features increase $Q$ , consistent with physical intuition. The next most important features, median $((\boldsymbol{B}\times \boldsymbol{\nabla} B\boldsymbol{\cdot} \boldsymbol{\nabla} x)^2|\boldsymbol{\nabla} x|^8/B^8)$ and absFFTCoeff1 $($ ReLU $(\boldsymbol{B}\times \boldsymbol{\nabla} B\boldsymbol{\cdot} \boldsymbol{\nabla} x)/B^5)$ , both involve the geodesic curvature, consistent with the pattern in table 3. The colours indicate that larger average magnitude of geodesic curvature increases $Q$ , coinciding with the theoretical prediction of Xanthopoulos et al. (Reference Xanthopoulos, Mischchenko, Helander, Sugama and Watanabe2011) and Nakata & Matsuoka (Reference Nakata and Matsuoka2022). Physically, larger geodesic curvature means greater neoclassical damping of zonal flows, hence smaller average zonal flow magnitude, and therefore higher turbulence intensity.

The second-most-important geometric feature according to Shapely values, mean $(|\boldsymbol{\nabla} x|^4/B^6)$ , was the fifth most important geometric feature according to FSFS order. Other differences appear in the ranking of feature importance between Shapley values and FSFS after the two geodesic curvature features. Since Shapely values and FSFS order are different measures of importance, there is no guarantee that they will rank the features in the same order, particularly farther down in the list where the effect sizes become smaller.

5.5. Fine-tuning the top feature

From the previous subsections, there is a consistent and robust finding that the most important geometric feature is the flux surface compression in regions of bad curvature. With this in mind, we can repeat the analyses of the previous subsections with a new set of extracted features that focuses on this general combination of quantities, giving more variations on this theme. Optimisation can also be used to tune constants and exponents appearing within a proposed functional form. For example, considering features of the form $\mathrm{mean}([\Theta (\boldsymbol{B}\times \boldsymbol{\kappa }\boldsymbol{\cdot} \boldsymbol{\nabla} y) + \alpha ] |\boldsymbol{\nabla} x|^\beta / B^\gamma )$ , the parameters $\{\alpha ,\, \beta ,\, \gamma \}$ can be optimised to maximise $R^2$ . Both approaches – trying a set of variations on a theme and optimisation – result in broadly similar conclusions: slightly improved prediction of $Q$ is found if a shift is added to the Heaviside function and the exponent on $|\boldsymbol{\nabla} x|$ is modified, yielding the feature

(5.1) \begin{equation} f_Q = \mathrm{mean}([\Theta (\boldsymbol{B}\times \boldsymbol{\kappa }\boldsymbol{\cdot} \boldsymbol{\nabla} y) + 0.2] |\boldsymbol{\nabla} x|^3 / B). \end{equation}

The shift of 0.2 to the Heaviside function effectively combines the top two geometric features according to the Shapley values in figure 13 (though with different powers of $B$ .) For this single feature with the fixed-gradient dataset, the Spearman correlation is 0.788 and XGBoost regression on $\ln Q$ yields $R^2 =$ 0.737. These values are slightly increased from 0.775 and 0.669, respectively, without the shift to the Heaviside function. For the varied-gradient dataset, using models with the three features $\{ a/L_T, a/L_n, f_Q\}$ , regression on $\ln Q$ gives $R^2=0.887$ , slightly improved from $R^2=0.876$ from table 3.

Figure 14 shows a comparison of the single feature $f_Q$ and the GX heat flux for the fixed-gradient dataset. Note that there is no regression model applied here. While (5.1) does not explain the full variation in the heat flux, it clearly does explain some. It is possible that with additional work, a single geometric feature of comparable complexity could be found with even higher predictive accuracy.

Figure 14.

Comparing the true heat flux and the single geometric feature $f_Q$ in (5.1) for fixed gradients, it is clear that there is significant correlation. No regression model is used here.

The set of three features $\{ a/L_T, a/L_n, f_Q\}$ does reasonably well in the classifier for predicting stability: log-loss = 0.128, accuracy = 0.944, ROC-AUC = 0.989. However, the geometric feature that best predicts the stiffness of the heat flux above marginal stability is not necessarily the best feature for predicting the stability boundary. Using the same methods described earlier in this subsection, a single geometric feature can be fine-tuned to marginally improve the stability classifier’s accuracy. Again, adding a shift to the Heaviside function and modifying the exponents, we arrive at

(5.2) \begin{equation} f_{\mathrm{stab}} = \mathrm{mean}([\Theta (\boldsymbol{B}\times \boldsymbol{\nabla} B\boldsymbol{\cdot} \boldsymbol{\nabla} y) + 0.4] |\boldsymbol{\nabla} x| / \sqrt {B}).\end{equation}

Using XGBoost with the three features $\{ a/L_T, a/L_n, f_{\mathrm{stab}}\}$ , the classifier achieves log-loss = 0.111, accuracy = 0.953 and ROC-AUC = 0.991.

Figure 15.

Comparing several proposed ITG objectives using three scores. The classification and regression scores are computed using XGBoost with three features: the single geometric feature, $a/L_T$ and $a/L_n$ .