We are releasing the code for our recent paper on Fourier Neural Operators (FNOs) applied to magnetohydrodynamic turbulence. The repository contains a PyTorch implementation of a FNO trained to act as a surrogate for 2D MHD simulations of the Orszag–Tang vortex, a canonical plasma physics benchmark. Given a short initial window of simulation frames plus two physical parameters (kinematic viscosity $ν$ and Ohmic diffusivity $η$ ), the model predicts the future evolution of plasma fields—density, velocity, and magnetic field components—at 25× the speed of a traditional numerical solver.

<h2 id=what-are-mhd-simulations>What Are MHD Simulations?</h2><p>Magnetohydrodynamics (MHD) is the theoretical framework that describes plasma as a conducting fluid coupled to electromagnetic fields. A plasma is an ionized gas where free electrons and ions interact collectively through long-range electromagnetic forces. We find plasma in stellar atmospheres, accretion flows around black holes, the interstellar medium. And inside nuclear fusion experiments.</p>

The MHD equations are a system of nonlinear PDEs expressing conservation of mass, momentum, and energy, coupled with Maxwell's equations through the induction equation for the magnetic field:

Continuity equation
$\frac{\partial ρ}{\partial t} + \nabla \cdot (ρ 𝐯) = 0$

Momentum equation $ρ [\frac{\partial 𝐯}{\partial t} + (𝐯 \cdot \nabla) 𝐯] = - \nabla P + \frac{1}{μ_{0}} (\nabla \times 𝐁) \times 𝐁 + \nabla \cdot 𝐓$

Induction equation $\frac{\partial 𝐁}{\partial t} = \nabla \times (𝐯 \times 𝐁 - η \nabla \times 𝐁)$

Energy equation $\frac{\partial e}{\partial t} + \nabla \cdot (e 𝐯) = - P \nabla \cdot 𝐯$

Here, $μ_{0}$ is the vacuum permeability, $𝐓$ is the viscous stress tensor, $η$ is the magnetic diffusivity and $e$ is the internal energy density.

The nonlinearity of these equations — particularly the coupling between flow and magnetic field — makes analytic solutions rare and numerical simulations expensive. Modern MHD codes such discretize the domain on a grid and evolve these equations forward in time, often requiring thousands of CPU/GPU hours for a single parameter configuration. Parameter sweeps across viscosity and resistivity ranges multiply this cost further.

<h2 id=the-orszagtang-vortex-a-canonical-benchmark>The Orszag–Tang Vortex: A Canonical Benchmark</h2><p>The Orszag–Tang vortex, introduced by Orszag & Tang in 1979, is one of the most widely used standard tests in computational MHD. Its setup is deceptively simple: a 2D doubly-periodic domain initialized with sinusoidal velocity and magnetic field profiles,

𝐯_{0} = (- \sin 2 π y, \sin 2 π x),

𝐁_{0} = \frac{1}{\sqrt{4 π}} (- \sin 2 π y, \sin 4 π x)

at uniform density and pressure. From these smooth initial conditions, the system rapidly develops turbulence, shocks, and shock–shock interactions. Vortices form, merge, and cascade energy across scales. The magnetic field organizes into complex filamentary structures, and the competition between kinetic and magnetic energy dissipation—controlled by $ν$ and $η$ —shapes the entire long-term evolution.</p>

This richness makes the Orszag–Tang vortex physically interesting for several reasons. It is a compressible system transitioning toward supersonic MHD turbulence. The problem captures simultaneously the physics of MHD discontinuities, turbulent energy cascades, and dissipation — all phenomena that appear in real astrophysical settings.

From a code validation standpoint, the Orszag–Tang vortex is invaluable precisely because so much is known about it analytically and numerically. If a new MHD solver produces the right shock structures, the right energy dissipation curves, and the right spectral slopes, it has passed a demanding test of its robustness. The vortex is a built-in test case in FARGO3D, the GPU-accelerated code we used to generate the training data.

<h2 id=the-codebase>The Codebase</h2><p>The repository implements a 3D Fourier Neural Operator (FNO) surrogate that learns to map initial plasma states to future states, conditioned on the physical parameters of the simulation.</p>

