Generative modeling for fusion: a large-scale benchmark of stellarator fast-ion simulations

Misha Padidar; Teresa Huang; Mark Goldstein; Michael Czekanski; Jeffrey Soules

cs.LGOct 27, 2026

Generative modeling for fusion: a large-scale benchmark of stellarator fast-ion simulations

Authors: Misha Padidar, Teresa Huang, Mark Goldstein, Michael Czekanski, Jeffrey Soules

Organizations: Center for Computational Mathematics, Flatiron Institute · Department of Statistics and Data Science, Cornell University

Abstract

Stellarators offer a promising magnetic confinement approach for achieving controlled nuclear fusion. However, stellarator design remains challenging, as evaluating configuration performance requires computationally expensive, high-fidelity simulations. In particular, key performance metrics derived from guiding center simulations are costly to include directly in the design optimization loop, and are almost exclusively used as diagnostics. Machine learning has the potential to advance fusion research by enabling rapid generation and evaluation of new stellarator configurations. Effectively applying ML for fusion demands large-scale datasets that capture both design diversity and high-fidelity simulation outputs. In this work, we introduce the Benchmark of Stellarator Simulations (BoSS), a dataset of stellarator designs paired with guiding center simulation results and trajectories. The BoSS is the first large-scale dataset of guiding center simulations, containing simulation data for 180,000 distinct magnetohydrodynamic equilibria spanning a wide range of magnetic geometries and confinement properties. To exemplify the utility of ML to accelerate stellarator design, we formulate and address two sets of modeling tasks: generative modeling for designing stellarators with good particle confinement, and surrogate modeling for predicting stellarator heat deposition patterns and particle trajectories. The tasks reveal the potential impact ML can have in the stellarator field and important opportunities to improve the applicability of ML models in the sciences. Our baselines generate diverse, performant stellarators and provide useful surrogate predictions, and point to long-horizon trajectory forecasting and out-of-distribution generation as key directions for future work. The BoSS can be accessed through https://boss.flatironinstitute.org/.

Figures & tables

Panel of Figure 1 — Figure 1: A) The electromagnetic coils and plasma for stellarator- 0010198 from the QUASR database, as seen from above. The surface is colored by the strength of the magnetic field. B) A slice of the $n_{\text{\tiny{fp}}}=2$ -field period device, showing the nested magnetic surfaces, Boozer coordinates $s,\theta,\varphi$ , as well as standard cylindrical coordinates $R,Z$ . The red curve indicates a hypothetical magnetic field line. C) Segments of four particle trajectories (colored curves) when viewed in a cross-section in Boozer coordinates. The circle marks the start of the trajectory and the square marks the end. The four circles represent four magnetic surfaces with labels $s=0.1,0.25,0.5,1.0$ , the outermost being the plasma boundary. One particle intersects with the $s=1$ boundary and is considered lost.

Figure 2: Pipeline for data generation. Devices are selected from candidate datasets such as QUASR and Constellaration, assigned profile functions, and rescaled. VMEC is then used to compute the ideal MHD equilibria before running guiding center simulations with CATAPULT. Metrics are computed from the geometry, magnetic field, and guiding center simulation results.

Panel of Figure 3 — Figure 3: (A-D) (top) top-down view of four devices taken from the BoSS , colored by the magnetic field strength; (bottom) cross-sections of each device at three fixed values of cylindrical $\phi$ . (E) The fraction of particles lost by time $t$ for each device (A, B, C, D).

Figure 4: (A) Efficiency gain of the rejection sampler over uniform random sampling of the the BoSS test set when using three proxy metrics for $e_{\text{loss}}$ : $e_{A}$ and prediction error in $\overline{\iota},\max B$ . Markers show the thresholds used in active learning, $\tau=0.5\%$ (circle) and $1\%$ (square) at which the efficiency gain is roughly $20\%$ and $40\%$ . The band is a 95% bootstrap confidence interval, and the dashed line marks no gain. (B) Distribution of the loss fraction prediction error $e_{\text{loss}}$ on the test set and on the active learning devices, by collection method. The distribution of high prediction errors highlights the efficiency of the method. Errors below $0.01$ pp ( $4.1\%$ , $12.5\%$ , and $5.4\%$ of each group, respectively) are not shown; these are devices whose predicted and actual loss fractions are both saturated near $0$ or $1$ . The units pp indicate percentage points (difference in percentages). (C) The distribution of $|\overline{\iota}|$ and cross-section elongation for training set devices and ative learning devices.

