Learned Surrogates of Three-Dimensional Magnetohydrodynamic Turbulence
Authors: Semih Kacmaz, E. A. Huerta
Organizations: National Center for Supercomputing Applications, University of Illinois Urbana-Champaign, Urbana, Illinois, USA · Department of Physics, University of Illinois Urbana-Champaign, Urbana, Illinois, USA · Data Science and Learning Division, Argonne National Laboratory, Lemont, Illinois, USA · Department of Computer Science, The University of Chicago, Chicago, Illinois, USA · Department of Astronomy, University of Illinois Urbana-Champaign, Urbana, Illinois, USA
Three-dimensional magnetohydrodynamic (MHD) turbulence, the flow of astrophysical and fusion plasmas, is costly to simulate. We present, to our knowledge, the first learned surrogates of three-dimensional resistive MHD turbulence, and compare three standard backbones in that role: a convolutional UNet, a shifted-window transformer and a Fourier neural operator. They are trained as autoregressive emulators of decaying MHD turbulence at a resolution of 1283 and Reynolds numbers of 1000, 3000 and 6000, on 500 simulations per Reynolds number, each with one two-stage procedure, single steps and then short rollouts, and stopped by one rule. Over a window of about one large-eddy turnover time, the UNet and the transformer reach relative velocity errors of 0.018 to 0.053 and keep the current sheets, the enstrophy, the current density and the energy spectra of the simulation; which of the two is the more accurate depends on the Reynolds number, the UNet leading at 3000 and 6000 and the transformer at 1000. Applied with the same backbones and training procedure to two public benchmarks, the models reach the lowest error reported to date on the isotropic turbulence set of FactFormer, 39% below the best published model. On the compressible MHD set of The Well, their rollout errors are a third to a half below those of the best of The Well's baselines. They also match a 1.3 B-parameter pretrained foundation model in its rollout windows, with 20-62 times fewer parameters. The code, the configuration of every model, the evaluation scripts and the data generator with the seed of every simulation will be released upon publication.
Figures & tables
Figure 1: Three-dimensional turbulent MHD at Re=6000 after a rollout of about one eddy-turnover time, predicted by three backbones trained with one procedure. Mid-plane slices ( z=L/2 ) of the velocity ux , the magnetic field Bx , the vorticity magnitude ∣ω∣ and the current density ∣J∣ in the 1283 volume after 25 autoregressive steps from the initial frame ( τ=1 ), for the truth and for the UNet, the transformer and the FNO. Each row shares the truth’s colour scale. Errors in Table 2 , other Reynolds numbers in App. D .
Figure 2: Method. (a) The one-step map Φθ (blue panel; Equation 3 ). xt is the true frame, x^t the prediction and x~t the frame as the network receives it, noisy in training and the previous prediction in a rollout. The two received frames and three scalars, the noise level λ , the first-step flag and the time, form the input ct of the backbone fθ , which predicts the change; Σt scales it per component and step, and the most recent frame is added back. A rollout feeds every prediction back as the next input, T times from x0 , and the previous input moves to the second slot (grey arrow). At evaluation the prediction is first clipped and projected onto divergence-free fields ( Equation 5 ); training rollouts skip this. The volumes are isosurfaces of the vorticity magnitude of a training simulation at Re=3000 and of the signed change of one velocity component (red positive, blue negative). (b) The two training stages as computation graphs (dashed boxes: true frames; dots: predictions; red links: losses): single steps from noisy true frames, then rollouts of K steps with a loss at every step and the gradient through the whole window. Each Φθ receives the two most recent frames; the chain draws the latest. (c) Early stopping of the second stage on the validation rollout error, checked every U updates; the model before the first check that gains less than ϵ is reported. Symbols and values in Table 1 .
value
role and choice
[amin,amax]
[0.01,0.3]
input-noise levels; spans the errors of rollout inputs, from early to late steps (App. B)
U1
40,000
first-stage updates; twice as many improved only J , at 1.4 times the cost (App. B)
K
4
window length, as the unrolled baseline of Dai et al. (2026b) ; memory grows with K
U
6,600
updates between early-stopping checks
ϵ
5%
tolerance of the early-stopping rule, set in advance
Table 1: Training settings, shared by all backbones and Reynolds numbers.
relative error at τ=1
Params (M)
GPU-h
u
B
ω
J
Re=1000 ( Reλ=9 – 17 )
UNet
37.7
20
0.048
0.068
0.069
0.094
Transformer
66.3
31
0.018
0.031
0.034
0.051
FNO
37.8
18
0.127
0.285
0.272
0.491
Re=3000 ( Reλ=25 – 50 )
Table 2: Rollout errors on the 1283 MHD data sets. Relative L2 error ( Equation 10 ) of velocity, magnetic field, vorticity and current density at the end of the window, τ=1 , about one eddy-turnover time, on the test split (50 simulations). Lower is better; best per column and Reynolds number in bold. Params: trainable parameters. GPU-h: training time of the final model over both stages on GH200 GPUs. Reλ : Taylor-scale Reynolds number over the stored window (App. A ).
Figure 3: Energy spectra at τ=1 , Re=6000 (test split, mean over 50 simulations), against the wavenumber k=2πn up to the Nyquist shell: (a) kinetic, (b) magnetic, truth in black.
(a) Isotropic turbulence
L2 error (%)
Params (M)
u
p
Ours
UNet
9.88
6.54
9.88
UNet, small
4.93
7.67
11.49
UNet, tiny
1.68
9.43
14.35
Table 3: Public benchmarks (test splits; ours: mean of three seeds; lower is better, best in bold; standard deviations, all rows and the protocols in App. C ). (a) Isotropic turbulence ( Li et al., 2023 ) : relative L2 error (%) of velocity u and pressure p over 10 predicted frames from 10 given frames (ours use the last four). Small / tiny = at FactFormer’s / Transolver- σ ’s size; † = count not reported, ours from the paper. Published rows other than Transolver- σ ’s own are retrainings reported by Shangguan et al. (2026) (EddyFormer ( Du & Krishnapriyan, 2025 ) , FactFormer, DRIFT-Net ( Li & Salim, 2026 ) ) and by Li et al. (2023) (Dil-ResNet ( Stachenfeld et al., 2022 ) ). (b) MHD_64 ( Ohana et al., 2024 ) : VRMSE, mean over the seven fields. One-step: all windows of the test trajectories. The Well’s protocol: rollouts from the first four frames, mean over predicted steps 6–12 and 13–30 and over trajectories; CNextU-net ( Liu et al., 2022 ) also retrained by us with The Well’s code and budget. Walrus’ protocol ( McCabe et al., 2026 ) : rollouts from frames 7–10, median over trajectories of the mean over predicted steps 1–20 and 21–60. EddyFormer follows The Well’s protocol. Small = at CNextU-net’s size.
Figure 4: (a) Isotropic turbulence: velocity error against parameter count for our UNet (three widths) and transformer (two widths; the 4 M model is tabulated in App. C ) and for the published models (grey; dotted lines: no reported count; bar: the 2.0–2.2 M range of our count for Transolver- σ ). (b) MHD_64: VRMSE per rollout step, averaged over fields, the ten test trajectories and three seeds, for our UNet (38 M) and transformer (66 M) and for the retrained CNextU-net (grey). Shaded: The Well’s windows; black: the published means of its best baseline (FNO) over them; red, dashed: Walrus’ published medians over its windows, drawn at their frames (its rollouts start at frame 11, our step 8).
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
Re=1000
Re=3000
Re=6000
Reλ , τ=0→1
17 → 9
50 → 25
101 → 44
Ekin(1)/Ekin(0) (%)
31
54
66
turnover times in the window
0.9
1.2
1.3
minτkmaxη
4.94
2.59
1.63
Appendix
Table 4: Statistics of the stored windows ( τ∈[0,1] ; frame means over the first 16 training simulations of each set). Reλ at the first and last frame; kinetic energy at τ=1 relative to τ=0 ; large-eddy turnover times in the window; smallest value over the window of the largest resolved wavenumber times the Kolmogorov scale.
updates
blocks
stop gain
GPU-h
gap (%)
Params (M)
one-step
rollout
trained
reported
(%)
reported
all
u / B / ω / J
Re=1000 ( Reλ=9 – 17 )
UNet
37.7
40,000
0
1
0
12.4 †
20.0
49.9
190 / 76 / 170 / 107
Transformer
66.3
40,000
26,400
4
3
2.7
30.9
36.0
50 / 46 / 25 / 30
FNO
37.8
40,000
13,200
2
1
3.6
17.5
20.7
61 / 70 / 31 / 38
Re=3000 ( Reλ=25 – 50 )
Appendix
Table 5: Training of the 1283 models: parameters, updates of the one-step stage and of the rollout stage up to the reported model, rollout blocks trained (including the stop block), index of the reported block, the validation gain of the stop block over it (the stopping rule’s gain, Equation 9 , at τ=1 ; the rule stops below 5%, or, marked † , when the training-simulation error falls while the validation error lags), GPU-hours of the reported model and of all blocks, and the gap between the validation and the training-simulation error of the reported model (validation / training −1 ).
GPU-h
updates
u
B
ω
J
two-stage
41
40,000 + 8,800
0.030
0.077
0.064
0.173
rollout training from scratch
41
17,600
0.039
0.116
0.096
0.378
ratio from scratch / two-stage
1.32
1.50
1.50
2.18
from scratch worse in (of 12 simulations)
12
12
12
12
Appendix
Table 6: Development test of the two-stage procedure at equal GPU time (Re =3000 , UNet, 12 validation simulations, τ=1 ): two-stage training against rollout training from random initialization.
Isotropic turbulence ( Li et al., 2023 )
MHD_64 ( Ohana et al., 2024 )
Flow
forced incompressible turbulence, Reλ=84
driven compressible ideal MHD, ten regimes (sonic Mach 0.5–7, Alfvénic Mach 0.7 and 2), Reλ≈55 –80 (our estimate)
Grid, fields
603 ; u , p
643 ; ρ , u , B
Trajectories
1,100 of 20 frames
100 of 100 frames (97 in the released files)
Split
1,000 / 100, fixed by FactFormer’s code; no validation split, ours = the last 50 training trajectories
77 / 10 / 10, official
Input frames
10 given; ours use the last 4
4 (Walrus: 10 in its released MHD_64 configuration, 3 in its paper)
Metric
relative L2 error (%) per field over space and the 10 predicted frames
VRMSE per field, mean over the seven fields
Appendix
Table 7: Benchmark protocols. Ours = the choices of our rows; the rest follows the benchmark. Cost = GPU-hours of the final model’s stages on one GH200, in-training validation included, over all rows of Tables 8 and 10 .
with projection
without projection
params (M)
u
p
u
p
Ours (test, mean ± standard deviation over 3 seeds)
UNet
9.88
6.54 ± 0.02
9.88 ± 0.03
6.71 ± 0.02
9.98 ± 0.03
UNet, small
4.93
7.67 ± 0.03
11.49 ± 0.05
7.85 ± 0.03
11.61 ± 0.05
UNet, tiny
1.68
9.43 ± 0.01
14.35 ± 0.06
9.59 ± 0.01
14.43 ± 0.06
Transformer, 643
66.27
8.54 ± 0.13
14.02 ± 0.34
8.94 ± 0.27
14.20 ± 0.57
Appendix
Table 8: Isotropic turbulence, all rows: relative L2 error (%) on the test split, velocity / pressure, with and without the divergence-free projection of the predicted velocity (ours: mean ± standard deviation over three seeds; Transolver- σ : as reported). Small / tiny as in Table 3 ; 643 = the native transformer layout on the resampled grid, 603 = the layout fitted to the 603 grid. Superscripts: p = preprint; r = retrained and reported by another paper; † = our count.
params (M)
u
p
Δu (%)
GPU-h
Reference
4.93
7.67 ± 0.03
11.49 ± 0.05
–
2.5
No rollout stage
4.93
8.69 ± 0.04
13.43 ± 0.10
+13.3
1.7
Rollout stage only, equal GPU time
4.93
8.97 ± 0.02
13.53 ± 0.05
+16.9
2.4
No input noise
4.93
7.53 ± 0.06
11.32 ± 0.09
−1.9
2.4
Two input frames
4.85
7.85 ± 0.06
11.78 ± 0.15
+2.3
2.3
No self-attention
4.81
7.63 ± 0.01
11.43 ± 0.03
−0.6
2.3
Appendix
Table 9: Isotropic turbulence, ablation of the UNet at FactFormer’s size ( Section 5.1 ): relative L2 error (%) on the test split with the divergence-free projection, mean ± standard deviation over three seeds; Δu = change of the velocity error relative to the reference, the row UNet, small of Table 3 . Each variant changes one element of the reference and keeps its seeds, data and evaluation; the learning rate is not re-tuned. No rollout stage: the reference’s own model at the end of its one-step stage. Rollout stage only: 13,200 updates on four-step windows from random initialization, the reference’s training time. No self-attention: the attention block at the coarsest level removed. GPU-h: training of the final model on one GH200. Thresholds fixed before the runs: at least 10 % higher = the element matters; within 5 % = it does not.
windows (a)
windows (b)
Walrus (median)
params (M)
one-step
6–12
13–30
6–12
13–30
1–20
21–60
Ours (test, mean ± standard deviation over 3 seeds)
UNet
37.93
0.140 ± 0.018
0.718 ± 0.007
1.041 ± 0.010
0.749 ± 0.006
1.050 ± 0.010
0.701 ± 0.011
1.134 ± 0.017
UNet, small
21.48
0.154 ± 0.003
0.720 ± 0.003
1.026 ± 0.003
0.750 ± 0.003
1.034 ± 0.004
0.707 ± 0.004
1.101 ± 0.009
Transformer
66.30
0.161 ± 0.021
0.689 ± 0.008
1.031 ± 0.006
0.720 ± 0.007
1.040 ± 0.006
0.683 ± 0.008
1.156 ± 0.011
Transformer, small
21.09
0.177 ± 0.005
0.694 ± 0.011
1.030 ± 0.003
0.723 ± 0.011
1.039 ± 0.004
0.690 ± 0.003
1.151 ± 0.027
Appendix
Table 10: MHD_64, all rows: test VRMSE under The Well’s protocol with its windows read as predicted steps 6–12 and 13–30 (a) or 7–12 and 14–30 (b), and under Walrus’ protocol (ours: mean ± standard deviation over three seeds). Small as in Table 3 ; w = selected on validation under Walrus’ protocol; r = retrained and reported by another paper.
τ
u
B
ω
J
Ω ratio
⟨∣J∣2⟩ ratio
Re=1000 ( Reλ=9 – 17 )
UNet
0.04
0.004 ± 0.001
0.007 ± 0.001
0.010 ± 0.001
0.017 ± 0.001
1.000
1.002
0.48
0.027 ± 0.012
0.034 ± 0.011
0.041 ± 0.017
0.055 ± 0.018
0.997
1.000
1
0.048 ± 0.022
0.068 ± 0.021
0.069 ± 0.030
0.094 ± 0.034
0.997
1.001
Transformer
0.04
0.004 ± 0.000
0.007 ± 0.000
0.039 ± 0.001
0.056 ± 0.001
1.001
1.003
0.48
0.011 ± 0.002
0.020 ± 0.004
0.026 ± 0.003
0.040 ± 0.005
1.000
0.999
Appendix
Table 11: Rollout errors at τ=0.04 , 0.48 and 1 (steps 1, 12 and 25) on the 1283 data sets (test split; mean ± standard deviation over 50 simulations) and the ratios of predicted to true enstrophy Ω=⟨∣ω∣2⟩ and mean squared current density.
noise-free
input noise 0.01
u
B
ω
J
u
B
ω
J
J ratio
Re=1000 ( Reλ=9 – 17 )
UNet
0.063
0.081
0.088
0.117
0.063
0.081
0.088
0.117
1.00
Transformer
0.021
0.036
0.037
0.057
0.021
0.037
0.075
0.079
1.38
FNO
0.125
0.281
0.270
0.488
0.127
0.288
0.474
0.878
1.80
Re=3000 ( Reλ=25 – 50 )
Appendix
Table 12: Sensitivity to input noise: rollouts of the final models with Gaussian noise of level 0.01 added to the received frames at every step, against the noise-free rollouts, on the first four validation simulations ( τ=1 ).
Figure 5: Magnetic field at τ=1 , Re=6000 , same simulation as Figure 1 : isosurfaces of ∣B∣ at 0.45 (translucent) and 0.8 (opaque) of the truth’s 99.5th percentile, the same absolute levels in every panel; the truth and the three models’ rollouts after 25 steps.
Figure 6: Mid-plane slices at τ=1 , Re=1000 , as Figure 1 .
Figure 7: Mid-plane slices at τ=1 , Re=3000 , as Figure 1 .
Magnetohydrodynamics (MHD) is central to plasma modeling in astrophysics, space science, fusion, and engineering, but resolving multiscale MHD dynamics is computationally expensive. Machine-learning surrogates enable fast inference by learning reusable solution operators, yet existing models require separate training for each physical regime, limiting generalization across varying parameter settings. We introduce PHASE, a PHysics-Adaptive Scalable operator with residual Error correction, designed to model incompressible MHD across varying physical parameters with a single model. PHASE combines transfer learning, regime-aware adaptation, physics-centered learning, and residual refinement to improve both physical fidelity and generalization across MHD regimes. Together, these improvements achieve state-of-the-art prediction accuracy on two-dimensional MHD turbulence by reducing relative L2 errors on physical fields by more than an order of magnitude compared to prior MHD neural-operator baselines. Moreover, PHASE generalizes successfully to unseen parameter values without retraining, demonstrating the cross-regime adaptability expected from operator learning. We evaluate PHASE beyond point-wise prediction errors using derived physical fields, spectral analysis, and distribution statistics, consistently observing improved physical fidelity. We further show that our framework can accurately simulate MHD instabilities by testing it on the Kelvin--Helmholtz instability, demonstrating the robustness of our method.
Radhika Achikanath Chirakkara, Rajdeep Haldar, Zezheng Song +1
Canadian Institute for Theoretical Astrophysics, University of Toronto, ON, Canada, M5S 3H8 · Department of Statistics, Purdue University, IN, United States of America, 47907 · University of Maryland, College Park, MD, United States of America, 20742 +1
Machine-learning surrogates are increasingly used to accelerate computational fluid dynamics, yet progress is limited by the lack of benchmarks capturing realistic, time-dependent turbulent flows. We introduce TurtLES, a 14 TB dataset of 3D turbulent cylinder wakes in the shear-layer transition regime, computed via high-fidelity implicit large-eddy simulations. The dataset contains 400 trajectories of 400 time steps each, with 3-11 million points per frame, spanning variations in geometry, Reynolds number, and angle of attack. Unlike existing datasets, TurtLES combines (i) physically realistic turbulence with an active energy cascade driven by vortex stretching, (ii) dense irregular meshes, and (iii) long-horizon temporal dynamics. We benchmark state-of-the-art neural operators on autoregressive rollouts and identify key failure modes, including temporal instability, loss of small-scale energy, and challenges in learning on large irregular domains. TurtLES provides a new benchmark for evaluating surrogate models on 3D turbulent wakes at industrially relevant mesh resolution.
Sorbonne Université, CNRS, ISIR, Paris, France · Delft University of Technology, Faculty of Aerospace Engineering, Delft, The Netherlands · Sorbonne Université, CNRS, Institut Jean Le Rond d’Alembert, Paris, France +1
Machine learning surrogate models offer a promising path toward accelerating plasma turbulence simulations. We present PreVAE-Turb, a surrogate modeling framework that leverages pre-trained variational autoencoders (VAEs) from the Stable Diffusion image generation model for efficient spatial compression of turbulence fields. The pre-trained VAE is fine-tuned on turbulence data using a physics-informed loss function that includes a spectral loss operating in Fourier space to enforce spectral accuracy across scales. The VAE is combined with convolutional long short-term memory (ConvLSTM) networks to learn temporal dynamics in latent space, with a manifold consistency error metric that monitors encode--decode consistency during autoregressive rollouts. We validate the framework on two-dimensional Hasegawa-Wakatani drift-wave turbulence and extend it to gyrokinetic turbulence from the GENE code, where a four-channel adaptation simultaneously predicts electrostatic potential, density, and parallel/perpendicular temperature fluctuations without requiring architecture redesign. Once trained, inference generates thousands of time steps in seconds on a single GPU, providing substantial computational acceleration compared to direct numerical simulation. The pre-trained approach offers a transferable methodology broadly applicable to various turbulence simulation codes.
Minglei Yang, Marshall Nicholson, Diego Del-Castillo-Negrete +2
Oak Ridge National Laboratory, Oak Ridge, TN, 37831, USA · Grand Valley State University, Allendale, MI, 49401, USA · University of Texas at Austin, Austin, TX, 78712, USA