Machine-learned spectral PDE solvers for fluids and plasmas

Mellow Baixuan Chen; Alan A. Kaptanoglu

physics.comp-phOct 19, 2026

Machine-learned spectral PDE solvers for fluids and plasmas

Authors: Mellow Baixuan Chen, Alan A. Kaptanoglu

Organizations: Courant Institute of Mathematical Sciences, New York University, New York, 10012, NY, USA · William E. Boeing Department of Aeronautics and Astronautics, University of Washington, Seattle, 98195, Washington, USA

Abstract

Sparse regression methods for discovering governing partial differential equations (PDEs) from data are sensitive to observation noise and are not constrained to reproduce the data under forward integration. We address both limitations by performing weak-form discovery in the spatial discretization in which the PDE is solved, learned by differentiating through multistep rollouts anchored at the observed states. The new method, Weak-form Anchored Rollout for PDEs (WARP), produces a machine-learned PDE solver. In this work, we focus on systems amenable to a spectral Galerkin solver. Across eight benchmark fluid and plasma systems, including high-dimensional turbulent simulations, WARP produces numerically stable rollouts from new random initial conditions and substantially improves coefficient recovery, rollout accuracy, and statistical measures compared with a single-step weak-form SINDy baseline.

Figures & tables

Benchmark	Governing equation and true coefficients	Training initial condition and identification window	Spatial state and time sampling	$H$ (steps)
Deterministic benchmarks
Burgers	$u_{t}=buu_{x}+\nu u_{xx}$ , $(b^{\ast},\nu^{\ast})=(-1,0.05)$	$u_{0}=\sin x+0.5\sin(2x+0.3)$ on $[0,2\pi]$ , $t\in[0,0.4]$	64 grid points, $M=12$ , $K=d=25$ , $\Delta t_{\rm s}=10^{-3}$ , $T=401$	$50$ bi
KdV	$u_{t}=-\alpha uu_{x}-\beta u_{xxx}$ , $(\alpha^{\ast},\beta^{\ast})=(6,1)$	$u_{0}=0.20\sin x+0.10\cos(2x-0.3)+0.06\sin(3x+0.7)+0.04\cos(5x-0.2)$ , $t\in[0,0.1]$	128 grid points, $M=15$ , $K=d=31$ , $\Delta t_{\rm s}=10^{-5}$ , $T=10{,}001$	$50$ bi
ZK	$\begin{aligned} u_{t}={}&-\alpha uu_{x}\\ &-\beta\partial_{x}\Delta u,\end{aligned}$ $(\alpha^{\ast},\beta^{\ast})=(6,1)$	$u_{0}=0.20\sin x\cos y+0.10\cos(2x-0.3)\cos(y+0.5)+0.08\sin(x+0.4)\sin(2y-0.2)+0.05\cos(2x)\sin(2y)$ , $t\in[0,0.016]$	$64^{2}$ grid, $M=4$ , $K=d=81$ , $\Delta t_{\rm s}=2\!\times\!10^{-5}$ , $T=801$	$50$ bi
Decaying 2D vorticity	$\begin{aligned} \omega_{t}={}&-\alpha(u\!\cdot\!\nabla)\omega\\ &+\nu\Delta\omega,\\ u={}&(\chi_{y},-\chi_{x}),\hskip 8.50012pt\Delta\chi=-\omega,\end{aligned}$ $(\alpha^{\ast},\nu^{\ast})=(1,0.02)$	$\omega_{0}=\sin x\cos y+0.8\cos(x+0.4)\sin(y-0.2)+0.5\sin(2x-0.3)\cos(y+0.5)+0.4\cos x\sin 2y$ , $t\in[0,0.2]$	$64^{2}$ grid, $M=4$ , $K=d=81$ , $\Delta t_{\rm s}=5\!\times\!10^{-4}$ , $T=401$	$50$ bi
Orszag–Tang MHD vortex	$\begin{aligned} \omega_{t}={}&-a_{u}(u\!\cdot\!\nabla)\omega\\ &+a_{B}(B\!\cdot\!\nabla)j+\nu\Delta\omega,\\ A_{t}={}&-a_{A}(u\!\cdot\!\nabla)A+\eta_{B}\Delta A\end{aligned}$ $\begin{aligned} (a_{u}^{\ast},a_{B}^{\ast},\nu^{\ast})&=(1,1,0.025),\\ (a_{A}^{\ast},\eta_{B}^{\ast})&=(1,0.025)\end{aligned}$	$\chi_{0}=2(\cos x+\cos y)$ and $A_{0}=2\cos x+\cos 2y$ on $[0,2\pi]^{2}$ , $t\in[0.15,1.5]$	$64^{2}$ grid, $M=21$ , $K=1{,}849$ , $d=3{,}698$ , $\Delta t_{\rm s}=0.005$ , $T=271$	$200$ fwd

Table 1: Benchmark definitions and numerical protocols. All domains are periodic. Here

M

is the largest retained Fourier mode per axis,

K

