🔬 physics

AGNI: A differentiable MHD stability solver & optimizer for magnetic confinement fusion devices

This paper introduces AGNI, a GPU-accelerated, differentiable solver that extends ideal MHD stability optimization to finite toroidal mode numbers, enabling efficient gradient-based design of tokamaks and stellarators against interchange instabilities.

Original authors: Rahul Gaur, Sanket Patil, Prateek Gupta, Djin Patch, Tony Qian

Published 2026-08-04

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Rahul Gaur, Sanket Patil, Prateek Gupta, Djin Patch, Tony Qian

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 the universe's most ambitious kitchen experiment: trying to cook a star inside a bottle. Scientists call this "magnetic confinement fusion." Instead of a pot and a stove, they use invisible, super-strong magnetic fields to hold a soup of super-hot gas (plasma) that is millions of degrees hot. If they can keep this soup from splashing out and cooling down, they could unlock a nearly limitless source of clean energy. But there's a catch. Just because you can build a magnetic bottle doesn't mean the soup inside will stay calm. Like a shaken soda can, the plasma can develop violent, invisible ripples that grow until the whole experiment collapses. For decades, scientists have had tools to check if their magnetic bottles are stable, but those tools were like looking at the ocean from a satellite: they could see the big waves, but they missed the tiny, dangerous ripples that could sink the ship.

Enter AGNI, a new digital tool created by researchers at the University of Wisconsin–Madison and the Indian Institute of Technology-Delhi. Think of AGNI as a high-speed, super-smart "what-if" machine for fusion reactors. In the past, checking for these tiny ripples was like trying to find a needle in a haystack by looking at one grain of hay at a time; it took forever and required massive computers. AGNI changes the game by using a special kind of math that lets it calculate how the plasma reacts to any change in the reactor's shape instantly. It's like having a video game where you can tweak the design of the magnetic bottle, press a button, and immediately see if the plasma inside will stay calm or explode, all while running on powerful graphics cards (GPUs) that make the calculations lightning-fast.

The Problem: The Invisible Ripples

To understand why AGNI is a big deal, we need to look at how fusion reactors work. They trap plasma in a donut-shaped (or twisted donut-shaped) magnetic cage. Scientists need to make sure this cage is stable. If the plasma gets wobbly, it can crash into the walls of the reactor, stopping the fusion process.

For a long time, scientists had two types of tools to check for wobbles:

The "Big Wave" Check: These tools looked at huge, slow-moving instabilities. They were good, but they missed the smaller, faster ripples.
The "Infinite" Check: Some tools assumed the ripples were so small they were basically infinite in number. This made the math easier but wasn't realistic for actual reactors.

The missing piece was a tool that could check for finite ripples—the specific, real-world wobbles that happen in actual machines like tokamaks and stellarators. Without this, designing a perfect reactor is like trying to build a bridge without checking if the wind will make the cables vibrate.

The Solution: AGNI's Superpowers

The authors present AGNI (Analysis of Global Normal modes in Ideal MHD). It's a computer program designed to be a "differentiable" solver. That's a fancy math word that essentially means: "If you change the shape of the reactor, AGNI can tell you exactly how the stability changes, without having to start the whole calculation over from scratch."

Here is how AGNI works, using some fun analogies:

The Digital Twin: AGNI takes a digital model of a fusion reactor (specifically a "stellarator," which looks like a twisted pretzel) and simulates the physics of the plasma inside.
The Pixelated Ocean: Instead of trying to solve the physics equations for every single atom (which is impossible), AGNI breaks the plasma down into a grid of points, like pixels on a screen. It uses "differentiation matrices" to figure out how the plasma moves at each pixel.
The Speedster: The real magic is that AGNI is built on a software framework called JAX, which is designed for artificial intelligence. This allows the program to run on GPUs (the same chips that power video games). While older computers might take hours to check one design, AGNI can do it in seconds.
The Gradient Guide: Because AGNI is "differentiable," it doesn't just say "This design is bad." It says, "This design is bad, and if you twist this part of the magnetic field just a tiny bit to the left, it gets 10% better." This allows scientists to use computers to automatically "optimize" the reactor shape, hunting for the most stable design like a hiker following a trail downhill to the lowest point.

