Fast solvers for Tokamak fluid models with PETSC

This paper presents a novel semi-coarsening geometric multigrid solver implemented in PETSc for the M3D-C1 tokamak code, which addresses the convergence limitations of the existing block Jacobi preconditioner by leveraging the toroidal grid structure to achieve superior robustness and performance on complex magnetohydrodynamics models.

The Gist

Imagine you are trying to predict how a swirling, super-hot soup of charged particles (plasma) behaves inside a giant, donut-shaped machine called a tokamak. This machine is designed to create fusion energy, like the power of the sun. However, this "soup" is incredibly chaotic. If you try to calculate its movement step-by-step using a computer, the math gets so complicated that the computer gets stuck, or takes so long that the answer is useless by the time it arrives.

This paper is about building a smarter, faster calculator for this specific type of problem.

Here is the breakdown of what the authors did, using simple analogies:

1. The Problem: The "Traffic Jam" of Math

The computer code they use (called M3D-C1) tries to solve equations that describe how the plasma moves. To do this, it has to solve a massive puzzle millions of times.

• The Old Way (Block Jacobi): Imagine you have a huge map of a city with traffic jams. The old method was like asking a different person to fix traffic on just one street at a time, while ignoring the rest of the city. If the city is small, this works. But as the city gets bigger (more "planes" or slices of the donut shape), the people fixing the streets can't talk to each other fast enough. The traffic jams get worse, and the solution slows down or stops working entirely.
• The Specific Challenge: The plasma in these machines is "anisotropic." Think of it like a stack of paper. It's very easy to slide a piece of paper along the surface (easy direction), but very hard to push it through the stack (hard direction). The old math solver didn't understand this "stack of paper" structure, so it tried to solve the hard direction and the easy direction with the same clumsy method.

2. The Solution: The "Multigrid" Elevator

The authors built a new solver using a method called Multigrid (MG).

• The Analogy: Imagine you are trying to find a lost toy in a massive, multi-story mansion.
  • The Old Way: You check every single room, every single drawer, and every single corner on the ground floor before moving up. It takes forever.
  • The Multigrid Way: You first look at a miniature model of the whole mansion from a bird's-eye view. You quickly spot the general area where the toy is missing (the "coarse" grid). Then, you zoom in to a medium-sized map to narrow it down. Finally, you go to the actual room (the "fine" grid) to pick up the toy.
  • By solving the problem on the "big picture" levels first, the solver knows exactly where to look when it gets down to the tiny details. This makes it incredibly fast.

3. The "Secret Sauce": Semi-Coarsening

The authors realized that the tokamak is like a stack of 2D slices (poloidal planes) twisted into a 3D donut.

• They applied their "Multigrid Elevator" specifically to the stacking direction (the toroidal direction).
• Instead of trying to simplify the whole 3D mess at once, they kept the 2D slices detailed (because the shape of the tokamak wall is complex) but made the stack of slices simpler as they went up the levels.
• This is like taking a thick book and only reducing the number of pages while keeping the text on each page clear. It's a perfect fit for the shape of the machine.

4. The Results: Speed and Reliability

The team tested this new solver on two very difficult scenarios:

• Scenario A: The "Runaway Electron" (SPARC): This simulates a dangerous event where particles accelerate uncontrollably.
  • Result: The new solver was competitive with the old one on smaller setups and much faster on the largest, most complex setups. It solved the problem in fewer steps, saving time.
• Scenario B: The "Stellarator" (A different, more twisted machine): This geometry is even more twisted and irregular than a standard donut.
  • Result: The old solver completely failed and couldn't find an answer. The new Multigrid solver succeeded. It was robust enough to handle the twisted geometry that broke the old tool.

5. The Hardware: Using Supercomputers

They ran these tests on Perlmutter, one of the world's fastest supercomputers, which uses both powerful CPUs and GPUs (graphics cards).

• They found that while the "setup" (building the miniature models) was expensive, the actual solving was incredibly fast on the GPUs.
• They discovered that for the hardest problems, they needed to use a "heavy-duty" smoother (a specific math trick) to keep the solver from getting stuck, which required a bit more computing power but paid off in speed.

Summary

The paper claims that by understanding the specific "stack of paper" shape of fusion plasma machines, they created a new math tool (Multigrid) that:

1. Solves problems faster than the current standard method on large, complex simulations.
2. Doesn't crash on twisted, complex shapes where the old method fails.
3. Is a crucial first step toward making fusion energy simulations practical and fast enough to help design real power plants.

They did not claim this solves fusion energy itself, but rather that it provides the fast, reliable calculator needed to simulate the physics that will eventually lead to fusion energy.

Technical Summary

Technical Summary: Fast Solvers for Tokamak Fluid Models with PETSc

Problem Statement
The M3D-C1 code is a primary tool for simulating extended-magnetohydrodynamic (MHD) plasmas in fusion energy research, specifically addressing large-scale instabilities and disruptions in tokamaks and stellarators. The code employs a semi-implicit time integrator that requires solving numerous linear systems per time step. Among these, the implicit advance of the momentum equation is identified as the most challenging and least robust component.