is the basis size per field,

d=N_{f}K

is the total modal-state dimension,

\Delta t_{\rm s}

is the interval between stored snapshots, and

T

is the total number of snapshots. The last column gives the rollout horizon

H

in snapshot intervals. The abbreviations bi and fwd denote bidirectional and forward-only rollouts. Adaptive Dormand–Prince integration is used throughout.

Figure 9: Solver validation for Kolmogorov flow at $64^{2}$ ( $M=21$ ) and $128^{2}$ ( $M=42$ ). The spectrum is taken at the end of the run; the two vorticity fields are different realizations of the chaotic state

Figure 12: Kolmogorov-flow statistics from five fresh initial conditions versus training noise

Figure 13: Kolmogorov flow at

\eta=5\%

: energy and dissipation along the deployment of Fig. 11 , and the enstrophy spectrum averaged over the discovery window. Black: clean trajectory. Grey: true coefficients

\xi^{\ast}

restarted from the same noisy first snapshot as WARP and FD

Figure 14: Kolmogorov-flow vorticity at $\eta=5\%$ (seed 42), with WARP and FD launched from the noisy snapshot at $t=500$ . Color scales follow the reference in each column; panels that exceed the range are annotated with their peak

Figure 15: Saturated HW data-generation diagnostics at $c_{1}=1$ on the $M=12$ band, with DNS-scale (top) and resolution-matched (bottom) hyperdiffusion. Other $c_{1}$ : Appendix A

Figure 18: Saturated HW at

c_{1}=1

and

\eta=5\%

: particle flux and energy along the deployment of Fig. 17 , and the energy spectrum averaged over the fresh initial conditions of Fig. 19 . In the time traces, black is the clean trajectory and grey is the true model restarted from the noisy first snapshot; in the spectrum, grey is the true-model average

Figure 19: Saturated HW at

c_{1}=1

: statistics from five fresh initial conditions versus training noise. Error bars show one standard deviation over the initial conditions

Figure 20: Independent

512^{2}

hw2d DNS at

c_{1}=1

, with the reference flux range of Gahr et al. [57] . Other

c_{1}

: Appendix A

Figure 23: mHW coefficient error $\|\xi-\xi^{\ast}\|_{2}$ versus noise level at four adiabaticities, with FD dashed and WARP solid

Figure 24: mHW: statistics from five fresh initial conditions across the zonal-flow transition, for models trained at $\eta=0.05$ and $0.1$

Panel of Figure 25 — Figure 25: mHW at $\eta=5\%$ (seed 42): rollout error against the clean reference, with all three models launched from the same noisy first snapshot at $\gamma_{\max}t=60$ . Dotted lines mark one and five growth times after launch

Figure 26: mHW at

C=1

and

\eta=5\%

: particle flux and zonal kinetic-energy fraction along the deployment of Fig. 25 , and the kinetic-energy spectrum averaged over the fresh initial conditions of Fig. 24

Figure 27: Orszag–Tang data-generation diagnostics on the $64^{2}$ grid with $M=21$ ; $t^{\ast}$ is the dissipation-peak time of Eq. ( 9 )

Figure 28: Solver validation of

J_{\max}(t)

for the

M=21

generator in blue and the

128^{2}

reference of Rosenberg et al. [84] in black

Panel of Figure 29 — Figure 29: Orszag–Tang robustness studies. (a) Horizon sweep: $E_{\xi}$ versus rollout duration $T_{\rm roll}=H\Delta t_{\rm s}$ . Open markers omit the per-field weights of Sec. 3.7 ; stars mark the production setting ( $H=200$ ), averaged over three noise seeds, and all other points use seed 42. (b) Coefficient error $E_{\xi}$ and testing trajectory error versus noise level

Figure 30: Orszag–Tang event errors of Eq. ( 10 ) from five fresh initial conditions, using the seed-42 fit. Error bars show one standard deviation over the initial conditions

Figure 31: Orszag–Tang deployment from a held-out initial condition on the production band with a $64^{2}$ grid and $M=21$ , for coefficients learned at $\eta=5\%$ . Rows within each block show the true model, the WARP deployment, and the FD deployment; the blocks are vorticity $\omega$ and current density $j$ . All three are integrated from the same new initial condition. Columns span the discovery window and include the current-sheet event of this initial condition, labeled $t^{\ast}$ in the panel titles. Color scales follow the reference at each time; panels that exceed the range are annotated with their peak

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 32: Additional band-closed HW data-generation diagnostics at $c_{1}=0.1$ (top) and $c_{1}=5$ (bottom), using resolution-matched hyperdiffusion. Columns and colors follow Fig. 15 : $\Gamma_{n}$ is particle flux, $\Gamma_{c}$ resistive energy dissipation, $E(k)$ shell energy, $n$ density fluctuation, and $\omega$ vorticity. Both flux traces are spatially integrated