What They Found

The researchers didn't just build the tool; they tested it to make sure it actually works.

The Test Drive: They compared AGNI against NIMSTELL, a well-known, trusted code that simulates plasma physics. They used a specific, complex magnetic shape called a "quasi-helically symmetric equilibrium."
The Match: AGNI and NIMSTELL agreed perfectly. Both found the same most dangerous ripple (a mode where the plasma twists 4 times around the short way and 4 times around the long way). They calculated the speed at which this ripple would grow, and the numbers matched to a high degree of precision.
The Speed Test: When they ran the numbers on a supercomputer, AGNI on a GPU was roughly 10 times faster than running on a standard CPU (the brain of a normal computer). For example, calculating the stability of a complex design took about 20 seconds on a GPU, compared to over 13 minutes on a CPU for the same resolution.
The "What-If" Test: They also tested AGNI's ability to give advice. They changed the size of the reactor slightly and asked, "How does the stability change?" AGNI's automatic calculations matched perfectly with the traditional, slower method of checking small steps one by one. This proves AGNI can be used to automatically design better reactors.

The Incompressible Trick

One tricky part of plasma physics is that the gas can be squished (compressible) or stay the same size (incompressible). Most dangerous ripples are incompressible. AGNI had to figure out how to simulate this without getting bogged down in heavy math.

The team tried two methods:

The Projection Method: A mathematically strict way to force the gas to be incompressible. It was accurate but slow and memory-hungry.
The "Speed of Sound" Trick: They found a clever shortcut. By pretending the "speed of sound" in the plasma was incredibly fast (by adjusting a number called the adiabatic constant, $\Gamma$ ), the math naturally forced the simulation to act like an incompressible fluid. This method was much faster and worked perfectly with their optimization tools.

Why This Matters

The paper concludes that AGNI is a major step forward. It is the first tool of its kind that is GPU-accelerated, automatically differentiable, and capable of handling finite-n (real-world size) instabilities for complex 3D reactors.

While the paper doesn't claim to have built a working fusion reactor yet, it provides the essential "flight simulator" engineers need. Before, designing a fusion reactor was like trying to navigate a maze in the dark. AGNI turns on the lights, showing the designers exactly where the walls are and how to twist the path to avoid them. With this tool, scientists can now rapidly test thousands of reactor shapes to find the one that keeps the star inside the bottle stable, bringing us one step closer to clean, limitless energy.

Technical Summary of AGNI: A Differentiable MHD Stability Solver & Optimizer

Problem Statement
The existence of an ideal MagnetoHydroDynamic (MHD) equilibrium in magnetic confinement fusion devices (tokamaks, stellarators, mirrors) does not guarantee stability. Finite toroidal mode number ( $n$ ) instabilities, such as interchange and ballooning modes, degrade plasma performance. While differentiable stability optimization tools exist, they have historically operated only in the infinite- $n$ limit. There is a critical need for a reactor design tool capable of solving and optimizing against finite- $n$ ideal MHD instabilities across multiple confinement concepts. Existing 3D stability solvers often lack differentiability or efficient integration with gradient-based optimization frameworks.

Methodology
The authors present AGNI (Analysis of Global Normal modes in Ideal MHD), a GPU-accelerated, automatically differentiable finite- $n$ ideal MHD stability solver and optimizer built on the JAX framework.