Data generation. We ran 50 FARGO3D simulations of the Orszag–Tang vortex on a 128×128 grid, each for 1,000 timesteps, sampling $η$ and $ν$ across the range $[10^{- 5}, 0.05]$ . Forty-eight simulations were used for training and validation; two were held out entirely for testing ( $ν = η = 5 \times 10^{- 5}$ and $ν = η = 3 \times 10^{- 4}$ ). The model is trained on the first $~ 160$ frames ( $\approx 0.7 t_{A}$ where $t_{A}$ is the Alfvén time) and asked to predict frames 160–1000 ( $\approx 0.7 - 4 t_{A}$ )— using only 20% of the simulation to forecast the remaining 80%.

Model architecture. The FNO3d model takes as input a block of shape (batch, 128, 128, 10, 7): 5 temporal snapshots of one field component plus viscosity ν and diffusivity η as the last two channels. Internally, spatial and temporal grid coordinates are appended, the 10-channel tensor is lifted to a hidden width of 30, and five Fourier layers process the data in the spectral domain. Each Fourier layer runs two parallel paths — a spectral convolution (3D real FFT → truncate to 64×64×5 modes → complex multiply → inverse FFT) and a pointwise 1×1×1 convolution — summed and passed through a GELU activation. After unpadding and projection through two linear layers, the output is (batch, 128, 128, 10): the predicted next 10 timesteps. Separate model instances are trained for each physical field (density, v_y, v_z, B_y, B_z, B_r).

Training. Models train for 10,000 epochs using a combined MAE + relative L2 loss, with Adam optimization.

Performance. On held-out test cases with viscosity/diffusivity values never seen during training, the model achieves mean-squared error of ~6×10⁻³ for the velocity field and ~10⁻³ for the magnetic field. Energy spectra and dissipation rates are reproduced within 96% accuracy. Compared to a UNet baseline, the FNO reduces error by 97%. The FNO predicts 800 frames in 8 seconds versus 198 seconds for FARGO3D on the same GPU—a 25× speedup.

A note on inference: the current implementation is teacher-forced, meaning the model receives ground-truth simulation frames as input for every prediction window rather than feeding its own predictions back as inputs. This means reported metrics reflect an upper bound on performance; true autoregressive rollout — where predictions accumulate over time — remains a goal for future work.

magfield Figure 1: Comparison between ground truth simulated data of MHD vortex (left panel) and FNO prediction (center) using this codebase including residuals (right panel). Colors indicate the magnetic field strength. Time in units of the Alfven time for the computational box.

Importance

Physics-based simulations of magnetized plasmas are expensive. Even with GPU-accelerated codes, a thorough exploration of parameter space demands substantial computational resources. Neural operator surrogates offer a data-driven shortcut: train once on an ensemble of simulations, then infer at a fraction of the cost.

FNOs are particularly well-suited to this problem because they operate in the Fourier domain, naturally aligning with the spectral structure of MHD turbulence. They are mesh-independent in principle, and they can capture long-range spatial correlations that local convolutional architectures miss, which is exactly why they can outperform UNets so dramatically here while UNets tend to smear out the turbulent structures.

The generalization result is the most scientifically exciting finding: the model accurately predicts plasma dynamics for viscosity/diffusivity combinations that were never part of the training set. This suggests the FNO is learning something about the underlying operator structure of the MHD equations, not just memorizing specific trajectories.

More broadly, this work contributes to an emerging toolkit for accelerating astrophysical simulations with machine learning. The same approach—training a neural operator surrogate on an ensemble of PDE solutions, then using it for rapid parameter sweeps—is applicable to protoplanetary disk dynamics, accretion flows, stellar convection, and other MHD-dominated systems where physical intuition is rich but simulation time is the bottleneck.

<h2 id=what-comes-next>What Comes Next</h2><p>The code is public, but the repository is not yet fully inference-ready. I am actively working to make it so. If you are interested in running the model or reproducing the results, check back soon. Planned additions include:</p>

Training dataset: making the full simulation ensemble publicly available so others can reproduce and extend our results
Inference guide with test dataset: a step-by-step walkthrough for running the pre-trained model on the held-out test cases
Link to raw FARGO3D simulation data and conversion scripts
Docker image: a containerized environment for fully reproducible training and inference
Autoregressive rollout: implementing the feedback loop where predicted frames are fed back as inputs, enabling true long-horizon forecasting and assessment of error accumulation over time

<h2 id=paper-and-code>Paper and Code</h2><p>The full details are in our paper:
Duarte, Nemmen & Lima-Santos (2025). Spectral Learning of Magnetized Plasma Dynamics: A Neural Operator Application. arXiv:2507.01388</p>

The code is available on GitHub: https://github.com/black-hole-group/fno-vortex.