yancc: A GPU-accelerated, differentiable solver for neoclassical transport in tokamaks and stellarators

Technical Summary of yancc: A GPU-accelerated, differentiable solver for neoclassical transport in tokamaks and stellarators

Problem Statement
Neoclassical transport theory describes the collisional transport of particles, heat, and momentum in magnetic confinement fusion devices (tokamaks and stellarators). Solving the drift kinetic equation (DKE) is essential for determining radial electric fields, bootstrap currents, and transport coefficients. However, numerical solution of the DKE is notoriously difficult due to strong advection dominance, recirculating flows, internal boundary layers (particularly between trapped and passing particle populations), severe anisotropy, and high dimensionality (four dimensions: two spatial, two velocity).

Existing solvers face a trade-off between speed and physical fidelity. Monoenergetic codes (e.g., DKES, MONKES) sacrifice interspecies coupling and energy scattering for performance. Full 4D solvers (e.g., SFINCS, NEO) offer high accuracy but suffer from high memory footprints and poor scalability to GPU hardware due to sparse matrix factorizations with irregular memory access. Furthermore, the lack of automatic differentiability in existing codes hinders gradient-based optimization and adjoint sensitivity analysis, which are increasingly required for stellarator design and uncertainty quantification.

Methodology
The authors present yancc, a new solver built on the JAX framework to leverage GPU acceleration and automatic differentiation. The code solves both the full 4D DKE (retaining speed-dependent collisions, energy scattering, and full interspecies coupling) and the reduced monoenergetic form.

• Discretization:

  • Speed ($x$): A collocation grid using weighted Maxwell polynomials (orthogonal with respect to $e^{-x^2}$), requiring only 5–10 points.
  • Pitch Angle ($\alpha$), Poloidal ($\theta$), and Toroidal ($\zeta$) Angles: Finite difference discretization.
  • Stencil Modification: To improve multigrid efficiency, the authors employ a modified upwind stencil (e.g., using points $\{0, 1, 4\}$ for 2nd order and $\{-2, 0, 1, 3, 4\}$ for 4th order) rather than standard upwind stencils. This increases the diagonal dominance of the resulting linear system matrices without significantly increasing the leading-order error constant.
  • Collision Operator: The field-particle scattering (Rosenbluth potentials) is handled analytically via a Green's function approach, solving the associated Poisson equations in velocity space using Legendre expansions and incomplete gamma functions, avoiding the need to solve for potentials as additional unknowns.

• Linear Solver:

  • The discretized DKE forms a large, sparse linear system $Af = s$.
  • Instead of sparse direct solvers (which fail to scale on GPUs), yancc uses a Krylov subspace method (GCROT with inner GMRES) preconditioned by a geometric multigrid scheme.
  • Multigrid Strategy:
    • Grid Hierarchy: Semi-coarsening is used; the speed coordinate ($x$) is kept at the finest resolution across all levels, while $\alpha$, $\theta$, and $\zeta$ are coarsened by a factor of 2.
    • Smoothing: A block-diagonal approximation of the operator is used, cycling through permutations that maintain full coupling along one coordinate at a time. The modified upwind stencils enhance the effectiveness of this block-diagonal smoother.
    • Coarse Grid Correction: Direct re-discretization is preferred over Galerkin projection to maintain low memory usage and avoid the high cost of matrix-free coarse-grid operations.
    • Cycle: A V-cycle (index 1) is used for the full DKE, while a W-cycle or higher (index 3) is found efficient for the monoenergetic case.

Key Contributions

1. GPU-Accelerated Full 4D Solver: yancc is the first solver to combine a full 4D DKE solution (including full interspecies coupling and energy scattering) with efficient GPU execution, overcoming the memory and scaling limitations of sparse-matrix-based codes like SFINCS.
2. Differentiability: Built in JAX, the solver is fully differentiable, enabling gradient-based optimization of magnetic geometry and adjoint sensitivity analysis without manual derivation of transport coefficients.
3. Algorithmic Innovations: The introduction of widened upwind stencils to enhance diagonal dominance and the use of semi-coarsening with block-diagonal smoothers tailored for the specific anisotropy of the DKE.
4. Analytic Rosenbluth Potentials: An efficient implementation of the field-particle collision operator using precomputed Green's functions and tensor contractions, avoiding the computational overhead of solving auxiliary Poisson equations on the grid.

Results and Benchmarks

• Accuracy: Benchmarks against MONKES (monoenergetic) and SFINCS (full 4D) show agreement within 1% across a wide range of collisionalities ($\nu^* \in [10^{-4}, 10^2]$), magnetic geometries (W7-X, NCSX-like), and plasma configurations (single and multi-species).
• Performance vs. SFINCS:
  • Speed: yancc achieves an order-of-magnitude speedup over SFINCS on a per-scan basis. While SFINCS runtime increases significantly at high collisionality (due to preconditioner degradation), yancc runtime remains nearly flat.
  • Memory: yancc uses an order of magnitude less memory (e.g., ~6 GB vs. >50 GB for SFINCS in tested scenarios), making it feasible for GPU-based workflows.
• Performance vs. MONKES: In monoenergetic benchmarks, yancc is 2–4 times faster than a JAX-reimplemented MONKES algorithm over most of the collisionality range, though MONKES may be faster for very simple geometries requiring low resolution.

Significance and Claims
The paper claims that the combination of speed, low memory footprint, and differentiability makes yancc uniquely suited for integration into modern stellarator optimization workflows, uncertainty quantification (via Monte Carlo sampling), and integrated profile prediction. Unlike current approaches that rely on precomputed databases or fit functions, yancc can compute neoclassical coefficients on-the-fly within transport solvers. The ability to compute gradients of transport coefficients with respect to input parameters enables efficient optimization over the full configuration space. The authors position yancc as a tool to bridge the gap between high-fidelity physics and the computational demands of design and analysis, noting that future work will include self-consistent radial electric field calculations and adaptive mesh refinement.