Figure 5: (A-D) (top) top-down view of four devices generated using active learning, colored by the magnetic field strength; (bottom) cross-sections of each device at three fixed values of cylindrical $\phi$ . Devices A,B were generated with rejection sampling, while devices C,D were generated with the optimization approach.

Figure 6: Basic properties for devices in the train and test sets: A) loss fraction for devices born on

s=0.1

; B) the square root of the quasisymmetry error; C) the aspect ratio; D) the mean rotational transform; E) the plasma-

\beta

for non-vacuum devices; F) and the number of field periods.

Metric	Diffusion	SurrogateOpt
Vacuum ( $\beta=0$ )	$N=417$	$N=97$
Low-loss rate- $10\%$ [%] $\uparrow$	43.6	52.6
Low-loss rate- $15\%$ [%] $\uparrow$	48.4	64.9
Mirror constraint [%] $\uparrow$	10.1	20.6
Mirror constraint- $1.1r_{\text{\tiny{mirror}}}^{*}$ [%] $\uparrow$	97.4	41.2
Aspect error (median) [%] $\downarrow$	0.3	3.7

Table 1: Inverse design at ARIES-CS scale, over successfully evaluated designs; metrics are defined in Section 6.1 . Bold = winner per row. Diffusion is conditioned on the property vectors of held-out test devices whose loss fraction is

\leq 10\%

and evaluated after rescaling to

a_{\text{\tiny{minor}}}=1.7

\mathrm{m}

, $\langle B\rangle=5.86$\mathrm{T}

; its

r_{\text{\tiny{mirror}}}^{*}

and

A^{*}

are the conditioning values of each source device, whereas SurrogateOpt uses

r_{\text{\tiny{mirror}}}^{*}=1.5

and

A^{*}\sim U[3.5,10]

Figure 7: Five designs generated by Diffusion that have $<10\%$ losses ( $5.8\%,4.5\%,4.9\%,0.0\%,0.7\%$ from left to right), mirror-ratio less than $1.65$ , and are far from the training set (normalized distance $\geq 0.17$ minor radii to its nearest training boundary, Appendix A ). The $n_{\text{\tiny{fp}}}=3$ device is a finite-pressure equilibrium, the others are vacuum. Top: boundary surface at the ARIES-CS scale coloured by $|\mathbf{B}|$ . Bottom: $|\mathbf{B}|$ in Boozer coordinates on the $s=0.24$ surface. The $n_{\text{\tiny{fp}}}=3,5,6$ devices display a helical symmetries, while the $n_{\text{\tiny{fp}}}=2$ device displays poloidal variation.

Metric	IID	PointNet	$\Delta$ %
Headline
ED (median)	0.0306	0.0358	+16.9
MC floor	0.019	0.019	-
HF $16\times 16$ (median)	0.3447	0.3824	+10.9
MC floor	0.187	0.187	-
HF $32\times 32$ (median)	0.5720	0.6134	+7.2

Table 2: Evaluation of IID vs PointNet for stellarators from the test set using particles born on

s=0.1,0.25,0.5

. HF and the loss-fraction errors are over the entire test set (

N=74763

); ED needs a lost-particle cloud and is restricted to devices with at least

10

lost particles (

N=72512

). The headline results show that the IID model outperforms the PointNet model in terms of median energy distance and heat flux. Lower is better for every reported metric. Bold = winner per row. The Monte Carlo floor (MC floor) denotes the median value of the metric under sampling noise variation when using the perfect predictor. Almost reaching the MC floor, the IID model demonstrates practical performance.

Panel of Figure 8 — Figure 8: Heat deposition prediction in a single stellarator: (top 3 rows) lost particle point cloud position and energy (color) for born on surfaces $s=0.10,0.25,0.50$ ; (bottom 3 rows) lost particle density per surface $s=0.10,0.25,0.50$ . Columns from left to right show the actual device, the nearest device to target from the training set (we compute the nearest device by first filtering devices with $[0.9\langle B\rangle,1.1\langle B\rangle]$ , followed by finding the nearest neighbor in the Euclidean norm of the flattened $\bm{\chi}_{s}$ vector), the prediction from the IID baseline, and the prediction from the PointNet baseline.

Figure 9: Example predictions of the three trajectory models over time for distinct trajectories (rows) across the five channels (columns). Trajectory A is a regular single frequency trajectory over the horizon, while B and C display multiple frequencies, and D undergoes a phase transition. The vertical dashed line indicates the training horizon for the autoregressive models (AR-MLP, AR-Transformer).

