How does magnetic geometry affect ITG turbulence?�Insights from data & machine learning

M Landreman, J Y Choi, C Alves, P Balaprakash, R M Churchill , R Conlin, G Roberg-Clark

Thanks to many others who gave suggestions Supported by the US DOE StellFoundry SciDAC

Regression

Turbulence simulations

Feature importance

Roberg-Clark (2022)

N = 8

Mackenbach (2022)

N = 15

Proll (2016)

N = 10

TEM proxy

This work:

N = 100,705

Motivations

Profile prediction

Optimization

Understanding

Optimize geometry for maximum fusion power

Mandell (2024)

Regression

Turbulence simulations

Feature importance

Equilibria group 1: random rotating ellipses

Nfp,

aspect ratio,

elongation,

axis torsion,

and beta are all random.

All configurations have same minor radius & toroidal flux, so same gyroBohm normalization

Equilibria group 2: QUASR QA & QH (Giuliani 2024)

Random pressure added for even more diversity

Equilibria group 3: random boundary modes

RBC and ZBS boundary Fourier modes sampled from normal distributions, fit to 44 “real” stellarators

Regression

Turbulence simulations

Feature importance

Nonlinear turbulence simulations were run with GX in every equilibrium

• Electrostatic, adiabatic electrons.
• 1 simulation in each tube with random dT/dx and dn/dx.
• 1 simulation in each tube with (a/T) dT/dx = 3, (a/n) dn/dx=0.9
• 8 minutes to get heat flux on 1 GPU
• 2×10⁵ nonlinear simulations took < 7000 node-hours (1/8 allocation)

Don’t predict the time-dependence,

just the mean

Raw feature space: 7x 1D functions that enter the turbulence simulations

Turbulence simulations

Feature importance

Regression

Raw features should not be directly fed to classical regression or fully-connected neural network, since model should be translation-invariant

• GK equation, hence heat flux, is invariant under periodic translation of the raw features in z.

​

• Similar to computer vision, where convolutional neural networks give approximate translation-invariance.

​

Convolutional neural networks give accurate prediction of the turbulence

Convolutional layers x5

Fully connected layers x2

z-dependent geometric features

a/L_T

a/L_n

Heat flux Q

Prediction in 0.001 sec for single network, 0.1 sec for ensemble

Turbulence simulations

Feature importance

Regression

Our interpretable models use a large library of candidate features, all translation-invariant

Spearman correlation is a quick tool to find the most important feature

• Spearman correlation is the regular Pearson correlation of the the sorted rank of the target with the sorted rank of the feature.
• Its magnitude is invariant to any monotonic nonlinear function, e.g. corr(x, exp(x)) = 1
• No regression model required.
• Features with highest correlation to heat flux Q at fixed dT/dx & dn/dx:

Heaviside function: Where there is bad curvature,

local temperature gradient in real space (to various powers)

Jacobian (maybe squared)

Extremely similar to Mynick (2010), Xanthopoulos (2014), Stroteich (2022), Goodman (2024)!

|∇T| = (dT/dx) |∇x|

Forward sequential feature selection: ∼3 features can be almost as predictive as all features

Stiffness

Critical gradient

Sequential feature selection allows closer fit to the data as more geometric features are included

Performance shown on 20% held-out test data

Most important features from sequential feature selection

Classification (stability vs instability)

Regression on heat flux

Xanthopoulos et al (2011), Nakata & Matsuoka (2022):

Larger geodesic curvature (= radial drift) ⇒ Stronger damping of zonal flows ⇒ higher heat flux

At each step, the top features are variations on a theme

Regression for the random-gradient dataset

Shapley values show the sign and magnitude of each feature’s effect

Decreasing importance

Shapley value: contribution to predicted ln(Q)

The first geometric feature can be fine-tuned for even better fit

Fixed-gradient dataset.

​

No regression model used here.

Feature fine-tuned for stability classifier:

Previously proposed proxies can be tested

Multiple lines of evidence agree that the most important geometric feature is |∇ψ| in regions of bad curvature

• Highest Spearman correlation at fixed gradients.

​

• Consistently the first geometric feature chosen in sequential feature selection:

​

• In regression on the heat flux above the critical gradient

​

• And in the classifier for stability vs instability (i.e. determines critical gradient)

​

• Chosen by both XGBoost and nearest-neighbors.

​

• Also the largest Shapley values

There are many extensions possible

• Try larger sets of possible features
• From the gyrokinetic equation, understand how these features affect turbulence.
• Kinetic electrons, magnetic fluctuations.
• Saliency maps to understand the features learned by the neural networks.
• Symbolic regression.
• Kolmogorov-Arnold Networks.
• Optimization, profile prediction.
• Include & test other physics-motivated features.

​

​

Data will be released on Zenodo with the paper, so have a go at it!

Turbulence simulations

Feature importance

Regression

Extra slides

Other machine learning regression methods work also

All using a/L_T, a/L_n, and the top 10 geometric features selected via XGBoost

XGBoost regression model with 1 feature

XGBoost fit

Fixed-gradient dataset

XGBoost regression model using only a/LT and a/Ln

Same plot, showing data as dots

XGBoost regression model for fixed gradients using 2 features

Same plot, showing data as dots