Panel of Figure 33 — Figure 33: Saturated HW turbulence at $c_{1}=1$ and $\eta=5\%$ . Rows show the clean reference and the WARP and FD deployments launched from the noisy snapshot at $t=250$ ; columns are $10$ and $20$ time units ( $20$ and $40$ training horizons) after launch. Color scales are set by the reference at each time and shared down a column; panels that leave the range are annotated with their peak

Figure 34: As Fig. 20 , for $c_{1}=0.1$ (top) and $c_{1}=5$ (bottom)

Figure 35: As Fig. 22 , identification from projected independent DNS at $c_{1}=0.1$ (left) and $5$ (right)

Figure 36: WARP coefficient error versus retained band

M

(ZK) and rollout horizon

H

(Burgers). Error bars show one standard deviation over three noise seeds

Panel of Figure 37 — Figure 37: Long-horizon deployment of the $\eta=10\%$ models from a new initial condition, for twice the original testing interval (dotted line). HW uses $c_{1}=1$ ; $|q|_{\max}$ is the largest absolute modal coefficient

Figure 38: Kolmogorov flow with the overcomplete library: recovered coefficient magnitudes at each noise level (top four rows), and true-term error and spurious-term norm versus noise (bottom)

Figure 39: Deployment of the overcomplete-library models trained at $\eta=10\%$ from a new initial condition. $D$ and $I$ are evaluated with the true $\nu$ and $A_{f}$ for all models

PDE Solving

PDE solving develops numerical and learned methods to approximate solutions of partial differential equations across scientific applications.

See all 250 papers

Weak-PDE-Net: Discovering Open-Form PDEs via Differentiable Symbolic Networks and Weak Formulation
Mar 24, 2026Xinxin Li, Xingyu Cui, Jin Qi +3PDE Discovery Symbolic Regression
School of Mathematical Sciences, East China Normal University, Shang Hai, 200241, China · Institute of Applied Physics and Computational Mathematics, 100094, Beijing, China · Shanghai Zhangjiang Institute of Mathematics, 201203, Shanghai, China +2
Warm-starting PDE solvers with any-dimensional machine learning
Sep 30, 2026Wilson G. Gregory, George A. Kevrekidis, Ben Blum-Smith +1Neural PDE Solvers PDE Surrogate Modeling
Department of Statistics and Irving Institute for Cancer Dynamics, Columbia University · Division of Applied Mathematics, Brown University · Department of Applied Mathematics and Statistics, Johns Hopkins University
FoundPDE: A generative foundation model for data-driven symbolic discovery of partial differential equations
Oct 18, 2026Langchen Liu, Junlin Luo, Min Zhu +2PDE Discovery PDE Foundation Models
Department of Statistics and Data Science, Yale University, New Haven, CT 06511, USA · Department of Mechanical Engineering, Stanford University, Stanford, CA 94305, USA · Department of Chemical and Environmental Engineering, Yale University, New Haven, CT 06511, USA

Mar 24, 2026J/K move · Enter open · S save

Weak-PDE-Net: Discovering Open-Form PDEs via Differentiable Symbolic Networks and Weak Formulation

Xinxin Li, Xingyu Cui, Jin Qi, Juan Zhang+2

School of Mathematical Sciences, East China Normal University, Shang Hai, 200241, China · Institute of Applied Physics and Computational Mathematics, 100094, Beijing, China · Shanghai Zhangjiang Institute of Mathematics, 201203, Shanghai, China+2

PDE Discovery Symbolic Regression

Open paper AskPDF arXiv

Discovering governing Partial Differential Equations (PDEs) from sparse and noisy data is a challenging issue in data-driven scientific computing. Conventional sparse regression methods often suffer from two major limitations: (i) the instability of numerical differentiation under sparse and noisy data, and (ii) the restricted flexibility of a pre-defined candidate library. We propose Weak-PDE-Net, an end-to-end differentiable framework that can robustly identify open-form PDEs. Weak-PDE-Net consists of two interconnected modules: a forward response learner and a weak-form PDE generator. The learner embeds learnable Gaussian kernels within a lightweight MLP, serving as a surrogate model that adaptively captures system dynamics from sparse observations. Meanwhile, the generator integrates a symbolic network with an integral module to construct weak-form PDEs, avoiding explicit numerical differentiation and improving robustness to noise. To relax the constraints of the pre-defined library, we leverage Differentiable Neural Architecture Search strategy during training to explore the functional space, which enables the efficient discovery of open-form PDEs. The capability of Weak-PDE-Net in multivariable systems discovery is further enhanced by incorporating Galilean Invariance constraints and symmetry equivariance hypotheses to ensure physical consistency. Experiments on several challenging PDE benchmarks demonstrate that Weak-PDE-Net accurately recovers governing equations, even under highly sparse and noisy observations.

Abstract

PDE Solving

Explore similar work

Weak-PDE-Net: Discovering Open-Form PDEs via Differentiable Symbolic Networks and Weak Formulation

Warm-starting PDE solvers with any-dimensional machine learning

FoundPDE: A generative foundation model for data-driven symbolic discovery of partial differential equations