⚛️ nuclear experiments

Millisecond-Scale Neural Operator Surrogates for Double-Null Free-Boundary Grad-Shafranov Equilibria

This paper demonstrates that a geometrically conditioned Fourier Neural Operator can serve as a highly accurate, millisecond-scale surrogate for free-boundary Grad-Shafranov equilibria, achieving speedups of up to 640 times over traditional solvers while maintaining physics-consistent precision in predicting poloidal flux and plasma boundary locations.

Original authors: Plamen G. Krastev (Harvard University)

Published 2026-08-07

📖 4 min read🧠 Deep dive

CC BY 4.0

Original authors: Plamen G. Krastev (Harvard University)

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine trying to keep a swirling, super-hot ball of gas—hotter than the sun itself—floating in mid-air without it touching the walls of its container. This is the dream of nuclear fusion, the process that powers stars and promises clean, limitless energy for Earth. To make this happen, scientists use giant magnetic cages called tokamaks. But here's the tricky part: the gas doesn't just sit there; it squirms, stretches, and changes shape constantly. To keep it stable, scientists need to know the exact shape of the magnetic "bubbles" holding it every single millisecond.

The math behind these magnetic bubbles is called the Grad–Shafranov equation. Think of it like a incredibly complex recipe that tells you exactly how the magnetic forces balance out. The problem is, calculating this recipe is slow and heavy, like trying to solve a giant jigsaw puzzle while running a marathon. If you need to solve it thousands of times a second to control the machine, or to test thousands of different shapes to find the best one, the old way of doing it is just too slow. It's like trying to drive a race car using a map that takes an hour to draw. Scientists have been looking for a "shortcut"—a way to predict the magnetic shape instantly without doing all the heavy math every time.

This is where a new study comes in, offering a clever solution using artificial intelligence. The researchers trained a special type of AI, called a "neural operator," to act like a super-fast crystal ball. Instead of solving the slow, heavy math puzzle from scratch every time, this AI learned the patterns of the magnetic shapes. It was fed thousands of examples of how the magnetic fields look under different conditions, like different amounts of pressure or electric current. Once it learned the patterns, it could predict the entire magnetic shape in the blink of an eye.

The results are impressive. The AI, running on a standard computer chip, could predict the magnetic shape in about 26 milliseconds. If it runs on a powerful graphics card (like the ones used for gaming), it's even faster, taking just 2.77 milliseconds. To put that in perspective, the old method took about 1.7 seconds to do the same job. That means the AI is roughly 640 times faster on a graphics card and 69 times faster on a regular processor. It's the difference between waiting for a slow, old dial-up internet connection and having lightning-fast fiber optics.

But speed isn't the only thing that matters; the prediction has to be accurate. The researchers checked this carefully. They found that the AI's predictions were incredibly close to the real math, with an error of only about 0.05%. To visualize this, imagine the magnetic field is a giant map. The AI's map was so precise that the "X" marks where the magnetic field lines cross (called X-points) were off by less than the width of a pencil eraser (about 0.2 cm). Even the center of the magnetic bubble was pinpointed to within 0.03 cm.

The team also made sure the AI wasn't just guessing. They checked if the AI's predictions actually followed the laws of physics by running a quick test to see if the magnetic forces balanced out correctly. The AI passed this test just as well as the slow, traditional method did. This suggests that the AI isn't just memorizing answers; it has genuinely learned the rules of the game.

However, there are some limits to this magic trick. The AI was trained on a specific type of magnetic shape (a "double-null" shape) and a specific machine design. It's like a master chef who can make the perfect pizza but hasn't learned how to make sushi yet. If you ask it to design a completely different shape or a different machine, it might get confused. The researchers are clear that this is a tool for a specific job right now, not a universal fix for every fusion problem.

In the end, this paper shows that we can trade slow, heavy calculations for fast, smart predictions. By using this AI "surrogate," scientists can now run simulations and control systems much faster than before. It opens the door to testing thousands of new designs in the time it used to take to test one, and it could help keep the magnetic cages stable in real-time. While it's not a magic wand that solves everything, it's a giant leap forward in making the dream of fusion energy a practical reality.

Technical Summary: Millisecond-Scale Neural Operator Surrogates for Double-Null Free-Boundary Grad–Shafranov Equilibria

Problem Statement
The Grad–Shafranov (GS) equation governs the ideal magnetohydrodynamic (MHD) equilibrium in axisymmetric tokamak plasmas. Accurate and rapid evaluation of these equilibria is critical for stability analysis, transport modeling, and plasma shape control. However, in free-boundary scenarios—where the plasma boundary is not known a priori but must emerge self-consistently from plasma profiles and external fields—solving the GS equation requires iterative nonlinear numerical algorithms (specifically Picard iteration). These solvers, such as FREEGS, exhibit sample-dependent convergence latency and high computational costs. This makes them prohibitive for many-query settings, including uncertainty quantification, outer-loop optimization, integrated modeling, and real-time control workflows where equilibria must be evaluated repeatedly across families of operating points without the benefit of warm-started iterations. While machine learning (ML) has been applied to GS equilibria, prior work has largely focused on inverse reconstruction from diagnostics or fixed-boundary formulations, rather than learning a forward map for constrained free-boundary design problems.

