Learned Surrogates of Three-Dimensional Magnetohydrodynamic Turbulence

Semih Kacmaz; E. A. Huerta

cs.LGOct 31, 2026

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 ·

Abstract

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 $128^3$ 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 $\mathrm{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 $u_{x}$ , the magnetic field $B_{x}$ , the vorticity magnitude $|\boldsymbol{\omega}|$ and the current density $|\mathbf{J}|$ in the $128^{3}$ volume after 25 autoregressive steps from the initial frame ( $\tau=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 $\Phi_{\theta}$ (blue panel; Equation 3 ). $\mathbf{x}_{t}$ is the true frame, $\hat{\mathbf{x}}_{t}$ the prediction and $\tilde{\mathbf{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 $\lambda$ , the first-step flag and the time, form the input $\mathbf{c}_{t}$ of the backbone $f_{\theta}$ , which predicts the change; $\Sigma_{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 $\mathbf{x}_{0}$ , 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 $\mathrm{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 $\Phi_{\theta}$ 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 $\epsilon$ is reported. Symbols and values in Table 1 .

	value	role and choice
$[a_{\min},a_{\max}]$	$[0.01,0.3]$	input-noise levels; spans the errors of rollout inputs, from early to late steps (App. B)
$U_{1}$	40,000	first-stage updates; twice as many improved only $\mathbf{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
$\epsilon$	5%	tolerance of the early-stopping rule, set in advance

Table 1: Training settings, shared by all backbones and Reynolds numbers.

$\mathrm{Re}=1000$ ( $\mathrm{Re}_{\lambda}=9$ – $17$ )
			relative error at $\tau=1$
	Params (M)	GPU-h	$\mathbf{u}$	$\mathbf{B}$	$\boldsymbol{\omega}$	$\mathbf{J}$
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
$\mathrm{Re}=3000$ ( $\mathrm{Re}_{\lambda}=25$ – $50$ )

Table 2: Rollout errors on the

128^{3}

MHD data sets. Relative

L_{2}

error ( Equation 10 ) of velocity, magnetic field, vorticity and current density at the end of the window,

\tau=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.

\mathrm{Re}_{\lambda}

: Taylor-scale Reynolds number over the stored window (App. A ).

Figure 3: Energy spectra at

\tau=1

\mathrm{Re}=6000

(test split, mean over 50 simulations), against the wavenumber

k=2\pi n

up to the Nyquist shell: (a) kinetic, (b) magnetic, truth in black.

(a) Isotropic turbulence
		$L_{2}$ error (%)
	Params (M)	$\mathbf{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

L_{2}

error (%) of velocity

\mathbf{u}

and pressure

p

over 10 predicted frames from 10 given frames (ours use the last four). Small / tiny = at FactFormer’s / Transolver-

\sigma

’s size;

\dagger

= count not reported, ours from the paper. Published rows other than Transolver-

\sigma

’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-

\sigma

). (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

	$\mathrm{Re}=1000$	$\mathrm{Re}=3000$	$\mathrm{Re}=6000$
$\mathrm{Re}_{\lambda}$ , $\tau=0\to 1$	17 $\to$ 9	50 $\to$ 25	101 $\to$ 44
$E_{\mathrm{kin}}(1)/E_{\mathrm{kin}}(0)$ (%)	31	54	66
turnover times in the window	0.9	1.2	1.3
$\min_{\tau}\,k_{\max}\eta$	4.94	2.59	1.63

Appendix

Table 4: Statistics of the stored windows (

\tau\in[0,1]

; frame means over the first 16 training simulations of each set).

\mathrm{Re}_{\lambda}

at the first and last frame; kinetic energy at

\tau=1

relative to

\tau=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	$\mathbf{u}$ / $\mathbf{B}$ / $\boldsymbol{\omega}$ / $\mathbf{J}$
$\mathrm{Re}=1000$ ( $\mathrm{Re}_{\lambda}=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
$\mathrm{Re}=3000$ ( $\mathrm{Re}_{\lambda}=25$ – $50$ )

Appendix

Table 5: Training of the

128^{3}

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

\tau=1

; the rule stops below 5%, or, marked

\dagger

, 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	$\mathbf{u}$	$\mathbf{B}$	$\boldsymbol{\omega}$	$\mathbf{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,

\tau=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, $\mathrm{Re}_{\lambda}=84$	driven compressible ideal MHD, ten regimes (sonic Mach 0.5–7, Alfvénic Mach 0.7 and 2), $\mathrm{Re}_{\lambda}\approx 55$ –80 (our estimate)
Grid, fields	$60^{3}$ ; $\mathbf{u}$ , $p$	$64^{3}$ ; $\rho$ , $\mathbf{u}$ , $\mathbf{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 $L_{2}$ 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 .

Ours (test, mean $\pm$ standard deviation over 3 seeds)
		with projection		without projection
	params (M)	$\mathbf{u}$	$p$	$\mathbf{u}$	$p$
UNet	9.88	6.54 $\pm$ 0.02	9.88 $\pm$ 0.03	6.71 $\pm$ 0.02	9.98 $\pm$ 0.03
UNet, small	4.93	7.67 $\pm$ 0.03	11.49 $\pm$ 0.05	7.85 $\pm$ 0.03	11.61 $\pm$ 0.05
UNet, tiny	1.68	9.43 $\pm$ 0.01	14.35 $\pm$ 0.06	9.59 $\pm$ 0.01	14.43 $\pm$ 0.06
Transformer, $64^{3}$	66.27	8.54 $\pm$ 0.13	14.02 $\pm$ 0.34	8.94 $\pm$ 0.27	14.20 $\pm$ 0.57

Appendix

Table 8: Isotropic turbulence, all rows: relative

L_{2}

error (%) on the test split, velocity / pressure, with and without the divergence-free projection of the predicted velocity (ours: mean

\pm

standard deviation over three seeds; Transolver-

\sigma

: as reported). Small / tiny as in Table 3 ;

64^{3}

= the native transformer layout on the resampled grid,

60^{3}

= the layout fitted to the

60^{3}

grid. Superscripts: p = preprint; r = retrained and reported by another paper;

\dagger

= our count.

Reference	4.93	7.67 $\pm$ 0.03	11.49 $\pm$ 0.05	–	2.5
	params (M)	$\mathbf{u}$	$p$	$\Delta\mathbf{u}$ (%)	GPU-h
No rollout stage	4.93	8.69 $\pm$ 0.04	13.43 $\pm$ 0.10	$+13.3$	1.7
Rollout stage only, equal GPU time	4.93	8.97 $\pm$ 0.02	13.53 $\pm$ 0.05	$+16.9$	2.4
No input noise	4.93	7.53 $\pm$ 0.06	11.32 $\pm$ 0.09	$-1.9$	2.4
Two input frames	4.85	7.85 $\pm$ 0.06	11.78 $\pm$ 0.15	$+2.3$	2.3
No self-attention	4.81	7.63 $\pm$ 0.01	11.43 $\pm$ 0.03	$-0.6$	2.3

Appendix

Table 9: Isotropic turbulence, ablation of the UNet at FactFormer’s size ( Section 5.1 ): relative

L_{2}

error (%) on the test split with the divergence-free projection, mean

\pm

standard deviation over three seeds;

\Delta\mathbf{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.

Ours (test, mean $\pm$ standard deviation over 3 seeds)
			windows (a)		windows (b)		Walrus (median)
	params (M)	one-step	6–12	13–30	6–12	13–30	1–20	21–60
UNet	37.93	0.140 $\pm$ 0.018	0.718 $\pm$ 0.007	1.041 $\pm$ 0.010	0.749 $\pm$ 0.006	1.050 $\pm$ 0.010	0.701 $\pm$ 0.011	1.134 $\pm$ 0.017
UNet, small	21.48	0.154 $\pm$ 0.003	0.720 $\pm$ 0.003	1.026 $\pm$ 0.003	0.750 $\pm$ 0.003	1.034 $\pm$ 0.004	0.707 $\pm$ 0.004	1.101 $\pm$ 0.009
Transformer	66.30	0.161 $\pm$ 0.021	0.689 $\pm$ 0.008	1.031 $\pm$ 0.006	0.720 $\pm$ 0.007	1.040 $\pm$ 0.006	0.683 $\pm$ 0.008	1.156 $\pm$ 0.011
Transformer, small	21.09	0.177 $\pm$ 0.005	0.694 $\pm$ 0.011	1.030 $\pm$ 0.003	0.723 $\pm$ 0.011	1.039 $\pm$ 0.004	0.690 $\pm$ 0.003	1.151 $\pm$ 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

\pm

standard deviation over three seeds). Small as in Table 3 ; w = selected on validation under Walrus’ protocol; r = retrained and reported by another paper.

$\mathrm{Re}=1000$ ( $\mathrm{Re}_{\lambda}=9$ – $17$ )
	$\tau$	$\mathbf{u}$	$\mathbf{B}$	$\boldsymbol{\omega}$	$\mathbf{J}$	$\Omega$ ratio	$\langle\|\mathbf{J}\|^{2}\rangle$ ratio
UNet	0.04	0.004 $\pm$ 0.001	0.007 $\pm$ 0.001	0.010 $\pm$ 0.001	0.017 $\pm$ 0.001	1.000	1.002
	0.48	0.027 $\pm$ 0.012	0.034 $\pm$ 0.011	0.041 $\pm$ 0.017	0.055 $\pm$ 0.018	0.997	1.000
	1	0.048 $\pm$ 0.022	0.068 $\pm$ 0.021	0.069 $\pm$ 0.030	0.094 $\pm$ 0.034	0.997	1.001
Transformer	0.04	0.004 $\pm$ 0.000	0.007 $\pm$ 0.000	0.039 $\pm$ 0.001	0.056 $\pm$ 0.001	1.001	1.003
	0.48	0.011 $\pm$ 0.002	0.020 $\pm$ 0.004	0.026 $\pm$ 0.003	0.040 $\pm$ 0.005	1.000	0.999

Appendix

Table 11: Rollout errors at

\tau=0.04

0.48

and

1

(steps 1, 12 and 25) on the

128^{3}

data sets (test split; mean

\pm

standard deviation over 50 simulations) and the ratios of predicted to true enstrophy

\Omega=\langle|\boldsymbol{\omega}|^{2}\rangle

and mean squared current density.

$\mathrm{Re}=1000$ ( $\mathrm{Re}_{\lambda}=9$ – $17$ )
	noise-free				input noise 0.01
	$\mathbf{u}$	$\mathbf{B}$	$\boldsymbol{\omega}$	$\mathbf{J}$	$\mathbf{u}$	$\mathbf{B}$	$\boldsymbol{\omega}$	$\mathbf{J}$	$\mathbf{J}$ ratio
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
$\mathrm{Re}=3000$ ( $\mathrm{Re}_{\lambda}=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 (

\tau=1

Figure 5: Magnetic field at $\tau=1$ , $\mathrm{Re}=6000$ , same simulation as Figure 1 : isosurfaces of $|\mathbf{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 $\tau=1$ , $\mathrm{Re}=1000$ , as Figure 1 .

Figure 7: Mid-plane slices at $\tau=1$ , $\mathrm{Re}=3000$ , as Figure 1 .

Neural Surrogate Modeling

Neural surrogate modeling uses learned models to approximate costly simulations, enabling faster prediction, design, control, and inference.

See all 440 papers

PHASE: Multi-Regime Modeling of Incompressible Magnetohydrodynamics
Sep 29, 2026Radhika Achikanath Chirakkara, Rajdeep Haldar, Zezheng Song +1Magnetohydrodynamics PDE Operator Learning
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
TurtLES: A Large-Scale Benchmark for Turbulent 3D Neural PDE Surrogates
Nov 6, 2026Jarno Platenburg, Armand Kassaï Koupaï, Paola Cinnella +1PDE Surrogate Modeling Computational Fluid Dynamics
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
A Pre-trained Variational Autoencoder for Gyrokinetic Plasma Turbulence Surrogate Modeling
Sep 29, 2026Minglei Yang, Marshall Nicholson, Diego Del-Castillo-Negrete +2Variational Autoencoders Physics-Informed Generative Modeling
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

Sep 29, 2026J/K move · Enter open · S save

PHASE: Multi-Regime Modeling of Incompressible Magnetohydrodynamics

Radhika Achikanath Chirakkara, Rajdeep Haldar, Zezheng Song, Jiequn Han

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

Magnetohydrodynamics PDE Operator Learning Physics-Informed Neural Operators Computational Fluid Dynamics

Open paper AskPDF arXiv

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

L_2

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.

Abstract

Neural Surrogate Modeling

Explore similar work

PHASE: Multi-Regime Modeling of Incompressible Magnetohydrodynamics

TurtLES: A Large-Scale Benchmark for Turbulent 3D Neural PDE Surrogates

A Pre-trained Variational Autoencoder for Gyrokinetic Plasma Turbulence Surrogate Modeling