MRX: A differentiable 3D MHD equilibrium solver without nested flux surfaces

TL;DR

This paper addresses the challenge of computing three-dimensional magnetohydrostatic equilibria without assuming nested flux surfaces, a capability essential for accurately modeling islands and chaotic field lines in toroidal plasmas.It introduces MRX, a differentiable, structure-preserving solver built on Finite Element Exterior Calculus and implemented in JAX, enabling robust 3D equilibria while preserving divergence-free fields and helicity under admissible variations.Key contributions include a discretization that preserves core geometric identities, a midpoint relaxation scheme with Leray projection and harmonic regularization, and comprehensive diagnostics and axisymmetric/stellarator test cases.The framework supports nonuniform meshes, island seeding, and differentiable optimization workflows, offering a path toward 3D equilibrium optimization and topology-aware plasma modeling with high performance on modern accelerators.

Abstract

This article introduces a new 3D magnetohydrodynamic (MHD) equilibrium solver, based on the concept of admissible variations of B, p that allows for magnetic relaxation of a magnetic field in a perturbed/non-minimum energy state to a lower energy state. We describe the mathematical theory behind this method, including ensuring certain bounds on the magnetic energy, and the differential geometry behind transforming to and from a logical domain and physical domain. Our code is designed to address a number of traditional challenges to 3D MHD equilibrium solvers, e.g. exactly enforcing physical constraints such as divergence-free magnetic field, exhibiting high levels of numerical convergence, dealing with complex geometries, and modeling stochastic field lines or chaotic behavior. By using differentiable Python, our numerical method comes with the additional benefits of computational efficiency on modern computing architectures, high code accessibility, and differentiability at each step. The proposed magnetic relaxation solver is robustly benchmarked and tested with standard examples, including solving 2D toroidal equilibria at high-beta, and a rotating ellipse stellarator. Future work will address the integration of this code for 3D equilibrium optimization for modeling magnetic islands and chaos in stellarator fusion devices.

Figures (14)

• Figure 1: A sketch of magnetic relaxation as a constrained minimization problem. In the configuration space of magnetic fields, the energy $\mathcal{E}$ is a quadratic functional. Starting from initial configuration $B_0$, admissible variations evolve along states of constant helicity (dotted) until they reach a stationary point (colored in orange), where the MHS equations hold.
• Figure 2: Left: Quadratic B-splines in one spatial dimension (dashed lines) and a function spanned by them (solid line). These functions are piece-wise quadratic and $\in C^1([0,1])$. Right: Derivatives of the functions on the left. They are piece-wise linear, $\in C^0([0,1])$, and crucially $\not \in C^1([0,1])$.
• Figure 3: Error convergence for the Poisson problem in 3D toroidal geometry.
• Figure 4: Optimization of the domain $\Phi_a(\hat{\Omega})$ to match the target spectrum of an ellipse with axes equal to $1.0$ and $0.6$. The four panels show different optimization iterates $\in \{0, 10, 50, 500\}$. For each iteration, we plot: The current 0th eigenfunction on $\Phi_a(\hat{\Omega})$, the current map $\theta \mapsto a(\theta)$ compared to the target function $\bar{a}(\theta)$, the spectra of the current and target shape, the relative error in eigenvalues, and the value of the objective function.
• Figure 5: Illustrating the maximum stable time-step size. The visible oscillations are the result of halving $\delta t$ when the Picard iterates diverge.

Theorems & Definitions (35)

• Definition 1: Function spaces
• Remark 2
• Remark 3
• Remark 4
• Definition 5: Push-forward
• Proposition 6: abraham_manifolds_1988, Theorem 6.4.4
• Proposition 7: Pull-backs preserve boundary conditions
• proof
• Example 8: Tokamak
• Example 9: Stellarator