Metric	Stationary	AR-MLP	AR-Transformer	Direct-MLP
nRMSE ( $\tau=10\,\mu$ s)	1.019	0.261	0.378	0.416
nRMSE ( $\tau=100\,\mu$ s)	1.042	0.846	0.807	0.761
nRMSE ( $\tau=999\,\mu$ s)	1.038	1.046	1.094	0.862
Loss fraction MAE	0.183	0.104	0.118	0.137
Confinement time $W_{1}$ [ $\mu$ s]	159.8	79.4	105.2	115.5

Table 3: Trajectory prediction metrics on the test set. Stationary predicts the initial condition,

\hat{\mathbf{z}}(t)=\mathbf{z}_{0}

, and so never loses a particle. Standard errors of each metric, computed via bootstrap subsampling, are below

0.0004

for nRMSE and loss fraction MAE and below

0.3\,

\mathrm{\SIUnitSymbolMicro s}$$ for

W_{1}

. Lower is better for every metric. Bold = best per row.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 10: Graphic illustrating the normalized distance. Frame A shows the

\phi=0

cross section of three surfaces: a reference surface (black), a surface with distinct shape (orange), and a surface with distinct aspect ratio (violet). Frame B shows the normalized cross sections between which the distance is calculated.

Figure 11: Diffusion at the ARIES-CS scale: the

588

designs behind the Diffusion column of Table 1 (all converged,

586

traced), by

n_{\text{\tiny{fp}}}

and vacuum (blue) / finite pressure (green). Panels A and B: traced loss fraction to

t_{\max}=0.1

\sec

and $1$\mathrm{ms}

. Panel C: distribution of mirror ratio; diamonds indicate the median mirror ratio condition. Panel D: distance to the nearest training boundary ( Section A.2 ). The

n_{\text{\tiny{fp}}}=1

and

n_{\text{\tiny{fp}}}=6

cells hold

23

and

7

designs.

Figure 12: (Top) Five new devices generated by Diffusion, $n_{\text{\tiny{fp}}}=2$ – $6$ from left to right; (bottom) the contours of $|\mathbf{B}|$ in Boozer coordinates on $s=0.24$ ; thin black lines are field lines. The devices have loss fraction $18.5\%,21.3\%,8.7\%,16.3\%,43.5\%$ at the ARIES-CS scale. The third device ( $n_{\text{\tiny{fp}}}=4$ ) is a finite-pressure equilibrium ( $\beta=3\%$ ), the others vacuum.

Figure 13: Error between achieved and conditioned properties for all

4{,}420

converged test-conditioned designs at native scale, by

n_{\text{\tiny{fp}}}

and vacuum (blue) / finite pressure (green): absolute error in loss fraction (

4{,}369

traced designs), and relative error in mirror ratio and aspect ratio. The vertical scale is set by a few outliers; medians are

0.026

2.0\%

and

0.5\%

Figure 14: Five stellarators found by surrogate optimization, colored by $|\mathbf{B}|$ (top row). (Bottom) Contours of $|\mathbf{B}|$ in Boozer coordinates.

Figure 15: Statistics for 209 solutions of the surrogate optimization problem with predicted loss fraction

\leq 10\%

t_{\max}=0.1

\sec

. Panels A and B show the actual loss fraction to $0.1$\sec

and

1

\mathrm{ms}$$ , showing that half of all solutions for

n_{\text{\tiny{fp}}}=4,5

have under

10\%

losses while having reasonable mirror ratios (panel C). The horizontal dashed line indicates target values for each metric. Panel D shows the minimum (normalized) distance between the new surfaces and surfaces in the training set. Most surfaces are fairly distinct, with distance exceeding 20%(of the minor radius).

Figure 16: Architecture of the neural network surrogate models for predicting scalar outputs such as loss fraction, energy loss, aspect ratio, etc.

Headline
Metric	IID	PointNet	$\Delta$ %
ED (median)	0.0306	0.0358	+16.9
MC floor	0.0193	0.0193	-
HF $16\times 16$ (median)	0.3447	0.3824	+10.9
MC floor	0.187	0.187	-
HF $32\times 32$ (median)	0.5720	0.6134	+7.2

Appendix

Table 4: Evaluation of IID vs PointNet for stellarators from the test set using particles born on

s=0.1,0.25,0.5

. HF and the loss-fraction errors are over the entire test set (

N=74763

); ED needs a lost-particle cloud and is restricted to devices with at least

10

lost particles (

N=72512

). The subsequent sections show the breakdown of ED and HF when only considering particles born on a given initial surface; or for the subset of configurations within specific loss fraction ranges. Lower is better for every reported metric. Bold = winner per row.

Figure 17: Model architecture for the IID and PointNet velocity models used in the conditional flow matching models for heat flux prediction.