The current production solver utilizes a Krylov method (FGMRES/GMRES) preconditioned by a Block Jacobi (BJ) method, where blocks group degrees of freedom on planes of constant toroidal coordinate. While effective for smaller configurations, the BJ solver's convergence degrades significantly as the number of toroidal planes increases due to the spectral properties of the preconditioned matrix. Furthermore, the BJ solver fails to converge entirely on non-axisymmetric stellarator configurations, limiting the code's applicability to complex fusion devices.

Methodology
This work develops a geometric multigrid (MG) solver tailored to the specific discretization and grid structure of M3D-C1. The methodology leverages the following structural characteristics:

• Grid Structure: M3D-C1 uses unstructured 2D grids in the poloidal plane (extruded around the torus) combined with a regular 1D grid in the toroidal direction.
• Semi-Coarsening: The solver employs geometric semi-coarsening specifically in the toroidal direction. This approach is chosen because the toroidal grid is structured, allowing for standard geometric MG techniques (2:1 coarsening) while leaving the complex poloidal plane uncoarsened for future work.
• Algebraic Interface: Although the application computes Jacobians on fine grids, the coarse grid operators are constructed algebraically using a Galerkin process ($A_{i+1} = P^T A_i P$), where $P$ is a 3-point stencil prolongation operator. This creates a geometric/algebraic hybrid solver with a purely algebraic interface compatible with PETSc.
• Smoothers: Two smoothers are investigated:
  1. Point-Block Jacobi (PBJ): Computes explicit dense inverses of diagonal blocks (size 36 per vertex). This is fast and memory-efficient, suitable for GPU acceleration.
  2. Block Jacobi (BJ): The existing production solver is repurposed as a smoother. It involves full LU factorizations on poloidal planes. While expensive, it is used as a post-smoother in a $V(0,1)$-cycle for the most difficult problems where PBJ stagnates.
• Implementation: The solver is implemented within the PETSc library, utilizing its Krylov solvers and abstract interfaces for linear algebra back-ends (including MUMPS and SuperLU for direct solves).

Key Results
The solver was evaluated on two test cases using the Perlmutter supercomputer (CPU and GPU nodes): a runaway electron model of a SPARC tokamak disruption and a quasi-axisymmetric stellarator model.

1. SPARC Tokamak (Scalability and Speed):

   • On configurations with 4 to 16 toroidal planes, the MG solver (using PBJ smoothing) converged in a single iteration when the preconditioner was refreshed, demonstrating excellent algorithmic scaling.
   • On the largest configuration (64 planes), the problem difficulty increased, requiring a tighter solver tolerance ($rtol = 10^{-7}$) and a more aggressive $V(6,6)$-cycle with frequent refreshes.
   • Performance: The MG solver achieved competitive performance compared to the BJ solver. On the 64-plane case, the MG solver reduced the total number of Krylov iterations by a factor of approximately 16x compared to the BJ solver. While the total solve time was roughly comparable on GPUs, the MG solver showed significantly faster solve times on CPUs.
   • Complexity: Theoretical analysis suggests the MG solver has an asymptotic work complexity of roughly 2x the BJ solver due to the geometric sum of work across grid levels. The observed performance indicates that coarse grid solves are currently performing poorly, suggesting room for optimization.

2. Stellarator (Robustness):

   • On a non-axisymmetric stellarator test case, the standard Block Jacobi solver failed to converge completely.
   • In contrast, the MG solver (using BJ as a smoother) converged successfully. This demonstrates a critical robustness advantage of the MG approach for non-axisymmetric geometries where the spectral properties of the matrix defeat the BJ preconditioner.

Significance and Claims
The paper claims to present the first step in applying multigrid methods to the complex, science-relevant MHD tokamak models within the M3D-C1 code. The primary contributions are:

• Demonstration of Feasibility: Proving that multigrid methods can provide a path toward faster, scalable, and robust solvers for highly anisotropic, ill-conditioned magnetized plasma models.
• Toroidal Semi-Coarsening: Establishing that semi-coarsening in the toroidal direction is a useful and effective first step for common tokamak codes with extruded grid structures.
• Robustness: Showing that MG solvers can solve problems (specifically stellarator disruptions) where the current production Block Jacobi solver fails.

Future Work and Limitations
The authors remain modest regarding the current state of the solver, noting several areas for improvement:

• Coarse Grid Performance: The current performance is limited by inefficient coarse grid solves, partly due to redundant parallelism requirements and MPI failures in direct solvers (MUMPS/SuperLU) on coarse grids.
• Full 3D Multigrid: The most impactful future optimization is the addition of 2D multigrid in the poloidal plane to create a full 3D MG solver. This would require handling unstructured geometric coarse grids.
• Smoothers: Future work includes optimizing block sizes for PBJ smoothers to better fit GPU thread groups and investigating FieldSplit preconditioners that treat the three decoupled MHD waves separately.
• Refresh Strategies: The current refresh model (periodic LU factorization) is simple; smarter algorithms are needed to handle quasi-steady states more efficiently.

In summary, this work establishes a viable, robust multigrid framework for M3D-C1 that outperforms the existing solver in terms of iteration counts and robustness on difficult geometries, while identifying specific algorithmic and implementation hurdles to be addressed to achieve full textbook multigrid efficiency.