Methodology
The author proposes a geometry-conditioned Fourier Neural Operator (FNO) to learn a computationally efficient representation of a constrained forward map. The model learns the mapping from:

Spatial coordinates $(R, Z)$ .
Scalar operating parameters: axis pressure ( $P_{axis}$ ), plasma current ( $I_p$ ), and a vacuum toroidal-field factor ( $f_{vac}$ ).
Prescribed X-point locations (upper and lower).

To the full poloidal-flux field $\psi(R, Z)$ .

Data Generation:

Solver: Data was generated using the open-source FREEGS solver for a single fixed machine geometry (FreeGS TestTokamak).
Configuration: A controlled family of constrained double-null free-boundary equilibria. The X-point targets were fixed at $(1.1, \pm 0.6)$ m with symmetric geometric jitter applied to preserve topology.
Parameters: Operating parameters were drawn uniformly from defined ranges ( $P_{axis} \in [200, 3000]$ Pa, $I_p \in [5\times10^4, 4\times10^5]$ A, $f_{vac} \in [0.5, 3.0]$ ).
Dataset: A base pool of 5,000 converged equilibria was generated. Scaling studies utilized nested subsets ( $N_{train} \in \{500, 1000, 2000, 5000\}$ ). A 100% acceptance rate was achieved across 6,000 total candidates (training, validation, and test).

Model Architecture:

Input: Nine channels on a $65 \times 65$ grid: spatial coordinates, broadcast scalar parameters, and broadcast X-point coordinates.
Architecture: A 2D FNO with four Fourier layers, retaining 16 modes per dimension and a hidden channel width of 64 (approx. 4.77M parameters).
Training: Trained using AdamW with mean squared error on normalized flux fields. Early stopping and learning rate scheduling were based on validation relative $L_2$ error.

Key Contributions and Results

Accuracy: The best model ( $N_{train}=5000$ ) achieved a mean relative $L_2$ error of 0.052% and a physical RMSE of $1.54 \times 10^{-5}$ Wb.
Scaling: Test error followed an empirical power law $\epsilon_{rel} \propto N^{-0.68}$ over the range $N_{train} \in \{500, 1000, 2000, 5000\}$ , indicating systematic improvement with dataset size.
Geometric Fidelity:
- Separatrix: Mean closest-point deviation of 0.072 cm; median Hausdorff distance of 0.967 cm.
- Critical Points: Both X-points recovered within $\sim 0.1\text{--}0.2$ cm (upper: 0.112 cm, lower: 0.161 cm). The magnetic O-point was localized to within 0.031 cm.
- Error Structure: Errors were spatially localized near the X-point pinch and separatrix, where small flux errors geometrically amplify. No topological failures (e.g., spurious islands) were observed.
Physics Consistency: The predicted fields satisfied an external finite-difference GS residual diagnostic with a mean normalized residual of 2.29, indistinguishable from the $2.29 \pm 0.06$ baseline obtained on the ground-truth FREEGS fields using the same diagnostic.
Computational Performance:
- Latency: The surrogate evaluates a single equilibrium in 2.77 ms on GPU and 25.6 ms on CPU.
- Speedup: This corresponds to speedups of $\sim 640\times$ (GPU) and $\sim 69\times$ (CPU) relative to the iterative FREEGS solver configured in the experiments.
- Determinism: The FNO exhibits near-deterministic latency ( $p95/\text{median} = 1.01$ ), contrasting with the variable convergence latency of the iterative solver ( $p95/\text{median} = 1.13$ ).

Significance and Claims
The paper claims that neural-operator surrogates can provide accurate, geometrically precise, and millisecond-scale equilibrium evaluations for magnetic-confinement fusion workflows, provided the topology and machine geometry are prescribed.

Paradigm Shift: The work demonstrates that constrained free-boundary equilibrium problems can be recast as operator-learning problems. Once trained, the neural operator replaces the conventional iterative nonlinear Picard scheme with a single deterministic forward evaluation.
Applicability: This approach is distinct from diagnostic reconstruction (inferring equilibrium from measurements) or actuator-conditioned forward models (predicting X-points from coil currents). Instead, it serves as a surrogate for constrained free-boundary design problems, where the target X-point geometry is prescribed or measurable. This makes it suitable for simulation-based design, optimization, and control workflows where equilibria at arbitrary operating points must be evaluated rapidly and independently.
Limitations: The author explicitly states that the results are scoped to a single machine geometry and a single double-null topology. The model is not a general actuator-level forward solver, nor does it currently handle cross-device transfer, resolution transfer, or coupled transport evolution. The reported speedups are relative to the specific FREEGS configuration used and may vary with optimized production solvers, though the qualitative contrast between iterative and deterministic inference is expected to hold.

In summary, the study validates that FNO surrogates can deliver physics-consistent, high-fidelity equilibrium solutions with predictable, millisecond-scale latency, addressing a critical bottleneck in real-time control and large-scale optimization for fusion research.

Technical Summary: Millisecond-Scale Neural Operator Surrogates for Double-Null Free-Boundary Grad–Shafranov Equilibria

More like this