Figure 18: Normalized test RMSE across each prediction channel for the trajectory prediction task. The vertical dashed line indicates the training horizon of the autoregressive models (AR-MLP, AR-Transformer). The two baselines “stationary” and “device mean” predict the initial condition, and the mean value of channel, respectively, for each device. No model is able to reliably predict any channel to

1

\mathrm{ms}$$ .

Figure 19: Trajectory model architectures.

Figure 20: Normalized test RMSE across each prediction channel for the trajectory prediction ablation. The vertical dashed lines indicate the training horizons

K_{\max}=10

and

K_{\max}=100

. The two baselines “stationary” and “device mean” predict the initial condition, and the mean value of channel, respectively, for each device. The

K_{\max}=100

models reliably outperform the

K_{\max}=10

models beyond a lead time of roughly

30

Metric	Stationary	AR-MLP $K_{\max}=10$	AR-MLP $K_{\max}=100$	AR-Transformer $K_{\max}=10$	AR-Transformer $K_{\max}=100$
nRMSE ( $\tau=10\,\mu$ s)	1.019	0.223	0.261	0.371	0.378
nRMSE ( $\tau=100\,\mu$ s)	1.042	0.905	0.846	1.026	0.807
nRMSE ( $\tau=999\,\mu$ s)	1.038	1.107	1.046	1.170	1.094
Loss fraction MAE	0.183	0.133	0.104	0.111	0.118
Confinement time $W_{1}$ [ $\mu$ s]	159.8	91.0	79.4	92.0	105.2

Appendix

Table 5: Training-horizon ablation, comparing the performance of the autoregressive models trained with

K_{\max}=10

and

K_{\max}=100

. Lower is better for every metric. Bold = best per row.

Generative Modeling

Generative modeling learns data distributions to create new samples, from images and text to molecules, simulations, and structured records.

See all 509 papers

A proof-of-concept for automated AI-driven stellarator coil optimization with in-the-loop finite-element calculations
Mar 16, 2026Alan A. Kaptanoglu, Pedro F. Gil
Courant Institute of Mathematical Sciences, New York University, New York, NY 10012, USA · Max-Planck-Institut für Plasmaphysik, 85748 Garching, Germany
AI Surrogate Modeling for Real-Time Tokamak Equilibrium Prediction: Benchmarking Neural Architectures and Validation on EXL-50U
Aug 24, 2026Guoyang Shi, Zitong Zhang, Siqi Ding +5Neural Surrogate Modeling PDE Surrogate Modeling
Beijing ENN Fusion Energy Science and Technology Co., Ltd., Beijing 101111, China · Beijing Key Laboratory of High Magnetic Field Spherical Torus Fusion Energy, Beijing 101111, China · Hebei Key Laboratory of Compact Fusion, Langfang 065000, China +3
IGNITE Tokamak World Model Architecture
Oct 1, 2026Peter Steiner, Azarakhsh Jalalvand, Nathaniel Chen +4World Model Learning Physics-Based Simulation
Department of Mechanical and Aerospace Engineering Princeton University Princeton, NJ · Princeton Plasma Physics Laboratory Princeton, NJ · Princeton University, Princeton Plasma Physics Laboratory Princeton, NJ

Mar 16, 2026J/K move · Enter open · S save

A proof-of-concept for automated AI-driven stellarator coil optimization with in-the-loop finite-element calculations

Alan A. Kaptanoglu, Pedro F. Gil

Courant Institute of Mathematical Sciences, New York University, New York, NY 10012, USA · Max-Planck-Institut für Plasmaphysik, 85748 Garching, Germany

Open paper AskPDF arXiv

Finding feasible coils for stellarator fusion devices is a critical challenge of realizing this concept for future power plants. Years of research work can be put into the design of even a single reactor-scale stellarator design. To rapidly speed up and automate the workflow of designing stellarator coils, we have designed an end-to-end ``runner'' for performing stellarator coil optimization. The entirety of pre and post-processing steps have been automated; the user specifies only a few basic input parameters, and final coil solutions are updated on an open-source leaderboard. Two policies are available for performing non-stop automated coil optimizations through a genetic algorithm or a context-aware LLM. Lastly, we construct a novel in-the-loop optimization of Von Mises stresses in the coils, opening up important future capabilities for in-the-loop finite-element calculations.

Abstract

Generative Modeling

Explore similar work

A proof-of-concept for automated AI-driven stellarator coil optimization with in-the-loop finite-element calculations

AI Surrogate Modeling for Real-Time Tokamak Equilibrium Prediction: Benchmarking Neural Architectures and Validation on EXL-50U

IGNITE Tokamak World Model Architecture