Discretization: AGNI discretizes the ideal MHD energy principle (the variational eigenvalue problem for plasma displacement $\tilde{\xi}$ ) using a pseudospectral approach in real space. It employs differentiation matrices and geometric coefficients derived from a DESC equilibrium. The energy integral is expressed as a quadrature over a collocation grid defined by radial, poloidal, and toroidal coordinates.
Coordinate System: The solver utilizes a nested flux surface coordinate system $(\rho, \theta, \zeta)$ . To handle singularities near the magnetic axis, the perturbation vector is rescaled. The radial basis uses a Legendre-Lobatto grid with a mapping function to cluster points near regions of interest (e.g., where eigenfunctions peak), while poloidal and toroidal directions utilize Fourier bases.
Eigenvalue Problem: The discretization yields a generalized Hermitian eigenvalue problem $A\xi = \lambda B\xi$ , where $A$ is the potential energy matrix and $B$ is the kinetic energy matrix. The eigenvalue $\lambda$ is proportional to the square of the growth rate $\gamma^2$ .
Numerical Solvers:
- Cholesky Decomposition: The kinetic energy matrix $B$ is block-diagonal, allowing efficient Cholesky decomposition to transform the generalized problem into a standard Hermitian eigenvalue problem $\hat{A}v = \lambda v$ .
- Shift-Invert Lanczos: To avoid computing the full spectrum (which includes dense stable and continuum modes), AGNI uses a shift-invert technique combined with the Lanczos algorithm. This targets the most unstable modes (largest $\lambda$ ) efficiently.
- Matrix-Free Options: The paper discusses matrix-free solvers to reduce memory footprint, though the default implementation uses dense LU factorization for robustness.
Differentiability: By leveraging JAX, AGNI computes reverse-mode gradients of the growth rate with respect to boundary shape and profile parameters without re-solving the equilibrium. This enables direct gradient-based optimization.
Incompressibility: Two methods are presented to impose the incompressibility condition ( $\nabla \cdot \xi = 0$ $\nabla \cdot ξ = 0$ ):
1. Projection Matrix: Projects out compressible modes using a dense Cholesky decomposition of the Gram matrix (high accuracy but computationally expensive for optimization).
2. Speed of Sound Penalty: Artificially increases the adiabatic constant $\Gamma$ (effectively increasing the speed of sound) to penalize compressible modes. This method is compatible with automatic differentiation and optimization.

Key Results and Benchmarks

Validation: AGNI was benchmarked against NIMSTELL, a non-linear MHD solver reduced to a linear ideal-MHD initial-value problem. Using a modified Landreman–Buller–Drevlak quasi-helically symmetric (QH) equilibrium, AGNI successfully recovered the dominant $m=n=4$ interchange mode. The growth rates and eigenfunction structures showed strong agreement between the two codes.
Gradient Verification: Automatic differentiation (AD) gradients were verified against central finite differences (FD). The AD gradients matched FD results with high precision (relative error $< 0.002$ ) and demonstrated robustness even in regions where the eigenvalue landscape was non-smooth or force balance was slightly violated.
Performance:
- Speed: GPU execution provides a significant speedup over CPU. For eigenvalue and gradient evaluation, GPUs are at least an order of magnitude faster. For example, a gradient calculation on a $32 \times 32 \times 12$ grid took 11.41 seconds on an NVIDIA A100 GPU compared to 265 seconds on an AMD EPYC CPU.
- Scalability: The GPU implementation allows for faster convergence in optimization loops, though memory constraints (VRAM) limit the maximum resolution compared to CPU runs.
Finite-Precision Limits: The study quantified the limits of resolving near-marginal eigenvalues. In double precision, eigenvalues with magnitudes below $\approx 10^{-10}$ (relative to the operator norm) are indistinguishable from numerical noise.

Significance and Claims
The paper claims that AGNI represents the first differentiable, GPU-accelerated, finite- $n$ ideal MHD stability solver and optimizer intended to span multiple confinement concepts. Its primary significance lies in:

Enabling Optimization: By providing reverse-mode gradients of stability metrics with respect to design parameters, AGNI allows for the direct optimization of tokamaks, stellarators, and mirrors against ideal MHD instabilities.
Finite- $n$ Capability: It fills the gap left by existing tools that are limited to the infinite- $n$ limit, addressing the specific finite- $n$ instabilities that degrade performance in stellarators and limit pedestal pressure in tokamaks.
Efficiency: The combination of JAX, GPU acceleration, and shift-invert methods makes high-resolution stability analysis and gradient-based optimization computationally feasible.

The authors note that future work will focus on developing efficient preconditioners for matrix-free solvers to further reduce memory usage, incorporating external (vacuum) modes to ensure full stability, and extending the formalism to equilibria with anisotropic pressure for open configurations like magnetic mirrors.

The Problem: The Invisible Ripples

The Solution: AGNI's Superpowers

What They Found

The Incompressible Trick

Why This Matters

More like this