category: literaturenote citekey: zhangomnifluidsunifiedphysicspre2025 title: "OmniFluids: Unified Physics Pre-trained Modeling of Fluid Dynamics" authors: "Zhang, Rui; Meng, Qi; Wan, Han; Liu, Yang; Ma, Zhi-Ming; Sun, Hao" year: 2025 date: 2025-00-00 2025 url: "https://arxiv.org/abs/2506.10862v1" zotero_key: 227UVNI3 zotero_storage: 7ZZPHQ77 collections: transpordi_planeerimine folder: 001_artiklid firstAuthor: "Zhang, Rui"

status: converted

OmniFluids: Unified Physics Pre-trained Modeling of Fluid Dynamics

Rui Zhang1,: , Qi Meng2,: , Han Wan¹ , Yang Liu³ , Zhi-Ming Ma² , Hao Sun1,˚

¹Gaoling School of Artificial Intelligence, Renmin University of China, Beijing, China ²Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing, China ³School of Engineering Science, University of Chinese Academy of Sciences, Beijing, China

:Equal contribution *Corresponding author

Abstract

High-fidelity and efficient simulation of fluid dynamics drive progress in various scientific and engineering applications. Traditional computational fluid dynamics methods offer strong interpretability and guaranteed convergence, but rely on fine spatial and temporal meshes, incurring prohibitive computational costs. Physics-informed neural networks (PINNs) and neural operators aim to accelerate PDE solvers using deep learning techniques. However, PINNs require extensive retraining and careful tuning, and purely data-driven operators demand large labeled datasets. Hybrid physics-aware methods embed numerical discretizations into network architectures or loss functions, but achieve marginal speed gains and become unstable when balancing coarse priors against high-fidelity measurements. To this end, we introduce Omni-Fluids, a unified physics pre-trained operator learning framework that integrates physics-only pre-training, coarse-grid operator distillation, and few-shot fine-tuning, which enables fast inference and accurate prediction under limited or zero data supervision. For architectural design, the key components of OmniFluids include a mixture of operators, a multi-frame decoder, and factorized Fourier layers, which enable efficient and scalable modeling of diverse physical tasks while maintaining seamless integration with physics-based supervision. Across a broad range of two- and three-dimensional benchmarks, OmniFluids significantly outperforms state-of-theart AI-driven methods in flow field reconstruction and turbulence statistics accuracy, delivering 10-100x speedups compared to classical solvers, and accurately recovers unknown physical parameters from sparse, noisy data. This work establishes a new paradigm for efficient and generalizable surrogate modeling in complex fluid systems under limited data availability.

Keywords: Neural operator learning, AI for physics, Computational fluid dynamics, Physicsinformed machine learning.

Introduction

Fluid dynamics plays a fundamental role in various scientific and engineering fields. Understanding and predicting fluid behavior is essential for designing aircraft [1], forecasting weather [2], and simulating biological flows [3]. To solve these complex problems, computational fluid dynamics (CFD) has been developed over the past century as a primary approach for modeling and simulating fluid dynamics. Traditional numerical methods, such as the finite difference method (FDM), the finite element method (FEM), and the finite volume method (FVM), have been central to the advancement of CFD [4]. These methods are based on rigorous mathematical theory and offer clear interpretability, making them reliable tools for solving various CFD problems. However, traditional numerical methods often require fine-grained spatial and temporal discretization to achieve high accuracy and stability. This requirement leads to substantial computational costs, particularly for strongly nonlinear problems and large-scale three-dimensional (3D) simulations. The computational burden becomes even more pronounced in applications that require rapid inference and real-time responses, such as inverse design and system nowcasting [5].

Over the last ten years, advances in artificial intelligence (AI) have encouraged researchers to explore AI-driven ways to solve partial differential equations (PDEs) in fluid mechanics, with the goal of improving the efficiency of traditional numerical methods. Among these methods, one of the most representative developments is the physics-informed neural network (PINN) [6–8]. PINNs use automatic differentiation to construct loss functions from PDE formulations, thereby integrating physical laws into neural network models. Compared to traditional numerical methods, PINNs avoid the complicated procedures of mesh generation and discretization, offering remarkable flexibility when dealing with complex geometries and high-dimensional problems [9, 10]. Moreover, PINNs show advantages in inverse problems, where they can infer equation parameters or structures from high noise or sparse data [11, 12]. However, applying PINNs to new CFD tasks requires significant time to retrain the neural network. The training process is also sensitive to hyperparameter tuning, and optimizing the loss function can be unstable, which limits their ability in real-world scenarios [13, 14].

Unlike PINNs, which learn the mapping from spatiotemporal coordinates to corresponding values, neural operators aim to learn mappings between function spaces, such as mapping the initial fields to their solutions [15]. For example, DeepONet [16] and its variants [17, 18] use branch and trunk networks to embed the initial conditions and spatiotemporal coordinates, respectively. Fourier Neural Operator (FNO) [19] and its variants [20, 21] use Fast Fourier Transformation (FFT) to conduct the inference stage in the frequency domain, which can handle different frequency components efficiently. Furthermore, several works use graph neural networks [22, 23] or Transformers [24, 25] as backbone models to extend neural operator frameworks to complex geometry problems. Recently, with the development of AI foundation models, researchers are working on unified neural operators capable of handling various physical simulation tasks, such as TL-DeepONet [26], MPP [27], OmniArch [28], DPOT [29], and the In-Context Operator Network [30]. Once the training stage is completed, these neural operators can rapidly solve physical equations through surrogate modeling, with inference times typically measured in milliseconds. As a result, neural operators can significantly improve the efficiency of fluid simulations and be successfully applied in fields such as weather forecasting [31] and system control [32]. However, most existing neural operators require large amounts of data for training to achieve reasonable performance, and acquiring and storing such scientific computing data can be extremely costly. Additionally, these methods primarily focus on 2D scenarios, and extending them to 3D often faces computational efficiency bottlenecks.

To reduce the data dependency of neural operators, physics-aware operator learning seeks to enable few-shot or even zero-shot learning by integrating physical knowledge [33]. Depending on how they incorporate physical priors, these methods can be divided into physics-informed and physics-encoded methods. Physics-informed methods embed PDE residuals or classical discretization schemes as soft constraints directly within the loss function (e.g., PI-DeepONet [34, 35], PINO [36, 37], MCNP [38]). Meanwhile, physics-encoded methods incorporate the core of a traditional numerical solver (FDM [39–42], FVM [43–45], or FEM [46, 47]) as part of the network architecture and then employ a neural network to refine the coarse-grid solution. While these physics-aware methods can train CFD surrogates with limited data, they encounter the following three principal challenges in practical applications. (1) Trade-off between physical prior and training/inference speed. Both physics-informed (soft-constraint) and physics-encoded (hard-constraint) approaches couple neural networks with traditional CFD solvers, which prevents substantial acceleration in training or inference without compromising stability. As a result, most physics-aware methods achieve only marginal speed-ups compared to their underlying numerical schemes [48]. (2) Mismatch between high-fidelity data and coarse-grid priors. In realworld scenarios, we often have access to highly accurate measurements, while the numerical prior embedded in physics-aware operators typically reflects a relatively low-order, coarser discretization. Consequently, there is a mismatch between the high-fidelity data and the low-fidelity physical prior. When jointly optimizing data and physics losses, as in PINO [36, 37], balancing these competing objectives becomes challenging, leading to unstable training dynamics and suboptimal surrogate performance. (3) Conflict between model capacity and inference efficiency for diverse CFD tasks. Expanding a single surrogate to handle diverse CFD problems (e.g., various initial conditions, parameters, and forcing) demands a significantly larger parameter space to ensure sufficient model capacity. However, this growth in model size inevitably degrades inference speed, thereby limiting the practical efficiency of AI-driven methods in real-world tasks.

To this end, we propose OmniFluids, a unified operator learning framework that relies purely on physics to pre-train neural operators for efficiently solving a wide range of CFD forward and inverse problems. OmniFluids integrates a novel training paradigm with a carefully designed model architecture. On the training side, it combines physics-only pre-training, coarse-grid operator distillation, and few-shot fine-tuning, enabling fast inference, accurate prediction, and strong generalization under limited or even zero data supervision. On the architectural side, OmniFluids incorporates key components such as a mixture of operators (MoO), a multi-frame decoder, and factorized Fourier layers, enabling efficient and scalable modeling of diverse physical tasks–especially in 3D scenarios– while maintaining seamless integration with physics-based supervision. Compared to other physicsaware AI-driven approaches, OmniFluids offers the following three key advantages. (1) Efficient and zero-data pre-training. In the pre-training stage, we train neural operators solely based on governing equations. By predicting multiple time frames in parallel via the multi-frame decoder and employing operator distillation, the OmniFluids can break free from the constraints of finegrid spatiotemporal discretization and offer significantly faster inference without any data-driven supervision. (2) Seamless integration of physical priors and high-fidelity data. Rather than jointly optimizing physics and data losses, the training strategy in OmniFluids first leverages large-scale, physics-only pre-training to embed the knowledge of various fluid dynamics into the pretrained model. In the subsequent few-shot fine-tuning phase, the OmniFluids adapts the operator to task-specific, high-fidelity data using minimal labeled data, which can effectively bridge the gap between coarse physical priors and fine-grained datasets. (3) Expressive yet efficient modeling via MoO. To effectively handle a wide range of physical tasks within a unified model, we assemble a rich library of neural operators to ensure expressive capacity yet deploy only lightweight personalized operators at inference. This MoO design balances model capacity with inference efficiency, resolving the aforementioned trade-off when handling various CFD tasks.

We validate OmniFluids across a diverse set of 2D and 3D fluid dynamics benchmarks, encompassing a wide range of Reynolds numbers, initial conditions, and external forces. OmniFluids consistently outperforms existing AI-driven CFD methods in predictive accuracy and consistency with turbulence statistics. Even in the zero-shot regime, OmniFluids can accurately predict flow fields and their corresponding kinetic energy spectrum after physics-only pre-training. Compared to classical numerical solvers, OmniFluids achieves 10-100x speedups by eliminating the need for fine-grained spatiotemporal discretization. Moreover, it accurately recovers unknown PDE parameters from sparse and noisy observations, delivering high-fidelity forecasts and demonstrating strong

Figure 1: Schematic framework of OmniFluids. a. Teacher/student network architecture. OmniFluids employs a mixture of operators (MoO) backbone to flexibly handle diverse CFD tasks, alongside a multi-frame decoder that predicts K future steps in parallel. Factorized Fourier layers enable efficient and scalable modeling, especially in 3D. b. Training workflow for OmniFluids. (1) Physics-only pre-training. The teacher network is trained using purely physics-based loss functions with no reliance on data. (2) Operator distillation. A compact student model is distilled from the pre-trained teacher at coarser spatial and temporal resolutions, enabling faster inference without sacrificing accuracy. (3) Few-shot fine-tuning. When a small amount of task-specific data becomes available, the student model is fine-tuned in a few-shot regime to bridge any remaining gap between pre-training priors and high-fidelity datasets. Here, u t ^k denotes the high-resolution prediction at the k th framework predicted by the teacher network; u˜ t ^k denotes the downsampled field of the teacher's high-resolution prediction. u˜ s ^k is the low-resolution prediction produced by the student network.

robustness and broad applicability.

Results

The overall framework of OmniFluids

Fig. 1a illustrates the schematic of OmniFluids. The operator architecture of OmniFluids is constructed on a mixture of operators (MoO) framework for solving various CFD problems. A routing network takes the PDE parameters as input and dynamically assigns the task to appropriate specialized operators. During training, an abundant library of specialized operators enhances model capacity, while during inference, we consolidate this ensemble into a singular personalized operator to ensure runtime efficiency. In the final layer, we use a multi-frame decoder to forecast multiple future steps in parallel, ensuring temporal continuity among subsequent physical fields and averting the expensive step-by-step iteration process. Furthermore, we incorporate factorized Fourier layers [20] for dimension-wise decomposition, enabling efficient extension to the 3D scenario. The architecture of OmniFluids achieves a favorable balance among expressiveness, inference speed, and scalability, which can naturally integrate with physics-based loss functions for robust training.

Fig. 1b illustrates OmniFluids' three-stage training workflow: physics-only pre-training, operator distillation, and few-shot fine-tuning. In the pre-training phase, the teacher network learns solely from diverse governing equations by predicting multiple future time steps in parallel. This multiframe loss breaks free from the sequential time-stepping bottleneck of traditional solvers and other physics-aware methods. Next, during operator distillation, we distillate knowledge from the highresolution teacher to a student model by matching its outputs at coarser spatial and temporal resolutions. This compression yields a student operator capable of rapid inference on a coarse grid without sacrificing accuracy. Finally, in the few-shot fine-tuning stage, the student model is briefly adapted using a small set of high-fidelity samples. Unlike joint physics-data training, this pre-train-then-fine-tuning strategy more robustly reconciles coarse-grid priors with fine-resolution datasets. This training paradigm establishes a unified framework for learning accurate, fast, and generalizable surrogate models under unsupervised or few-shot settings, capable of handling various physical problems. More details are given in Methods.

Benchmark PDEs for evaluating OmniFluids

To evaluate the performance of OmniFluids, we conduct experiments on a diverse set of PDEs that represent classical problems in fluid dynamics, ranging from chaotic pattern formation to realistic incompressible flows. Specifically, we consider the 2D Kuramoto-Sivashinsky equations (KSEs), the 2D incompressible Navier-Stokes equations (NSEs) in velocity-vorticity form, and the 3D incompressible NSEs in velocity-pressure form. These systems are physically interrelated and capture a broad spectrum of fluid phenomena–from the emergence of spatiotemporal chaos to 3D turbulence–thus offering a comprehensive evaluation framework for CFD modeling and simulation. All datasets are generated with high-resolution spectral solvers, and the detailed settings for each system are provided in Supplementary Note A.

2D KSE. The 2D KSE is a fourth-order nonlinear PDE widely used as a canonical model for spatiotemporal chaos in dissipative systems, such as the propagation of flame [49]. It exhibits complex behaviors, including pattern formation, wave turbulence, and chaotic attractors [50], making it an ideal benchmark for understanding chaotic systems. The 2D KSE is given by

$$\frac{\partial u}{\partial t} + \alpha(\Delta u + \Delta^2 u) + \frac{\beta}{2} |\nabla u|^2 = 0, \mathbf{x} \in [0, 8\pi)^2, t \in [0, 5],$$ (1)

where upx, tq P R denotes the solution field, ∆ and ∆² represent the Laplacian and biharmonic operators, and ∇ is the spatial gradient operator. The parameters α and β control the strength of the linear and nonlinear terms. In the physics-only pre-training phase of OmniFluids, we sample u⁰ P G2DKS, α P r0.1, 0.5s and β P r0.1, 0.5s to cover a wide range of chaotic regimes, where G2DKS is a 2D Gaussian random field.

2D incompressible vorticity-velocity NSE. The 2D incompressible NSEs in the vorticityvelocity formulation serve as a compact yet expressive model for laminar and turbulent flows. The vorticity form eliminates pressure, reducing the system to a single scalar PDE that maintains essential nonlinear convection and viscous diffusion effects [51]. The governing equation is

$$\frac{\partial \omega}{\partial t} + \mathbf{u} \cdot \nabla \omega = \frac{1}{\text{Re}} \Delta \omega + f, \quad \omega = \nabla \times \mathbf{u}, \mathbf{x} \in [0, 2\pi)^2, \ t \in [0, 10],$$ (2)

3D incompressible velocity-pressure NSE. The 3D NSEs are a cornerstone of CFD, encapsulating the fundamental laws governing the motion of viscous, incompressible fluids [52]. Unlike 2D cases, the 3D form retains the full spatial complexity and nonlinear interactions that characterize real-world turbulent flows. Solving these equations accurately is critical for simulating real-world flows in various engineering and geophysical systems. The governing equation is

$$\frac{\partial \mathbf{u}}{\partial t} + \mathbf{u} \cdot \nabla \mathbf{u} = -\nabla p + \frac{1}{\text{Re}} \Delta \mathbf{u} + 0.1 \mathbf{u}, \quad \nabla \cdot \mathbf{u} = 0, \mathbf{x} \in [0, 2\pi)^3, \ t \in [0, 10],$$ (3)

Physics-only pre-training for CFD tasks

In the physics-only pre-training stage, OmniFluids is trained in an unsupervised regime using only physics priors, which is an essential capability when data are scarce or costly to acquire. Thanks to the MoO architecture, OmniFluids simultaneously learns across diverse physical settings (e.g., varying parameters and forcing fields) without sacrificing accuracy. Furthermore, a multi-frame decoder predicts multiple steps in a single forward rollout, which allows computing the physicsbased loss in parallel and avoids the computational burden of step-by-step time iteration. The loss function combines a pseudo-spectral discretization of spatial derivatives with a Crank-Nicolson scheme for temporal evolution. Here, we evaluate the performance of the physics-only pre-trained model on the aforementioned three benchmark datasets and compare it against the representative physics-informed methods, i.e., PINO [36].

2D KSE system. In the physics-only pre-training stage, we discretize the spatial domain into a 64ˆ64 grid. For temporal evolution, we construct the loss function over the interval t P r0, 0.2s and use a multi-frame decoder to divide the time dimension into 100 intervals. During testing, we extend the predictions to t " 5 to evaluate the model's ability to perform long-term extrapolation. The error evolution curves indicate that OmniFluids outperforms PINO in error accumulation, maintaining a much lower relative ℓ² error over time (Fig. 2a). The corresponding predicted snapshots show that OmniFluids preserves the fine-scale structures of the solution, closely matching the reference. On the other hand, PINO begins to deviate as the prediction progresses, struggling to capture important subsequent information during the long-term extrapolation.

2D vorticity-velocity NSE system. In this case, we train OmniFluids using a spatial discretization of 256 ˆ 256 grid points. For the temporal domain, the loss function is constructed over the interval t P r0, 0.2s, divided into 50 bins. Meanwhile, the predictions are extrapolated to

Figure 2: Multi-step error and predicted snapshots from PINO and OmniFluids across diverse CFD systems. a. 2D KSE. b. The vorticity field (ω) of 2D NSE. c. The first velocity component (u) of 3D NSE. Note that only a corner of the flow field is shown for the 3D NSE in c for better visualization.

t " 10 to evaluate the long-term forecasting ability of each model. This case needs the model to handle different levels of Reynolds numbers, external forces, and initial conditions, making it more challenging for neural operators to approximate and generalize capacity.

As shown in Fig. 2b, the error evolution curves reveal that the performance gap between PINO and OmniFluids becomes increasingly pronounced as time progresses. OmniFluids successfully captures the fine-scale structures of the vorticity field–such as coherent vortices, eddies, and turbulent filaments–and preserves their spatial integrity throughout the temporal evolution. Even at the final step, the predicted vorticity fields from OmniFluids remain sharply defined and visually indistinguishable from the reference. In contrast, PINO struggles with long-term extrapolation. Its predictions begin to diverge early from the reference solution, losing essential physical structures such as vortex cores.

Furthermore, we check how well the models can keep the flow's spectral and statistical features at the last time step (Fig. 2d). The kinetic energy spectrum Epkq characterizes the energy distribution across spatial frequencies. OmniFluids closely follows the reference across all wavenumbers, accurately capturing both large-scale motions and small-scale turbulent features. By contrast, PINO underestimates the energy at high wavenumbers, reflecting its inability to maintain smallscale structures, indicating spectral dissipation. The probability density function (PDF) of the normalized vorticity ω further illustrates the statistical fidelity of each model. OmniFluids matches the reference PDF well, reproducing the sharp peaks and heavy tails associated with intermittent turbulent structures. Meanwhile, PINO's PDF deviates near the center and tails, indicating a narrower and less accurate distribution of vorticity values.

3D velocity-pressure NSE system. For the 3D case, we extend OmniFluids to a spatial resolution of 128³ grid points. The training time domain spans t P r0, 0.2s, discretized into 40 intervals, and predictions are extrapolated to t " 10. Compared to 2D settings, 3D fluid modeling is inherently more challenging due to the vortex stretching, anisotropic multi-scale dynamics, and a forward energy cascade, all of which require capturing intricate spatial and temporal correlations [53].

As shown in Fig. 2c, PINO fails to maintain physically consistent predictions beyond the shortterm horizon. The velocity fields become increasingly diffused and decor-related, losing essential 3D structures. On the other hand, OmniFluids accurately preserves the fine-scale features of the flow, like aligned turbulent filaments, remaining nearly indistinguishable from the reference even at t " 10. Quantitatively, OmniFluids maintains a stable relative ℓ² error below 0.25 over the entire extrapolation domain. In contrast, PINO's error grows rapidly, exceeding 1.0 after t ą 4, reflecting a total loss of pattern coherence and physical fidelity.

Spectral analysis further reveals the superior accuracy of OmniFluids (Extended Data Fig. 1). OmniFluids matches the reference solution across all resolved wavenumbers, with only minor deviations at the highest frequencies. In contrast, PINO deviates significantly across all spectrums, from low to high frequencies, indicating an inability to capture energy distribution across scales. Additionally, the PDF of normalized velocity magnitude shows that OmniFluids closely replicates the reference statistical distribution, while PINO exhibits a pronounced shift and fails to reproduce the correct skewness characteristics of turbulent velocity fields.

Adaptation to coarser resolutions via distillation and fine-tuning

After the physics-only pre-training stage, we perform operator distillation to transfer knowledge from a high-resolution teacher model to a student model designed to operate at coarser spatial and temporal resolutions. For a specific downstream task, the distilled student model is then fine-tuned using a small number of data samples from the target distribution. This two-stage adaptation strategy enables OmniFluids to generalize efficiently to new tasks with minimal data, leveraging prior physical knowledge while remaining adaptable to novel regimes. In this section, we evaluate the adaptability of OmniFluids on 2D KSE, 2D NSE, and 3D NSE, including both in-distribution (ID) and out-of-distribution (OOD) adaptation scenarios. For the 2D cases, we compare the performance of OmniFluids with various other approaches, including state-of-the-art neural operators such as PINO [36], DeepONet [16], FNO [19], FFNO [20], FactFormer [24], and CNO [54]; a large-scale datadriven pre-trained foundation model (DPOT [29]); and direct numerical simulation (DNS) based on spectral methods, as detailed in Supplementary Note B. For the 3D case, we use PINO [36],

Figure 3: Predicted snapshots, multi-step errors, and inference speed comparison of baseline methods and OmniFluids on 2D KSE. a and c. In-distribution snapshots (t " 5) and relative ℓ² errors with parameters α " 0.2, β " 0.5. b and d. Out-of-distribution snapshots (t " 5) and relative ℓ² errors with parameters α " 1.0, β " 1.0. e. Comparison of inference speed across different methods.

FNO [19], DPOT [29], and DNS as baselines for comparison.

2D KSE system. In the operator distillation stage, we compress the model to a coarser spatiotemporal resolution of 32 ˆ 32 with a time step ∆t " 1.0. We evaluate the model under indistribution (ID) and out-of-distribution (OOD) settings, using selected values of α " 0.2, β " 0.5 for ID, and α " 1.0, β " 1.0 for OOD.

The task is comparatively easier in the ID scenario, where the fourth-order dissipation and nonlinear advection are relatively weak. As shown in Fig. 3a and c, OmniFluids successfully reconstructs spatial structures and preserves fine-scale features, achieving the lowest relative ℓ² error and exhibiting stable long-term rollouts among all AI-driven methods. In contrast, purely data-driven baselines generally underperform due to their dependency on extensive training data. In particular, PINO struggles to enforce accurate physics-based constraints because inconsistencies often arise between the physics-informed loss and the high-fidelity data, leading to suboptimal training. DPOT achieves the best performance after OmniFluids among neural operator baselines, thanks to its pre-training on extensive datasets. However, the absence of physical knowledge limits its ability to capture delicate spatiotemporal dynamics.

Under the more challenging OOD conditions (Fig. 3b and d), baseline models' errors increase significantly, often growing by an order of magnitude as the system becomes more nonlinear and chaotic. Nevertheless, OmniFluids can accurately capture the dominant physical modes and maintain coherent spatial patterns. In contrast, DPOT exhibits grid-like artefacts, while FNO and DeepONet deteriorate rapidly into noise-like outputs, indicating a lack of generalization.

Finally, as shown in Fig. 3e, OmniFluids achieves a favorable balance between accuracy and computational efficiency. Its inference time is orders of magnitude faster than DNS (0.034 s vs. 3.495 s). We fix the spatial resolution across methods to ensure a fair comparison and adopt the largest stable time step for DNS. Because the KSE equation includes a fourth-order derivative, DNS

Figure 4: Predicted snapshots, flow statistics, temporal correlation, and inference speed of OmniFluids and baseline methods on 2D NSE. a and b. Predicted velocity (u " ru, vs) and vorticity fields (ω) at t " 10 under ID (Re " 2000) and OOD (Re " 4000) settings, respectively. c and d. Kinetic energy spectrum and normalized vorticity PDFs for ID and OOD scenarios. e and f. Temporal evolution of correlation coefficients between predicted and reference vorticity fields for ID and OOD scenarios. g. Inference time comparison across all methods.

requires a fine-grained step size to remain stable; otherwise, the numerical method quickly blows up. In contrast, Omnifluids overcomes the limitations of spatiotemporal discretization, achieving a 100x speedup compared to the DNS method.

2D vorticity-velocity NSE system. In the operator distillation stage, we compress the pretrained model to a coarser spatiotemporal resolution of 128 ˆ 128 with a time step ∆t " 1.0. We evaluate the model under ID and OOD settings, using selected Reynolds numbers of Re " 2000 for ID and Re " 4000 for OOD, respectively.

The predicted vorticity fields ω are shown in Fig. 4a and b. For a more comprehensive assessment, we solve a Poisson equation to reconstruct the velocity fields (u " ru, vs) and visualize them alongside the vorticity. In this part, even though the Reynolds number is fixed for each downstream task, this setup still demands strong generalization capability due to the variation in external forcing conditions. Among the baseline methods, DeepONet, FNO, and CNO fail to generate meaningful predictions due to limited training data and insufficient model expressivity. FactFormer and FFNO, while equipped with stronger representation capacity, only manage to localize a few primary vortices and still exhibit apparent discrepancies from the reference solution, especially under the more challenging Re " 4000 setting. DPOT, which benefits from extensive pre-training on large datasets, shows moderate adaptability to new regimes via fine-tuning. However, it fails to capture vortex positions and intensity accurately due to a lack of explicit physical priors. PINO performs better than most AI baselines, especially in representing the overall structure, but suffers from blurring and distortion when modeling fine-scale vortex features. In contrast, OmniFluids consistently delivers high-fidelity reconstructions, capturing intricate vorticity structures even in turbulence with a high Reynolds number, and demonstrates superior resolution of multiscale dynamics. Temporal correlation curves (Fig. 4e and f) illustrate OmniFluids' ability to maintain accurate long-term flow predictions. It consistently achieves the highest correlation with the reference solution over time among all AI-driven methods, whereas other models exhibit rapid degradation.

Fig. 4c and d present the kinetic energy spectra and normalized vorticity PDFs for the ID and OOD settings. OmniFluids accurately reproduces the reference spectra across the full range of wavenumbers, including high-frequency components, and even outperforms DNS evaluated at the same resolution. This superior spectral accuracy is attributed to the model's pre-training on highresolution simulations and subsequent fine-tuning on high-fidelity samples. Similar advantages are observed in turbulence statistics: OmniFluids faithfully recovers the reference vorticity distribution, precisely capturing both the sharp peak and the heavy-tailed behavior. In contrast, other AI-driven methods exhibit noticeable deviations, with shifted peaks and distorted tails, underscoring their limitations in modeling the underlying statistical structure of turbulent flows. Finally, OmniFluids achieves an inference time of only 0.136 s, compared to 3.627 s for DNS, yielding over 26x speedup (Fig. 4g). These results demonstrate that OmniFluids can generalize effectively across physical regimes and faithfully capture multiscale and statistical features of turbulent flows, offering a robust and efficient alternative to classical CFD solvers and existing AI-driven approaches.

3D velocity-pressure NSE system. We evaluate model performance on the 3D NSE with coupled velocity-pressure dynamics. Compared to the 2D case, modeling 3D turbulence is more challenging due to the increased degrees of freedom and nonlinearity, placing greater demands on the capacity and stability of learned surrogate models. During distillation, the pre-trained operator is compressed to a coarse spatiotemporal resolution of 128³ with a time step ∆t " 1.0. We consider both ID and OOD generalization settings for evaluation, corresponding to Reynolds numbers of Re " 2000 and Re " 3000.

As shown in Fig. 5a and b, OmniFluids achieves superior accuracy in predicting the 3D velocity field u " ru, v, ws over time, maintaining close agreement with the reference solution even at later stages. Streamtube visualizations of the normalized velocity magnitude in Extended Data Fig. 2 further highlight OmniFluids' ability to recover fine-scale coherent structures within the flow. In contrast, PINO, FNO, and DPOT exhibit rapid degradation in predictive accuracy. Their temporal correlation coefficients fall below 0.8 within the first one or two seconds, and performance deterio-

Figure 5: Predicted snapshots, flow statistics, temporal correlation, and inference speed of OmniFluids and baseline methods on 3D NSE. a and b. Temporal evolution of correlation coefficients and predicted velocity fields (u " ru, v, ws) at t " 8 under ID (Re " 2000) and OOD (Re " 3000) settings, respectively. Note that only a corner of the flow field is shown for better visualization. c and d. Kinetic energy spectrum for ID and OOD scenarios. e. Inference time comparison across all methods.

rates further under the more demanding Re " 3000 regime. Kinetic energy spectra in Fig. 5c and d reveal the model's ability to recover multiscale dynamics. Only OmniFluids and DNS accurately reproduce the reference spectrum across the full range of wavenumbers, including high-k components critical for capturing small-scale turbulence. Moreover, inference time comparisons in Fig. 5e show that OmniFluids achieves a nearly 11x speed-up over DNS (0.843 s vs. 9.382 s).

Inverse analysis of fluid systems with OmniFluids

Inverse problems in fluid dynamics involve inferring latent system parameters (e.g., viscosity or forcing) from limited, noisy observations of the flow field. These problems are both critical and chal-

Figure 6: Inverse analysis of fluid systems with OmniFluids. a. The OmniFluids framework for solving inverse problems. b-c. Inferenced viscosity coefficient and external forcing using two frames of flow field data under three noise levels (E1: 0% noise; E2: 5% noise; E3: 10% noise). d. Prediction of future flow evolution based on the inferred viscosity coefficient and external forcing. e. Error evolution and correlation analysis.

lenging: accurate parameter recovery is essential for applications ranging from subsurface reservoir characterization to aerodynamic design, yet the sparsity and noise inherent in real measurements often render classical inversion methods ill-posed and computationally prohibitive [55, 56].

Leveraging its broad generalization and differentiability, OmniFluids efficiently solves inverse problems by recovering unknown parameters even from sparse or noisy data across diverse physical regimes. As depicted in Fig. 6a, we begin with an initial guess of the unknown parameters, feed them into OmniFluids to generate a forward simulation, and then compare the simulated fields against observational data. The resulting difference is used to compute gradients with respect to the parameters, which are updated iteratively via gradient-based optimization until convergence. Equipped with a distilled, coarse grid student operator, OmniFluids achieves rapid forward inference, thereby accelerating the optimization loop for inverse problems. Once the parameters are inferred, OmniFluids can immediately forecast the subsequent evolution of the underlying systems with the newly recovered conditions.

We herein assess OmniFluids' performance in solving inverse problems for the 2D incompressible vorticity-velocity NSE under three levels of observational noise: 0 % (E1), 5 % (E2), and 10 % (E3). In these experiments, we assume access to only two frames of flow field data: the flow field at the initial time step (t " 0) and at a later time step (t " 1), providing a sparse temporal dataset for parameter inference. The spatial resolution of the observed flow fields is 128ˆ128, aligning with the student network. Using the observed data, we aim to infer two key unknown parameters, including the viscosity coefficient (ν " Re´¹ ) and the external forcing (f). As shown in Fig. 6b, the inferred viscosity values under different noise levels closely match the reference value, with a relative error of only around 2.8 % even in the 10 % noise scenario. Similarly, Fig. 6c compares the inferred external forcing fields with the reference forcing. Despite increasing noise, OmniFluids reliably reconstructs the spatially complex forcing patterns, demonstrating robustness to data corruption. Once parameters are inferred, we feed them back into the model to predict future flow evolution. Fig. 6d presents predicted flow fields from t " 2 to t " 10. The close visual agreement confirms that the inferred parameters enable accurate long-term forecasts of the dynamics under different noise levels. Quantitatively, we observe that OmniFluids' prediction accuracy degrades slightly as the noise level increases, which is to be expected; nevertheless, the correlation coefficient exceeds 0.9 in all scenarios, demonstrating that the model still captures the majority of flow patterns even under substantial observational corruption (Fig. 6e).

Discussion

This paper introduces OmniFluids, a unified physics-driven operator learning framework designed for efficient and generalizable modeling of complex fluid dynamics. At its core, OmniFluids leverages a carefully designed model architecture and a three-stage training strategy to overcome key limitations in current CFD surrogate methods. From an architecture perspective, OmniFluids adopts a multi-frame prediction strategy that enables parallel inference over time steps, avoiding costly sequential rollout and improving training stability. A factorized Fourier layer design enhances the scalability of 3D simulations by reducing the computational burden of spectral convolution. To support diverse physical scenarios, we introduce a mixture of operators architecture that maintains sufficient model capacity during training while enabling lightweight, task-specific operators for fast inference. In terms of training, OmniFluids first undergoes a physics-only pretraining phase, where it learns diverse surrogate models purely from the governing PDEs without any labeled data. OmniFluids efficiently integrates physical constraints as differentiable losses due to the multi-frame prediction strategy, avoiding iterative solvers and enabling stable, fully parallel training. We then apply operator distillation, compressing the high-resolution pre-trained model into a lower-resolution surrogate that retains fine-scale physical knowledge while supporting fast coarse-grid inference. Finally, a small amount of task-specific data (as few as 2-10 trajectories) is used for few-shot fine-tuning, allowing the model to reconcile discrepancies between the physics prior and real-world measurements. Furthermore, thanks to its differentiable design, support for multi-physics scenarios, and fast inference, OmniFluids is particularly well-suited for solving inverse problems, including identifying unknown PDE parameters or external forcings.

OmniFluids demonstrates strong performance across a broad range of challenging CFD benchmarks. It consistently outperforms neural operator baselines such as FNO, FFNO, and PINO, as well as large-scale data-driven PDE foundation models like DPOT, regarding accuracy, long-term stability, and spectral fidelity. It generalizes effectively across diverse PDE parameters, initial conditions, and external forces under few-shot adaptation, capturing fine-scale, nonlinear dynamics with minimal supervision. OmniFluids can conduct 3D turbulence simulations efficiently, a long-standing grand challenge due to the extreme dimensionality and multiscale nature. At Reynolds numbers up to 3000, it maintains accurate long-horizon forecasts while preserving spatiotemporal coherence and energy spectra, outperforming strong baselines in accuracy and physical fidelity. Beyond forward simulation, OmniFluids also excels at inverse problems such as system identification and parameter inference. Leveraging pre-trained operators, it efficiently infers critical system properties, such as viscosity coefficients and external forces, demonstrating its ability to robust downstream predictions and reconstructions in data-scarce or partially known environments. Furthermore, OmniFluids breaks the constraints of traditional CFD solvers, delivering 10-100x acceleration while retaining high-fidelity spectral characteristics and physical consistency at low spatial and temporal resolutions. Together, these results position OmniFluids as a unified, scalable, and data-efficient operator framework for both forward and inverse modeling of complex, multiscale fluid dynamics.

Despite its strong empirical performance and diversity, OmniFluids still has several limitations that suggest directions for future work. First, the current framework assumes that the governing equations are fully known and explicitly specified during pre-training. This assumption may not hold in gray-box or partially known systems, such as real-world settings with missing or uncertain physical terms. While OmniFluids demonstrates some robustness to such cases through adaptation, its performance could degrade without accurate physical priors. Second, the implementation of OmniFluids is tailored to structured Cartesian grids, which limits its direct applicability to domains involving irregular geometries or curved boundaries in other CFD problems. Extending OmniFluids to such settings will necessitate introducing geometric inductive biases through graph neural operators, mesh-aware architectures, or coordinate-free representations, which can better handle spatially complex domains while preserving physical accuracy. Third, although the current pretraining strategy leverages rich PDE dynamics, it does not explicitly encode additional physical constraints such as conservation laws (e.g., mass, momentum, energy), symmetries, or entropy conditions. These constraints are often crucial for ensuring robustness, stability, and generalization in real-world scenarios. Integrating such domain-specific inductive priors could further enhance the framework's reliability across broader applications.

Methods

We herein introduce the method of the proposed OmniFluids.

Network architecture

To address the challenges of modeling complex fluid systems, our network architecture is designed to satisfy three key requirements. First, it should enable rapid inference to efficiently integrate physics priors and scale across extensive spatiotemporal domains. Second, it should enjoy the sufficient capacity to capture diverse fluid behaviors under varying parameters, initial conditions, and external forcings. Third, the design can efficiently handle 3D CFD tasks, preserving essential turbulence statistics–such as the kinetic energy spectrum–at high fidelity. Guided by these principles, we design the OmniFluids architecture (Fig. 1a), which can deliver accurate, efficient, and generalizable modeling for CFD tasks. Since the teacher and student networks differ only in the decoder's resolution, we describe the teacher network in detail.

Spatially factorized spectral layer. We perform primary spatial operations in the Fourier domain to ensure the accurate preservation of high-frequency structures and turbulence statistics. Following the FFNO paradigm, we decompose the transform into two or three separable 1D spectral operations rather than using a 2D or 3D Fourier transformer directly [20]. For 3D cases, given a latent input vpx, y, zq, we compute its 1D Fourier coefficients along each coordinate axis as follows:

$$\hat{\mathbf{v}}_x(k;y,z) = \mathcal{F}_x\mathbf{v}(\cdot,y,z), \hat{\mathbf{v}}_y(\ell;x,z) = \mathcal{F}_y\mathbf{v}(x,\cdot,z), \hat{\mathbf{v}}_z(m;x,y) = \mathcal{F}_z\mathbf{v}(x,y,\cdot).$$ (4)

Each set of coefficients is then modulated by a learnable multiplier (Rxpkq, Rypℓq, Rzpmq) to capture spatial correlations in the frequency domain. The inverse transforms along each axis are summed to produce the updated field:

$$\mathbf{v}_{\text{update}}(x, y, z) = \mathcal{F}_x^{-1} \left[ \mathcal{R}_x \hat{\mathbf{v}}_x \right] + \mathcal{F}_y^{-1} \left[ \mathcal{R}_y \hat{\mathbf{v}}_y \right] + \mathcal{F}_z^{-1} \left[ \mathcal{R}_z \hat{\mathbf{v}}_z \right]. \tag{5}$$

After that, we wrap the separable spectral transform in a residual block with Layer Normalization [57] to stabilize training and improve representational capacity:

$$\mathbf{v}{\text{out}} = \mathbf{v} + \text{LayerNorm}(\mathbf{v}{\text{update}}).$$ (6)

Compared to the original FNO's full-rank 3D transform [19], this separable implementation processes each axis independently. By retaining d modes per dimension, we reduce parameter complexity from Opd 3 q to Opdq, which is particularly impactful for 3D fields.

We employ a patch-based embedding strategy to further accelerate inference and reduce computational costs for large-scale spatial domains. Specifically, the input physical field is downsampled into coarse spatial patches before passing through the factorized spectral layers. This operation reduces the spatial resolution during spectral processing, substantially lowering the computational and memory demand. After spectral modulation, we apply a pixel shuffle-based upsampling layer to reconstruct the original resolution from the latent feature space. This two-stage structure enables OmniFluids to operate effectively in the latent space while preserving critical fine-scale structures.

Multi-frame decoder. To avoid the computation overhead and error accumulation, we employ a 1D convolutional decoder [22] that generates multiple future time steps in a single forward pass. By predicting K future states at once, this decoder accelerates inference and enables fully parallel computation of the physics-based loss function across all frames. Moreover, the convolutional kernels enforce temporal continuity between adjacent predictions, significantly enhancing stability during long-term rollouts.

During physics-only pre-training, the decoder outputs all K frames in parallel, while at the inference or operator distillation stage, we typically require only the final (Kth) frame. Based on this observation, we introduce a split trick that accelerates inference without altering the pretrained model to eliminate this computational redundancy. Concretely, we factor the decoder into two MLP branches, MLP¹ and MLP2, followed by a 1D convolution. In the pre-trained stage, the model output K frames in parallel as follows:

$$\left[\mathbf{u}{1}^{t}, \mathbf{u}{2}^{t}, \dots, \mathbf{u}{K}^{t}\right] = \text{Conv1D}\left(\text{MLP}{1}(\mathbf{v}{\text{out}})||\text{MLP}{2}(\mathbf{v}_{\text{out}})\right), \tag{7}$$

where || denotes the concatenation operation along the channel dimension. While during the inference stage, we bypass MLP1, and directly use MLP² to predict only the final frame:

$$\mathbf{u}_K^t = \text{Conv1D}(\text{MLP}2(\mathbf{v}{\text{out}})).$$ (8)

This split trick reduces computation and memory overhead during inference and subsequent operator distillation, yielding remarkable speedups for 3D fluid simulations.

Mixture of operators. A core challenge in designing a network that handles diverse CFD tasks is balancing inference speed against representational capacity. Equipping a single model to address multiple complex regimes demands ample parameter space, which degrades inference efficiency. To this end, we introduce a mixture of operators (MoO) framework. Rather than relying on one monolithic network, we maintain a bank of P expert operator branches, each specialized for distinct physical regimes. At inference time, a lightweight routing network dynamically assigns the task to each operator and integrates its parameters into a unified model, thereby balancing inference speed with representational capacity.

As shown in Fig. 1a, the routing network takes the PDE parameters γ as input and encodes them into a set of logits tzpu P p"1 . These logits are normalized via a temperature-scaled softmax to produce non-negative weights twpu P p"1 satisfying ^ř^P ^p"¹ w^p " 1:

$$w_p = \frac{\exp(zp/\tau)}{\sum{j=1}^P \exp(z_j/\tau)},\tag{9}$$

where τ is a temperature hyperparameter that controls the sharpness of the weight distribution. The final operator used for prediction is a weighted combination of the expert parameters, i.e.,

$$\bar{\mathbf{\Theta}} = \sum_{p=1}^{P} w_p \,\mathbf{\Theta}_p,\tag{10}$$

where Θ^p denotes the p th operator in the operator pool. In this way, the large operator pool is compressed into a personalized operator with parameters Θ¯ that is responsible for tasks governed by PDE parameters γ. This task-adaptive mechanism preserves the expressiveness of a large expert ensemble while ensuring high inference efficiency. By dynamically selecting and integrating specialized subnetworks based on the underlying physical characteristics, the MoO architecture enables accurate and scalable modeling across diverse CFD tasks.

Training framework

OmniFluids is trained through a three-stage pipeline consisting of physics-only pre-training, operator distillation, and few-shot fine-tuning to address the key challenges of neural operator learning in fluid dynamics (Fig. 1b). In the physics-only pre-training stage, a high-capacity teacher network is trained exclusively on various PDEs without any observational data, thereby removing reliance on labeled datasets and embedding strong physics priors. Next, operator distillation compresses this teacher into a student model that operates on coarser spatial and temporal resolutions, further accelerating inference while retaining the teacher's predictive accuracy. Finally, when a small set of high-fidelity observations becomes available, few-shot fine-tuning adjusts the distilled student network with only a handful of samples, aligning the model with empirical data and delivering an operator that combines physical inductive biases with data guidance.

Physics-only pre-training. In the physics-only pre-training stage, we leverage the designed OmniFluids architecture to learn the fluid operator directly from the PDE information without any data. At each training epoch, we sample an initial field u⁰ from a prescribed distribution and draw PDE parameters γ from their respective priors, ensuring the model generalizes across diverse physical tasks. Given an initial field u⁰ on a regular grid, the network predicts K future states tu t k K k"1 in parallel via a multi-frame decoder. For spatial derivatives, we employ the pseudo-spectral method, which exhibits spectral (i.e., exponential) convergence and captures global flow features, making it the state-of-the-art choice on uniform grids [48, 58]. Concretely, for each predicted field u t k , we compute the right-hand side of the governing PDE (N puˆ t k q) in Fourier space to obtain high-accuracy estimates of advection, diffusion, and forcing terms.

We adopt the Crank- Nicolson scheme to integrate time more accurately than first-order methods. Given a time step ∆t, the physics loss LPDE is defined by the mean squared residual of the time-discretized PDE across all frames:

$$\mathcal{L}{\text{PDE}} = \frac{1}{K} \sum{k=1}^{K} \left| \frac{\hat{\mathbf{u}}{k}^{t} - \hat{\mathbf{u}}{k-1}^{t}}{\Delta t} - \frac{1}{2} \left( \mathcal{N}(\hat{\mathbf{u}}{k}^{t}) + \mathcal{N}(\hat{\mathbf{u}}{k-1}^{t}) \right) \right|_{2}^{2}, \tag{11}$$

where we abbreviate uˆ⁰ as uˆ t 0 for simplicity. Because OmniFluids outputs all K frames in a single forward pass, these computations can be performed fully in parallel, eliminating the need for stepby-step iteration and greatly simplifying training.

Operator distillation. In this stage, we aim to compress the high-capacity pre-trained teacher network into an efficient student operator via knowledge distillation [59]. The goal is to preserve the teacher's predictive fidelity while significantly reducing computational cost, notably by enabling inference on coarser spatial and temporal grids. To achieve this, we fix the pre-trained teacher and train the student network to match its outputs under downsampled conditions. We sample the same initial condition u⁰ and parameters γ as in the pre-training stage, but evaluate the difference between the teacher and student models on a coarsened grid and with a larger time step M∆t with M " 1. The distillation loss aligns the student's predictions uˆ s MK with the downsampled teacher outputs via a mean squared error:

$$\mathcal{L}{\text{Dist.}} = \left| \tilde{\mathbf{u}}{MK}^s - \mathcal{D}(\mathbf{u}_{MK}^t) \right|_2^2, \tag{12}$$

where Dp¨q denotes the spatial downsampling operator applied to the teacher's output. As a result, operator distillation yields a model adapted to coarse spatial and temporal resolutions, supporting fast deployment while preserving the key information from the high-resolution teacher network.

Few-shot fine-tuning. Once a small set of high-fidelity observations becomes available, we fine-tune the distilled student model to further enhance its accuracy. This stage requires only a handful of ground-truth trajectories (e.g., 2-10), enabling the model to align its predictions more closely with empirical data and correct any residual errors from the physics-driven prior. Finetuning is performed end-to-end using a standard mean squared loss on predicted fields (LData). We adopt a small learning rate during this stage to ensure stable adaptation and avoid overwriting pre-trained knowledge. This training stage preserves the inductive efficiency and generalization capabilities acquired during physics-only pre-training and operator distillation, while tailoring the model to the specific target task.

Inverse Problem

Inverse problems play a central role in scientific modeling, enabling the identification of unknown physical parameters from limited or noisy observations. In fluid systems, solving such problems is critical for system identification, model personalization, and downstream control or forecasting tasks. Leveraging the generalization capability of the distilled student model, OmniFluids provides an efficient and principled approach for inverse analysis. The key idea is to treat unknown physical parameters, such as viscosity ν or external forcing, as differentiable inputs to the model and optimize them by minimizing discrepancies between predicted and observed flow fields.

As illustrated in Fig. 6a, the process begins with an initial guess γ of the unknown parameters. Given an initial condition u˜0, the student model performs a rollout to predict the final state u˜ s T pγq after M steps. The inverse loss LDiff. measures the difference between the observed final state u˜^T and the predicted one:

$$\mathcal{L}_{\text{Diff.}}(\gamma) = |\tilde{\mathbf{u}}_T^s(\gamma) - \tilde{\mathbf{u}}_T|_2^2. \tag{13}$$

This loss is minimized via gradient descent:

$$\gamma \leftarrow \gamma - \eta \nabla{\gamma} \mathcal{L}{\text{Diff.}}(\gamma),$$ (14)

where η is the learning rate. Thanks to the fully differentiable architecture of OmniFluids, gradients with respect to the physical parameters can be computed via automatic differentiation. Compared to traditional simulation-based inverse solvers, OmniFluids is highly efficient for inverse problems, where all forward and backward steps are computed in single or few passes, allowing for rapid optimization over physical parameters. Once the optimal γ is recovered, it can be used for further prediction, simulation, or control, making OmniFluids a powerful tool for data-driven discovery in complex fluid systems.

Evaluation metrics

We use two primary metrics to quantitatively evaluate the prediction accuracy of each model, including the relative ℓ² error and correlation coefficient.

Relative ℓ² error measures the normalized difference between the predicted solution upre and the reference solution u. It is defined as

Relative $$\ell2$$ error = $\frac{|\mathbf{u}{\text{pre}} - \mathbf{u}|_2}{|\mathbf{u}|_2} = \frac{\sqrt{\sumi (\mathbf{u}{\text{pre}(i)} - \mathbf{u}_{(i)})^2}}{\sqrt{\sumi \mathbf{u}{(i)}^2}},$

where the sums run over all spatial points in the solution.

Correlation coefficient evaluates the linear correlation between the predicted and reference solutions, defined by

$$\operatorname{Corr}(\mathbf{u}{\operatorname{pre}}, \mathbf{u}) = \frac{\sum{i} (\mathbf{u}{\operatorname{pre}(i)} - \overline{\mathbf{u}}{\operatorname{pre}}) (\mathbf{u}{(i)} - \overline{\mathbf{u}})}{\sqrt{\sum{i} (\mathbf{u}{\operatorname{pre}(i)} - \overline{\mathbf{u}}{\operatorname{pre}})^{2}} \sqrt{\sum{i} (\mathbf{u}{(i)} - \overline{\mathbf{u}})^{2}}},$$

where upre and u¯ denote the mean values of upre and u, respectively.

Baseline models

In the physics-only pre-training stage, we use PINO [36] as a baseline, which can learn CFD surrogate models without data supervision. When transferring to specific tasks, for the 2D cases, we compare the performance of OmniFluids with various other approaches, including state-of-the-art neural operators such as PINO [36], DeepONet [16], FNO [19], FFNO [20], FactFormer [24], and CNO [54]; a large-scale data-driven pre-trained foundation model (DPOT [29]); and direct numerical simulation (DNS) based on spectral methods. For the 3D case, we use PINO [36], FNO [19], DPOT [29], and DNS as baselines for comparison. The introduction and detailed settings of the baseline models are found in Supplementary Note B.

Data availability

All the used datasets are available on GitHub at https://github.com/intell-sci-comput/ OmniFluids.

Code availability

All the source codes to reproduce the results in this study are available on GitHub at https: //github.com/intell-sci-comput/OmniFluids.

References

[1] L Masini, S Timme, and AJ Peace. Analysis of a civil aircraft wing transonic shock buffet experiment. J. Fluid Mech., 884:A1, 2020.
[2] Peter Bauer, Alan Thorpe, and Gilbert Brunet. The quiet revolution of numerical weather prediction. Nature, 525(7567):47–55, 2015.
[3] Jiangfan Yu, Dongdong Jin, Kai-Fung Chan, Qianqian Wang, Ke Yuan, and Li Zhang. Active generation and magnetic actuation of microrobotic swarms in bio-fluids. Nat. Commun., 10(1):5631, 2019.
[4] Harvard Lomax, Thomas H Pulliam, David W Zingg, and TA Kowalewski. Fundamentals of computational fluid dynamics. Appl. Mech. Rev., 55(4):B61–B61, 2002.
[5] George Em Karniadakis, Ioannis G Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nat. Rev. Phys., 3(6):422–440, 2021.
[6] Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. J. Comput. Phys., 378:686–707, 2019.
[7] Xiaowei Jin, Shengze Cai, Hui Li, and George Em Karniadakis. Nsfnets (navier-stokes flow nets): Physics-informed neural networks for the incompressible navier-stokes equations. J. Comput. Phys., 426:109951, 2021.
[8] Shengze Cai, Zhiping Mao, Zhicheng Wang, Minglang Yin, and George Em Karniadakis. Physics-informed neural networks (pinns) for fluid mechanics: A review. Acta Mech. Sin., 37(12):1727–1738, 2021.
[9] Maziar Raissi, Alireza Yazdani, and George Em Karniadakis. Hidden fluid mechanics: Learning velocity and pressure fields from flow visualizations. Science, 367(6481):1026–1030, 2020.
[10] Lorenz Richter, Leon Sallandt, and Nikolas Nüsken. Solving high-dimensional parabolic pdes using the tensor train format. In International Conference on Machine Learning, volume 139, pages 8998–9009. PMLR, 2021.
[11] Zhao Chen, Yang Liu, and Hao Sun. Physics-informed learning of governing equations from scarce data. Nat. Commun., 12(1):6136, 2021.
[12] Hao Xu, Junsheng Zeng, and Dongxiao Zhang. Discovery of partial differential equations from highly noisy and sparse data with physics-informed information criterion. Research, 6:0147, 2023.
[13] Aditi Krishnapriyan, Amir Gholami, Shandian Zhe, Robert Kirby, and Michael W Mahoney. Characterizing possible failure modes in physics-informed neural networks. In Advances in Neural Information Processing Systems, volume 34, pages 26548–26560, 2021.
[14] Tamara G Grossmann, Urszula Julia Komorowska, Jonas Latz, and Carola-Bibiane Schönlieb. Can physics-informed neural networks beat the finite element method? IMA J. Appl. Math., 89(1):hxae011, 2024.
[15] Kamyar Azizzadenesheli, Nikola Kovachki, Zongyi Li, Miguel Liu-Schiaffini, Jean Kossaifi, and Anima Anandkumar. Neural operators for accelerating scientific simulations and design. Nat. Rev. Phys., 6(5):320–328, 2024.
[16] Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nat. Mach. Intell., 3(3):218–229, 2021.
[17] Katiana Kontolati, Somdatta Goswami, George Em Karniadakis, and Michael D Shields. Learning nonlinear operators in latent spaces for real-time predictions of complex dynamics in physical systems. Nat. Commun., 15(1):5101, 2024.
[18] Alena Kopaničáková and George Em Karniadakis. Deeponet based preconditioning strategies for solving parametric linear systems of equations. SIAM J. Sci. Comput., 47(1):C151–C181, 2025.
[19] Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, pages 1–9, 2021.
[20] Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural operators. In International Conference on Learning Representations, pages 1–9, 2023.
[21] Yunpeng Wang, Zhijie Li, Zelong Yuan, Wenhui Peng, Tianyuan Liu, and Jianchun Wang. Prediction of turbulent channel flow using fourier neural operator-based machine-learning strategy. Phys. Rev. Fluids, 9(8):084604, 2024.
[22] Johannes Brandstetter, Daniel E. Worrall, and Max Welling. Message passing neural PDE solvers. In International Conference on Learning Representations, pages 1–9, 2022.
[23] Bocheng Zeng, Qi Wang, Mengtao Yan, Yang Liu, Ruizhi Chengze, Yi Zhang, Hongsheng Liu, Zidong Wang, and Hao Sun. PhyMPGN: Physics-encoded message passing graph network for spatiotemporal PDE systems. In International Conference on Learning Representations, pages 1–9, 2025.
[24] Zijie Li, Dule Shu, and Amir Barati Farimani. Scalable transformer for pde surrogate modeling. In Advances in Neural Information Processing Systems, volume 36, pages 28010–28039, 2023.
[25] Zhijie Li, Tianyuan Liu, Wenhui Peng, Zelong Yuan, and Jianchun Wang. A transformer-based neural operator for large-eddy simulation of turbulence. Phys. Fluids, 36(6), 2024.
[26] Somdatta Goswami, Katiana Kontolati, Michael D Shields, and George Em Karniadakis. Deep transfer operator learning for partial differential equations under conditional shift. Nat. Mach. Intell., 4(12):1155–1164, 2022.
[27] Michael McCabe, Bruno Régaldo-Saint Blancard, Liam Parker, Ruben Ohana, Miles Cranmer, Alberto Bietti, Michael Eickenberg, Siavash Golkar, Geraud Krawezik, Francois Lanusse, Mariel
Pettee, Tiberiu Tesileanu, Kyunghyun Cho, and Shirley Ho. Multiple physics pretraining for spatiotemporal surrogate models. In Advances in Neural Information Processing Systems, volume 37, pages 119301–119335, 2024.
[28] Tianyu Chen, Haoyi Zhou, Ying Li, Hao Wang, Chonghan Gao, Rongye Shi, Shanghang Zhang, and Jianxin Li. Building flexible machine learning models for scientific computing at scale. In International Conference on Machine Learning, 2025.
[29] Zhongkai Hao, Chang Su, Songming Liu, Julius Berner, Chengyang Ying, Hang Su, Anima Anandkumar, Jian Song, and Jun Zhu. Dpot: Auto-regressive denoising operator transformer for large-scale pde pre-training. In International Conference on Machine Learning, pages 17616– 17635. PMLR, 2024.
[30] Liu Yang, Siting Liu, Tingwei Meng, and Stanley J Osher. In-context operator learning with data prompts for differential equation problems. Proc. Natl. Acad. Sci., 120(39):e231014212, 2023.
[31] Yuchen Zhang, Mingsheng Long, Kaiyuan Chen, Lanxiang Xing, Ronghua Jin, Michael I Jordan, and Jianmin Wang. Skilful nowcasting of extreme precipitation with nowcastnet. Nature, 619(7970):526–532, 2023.
[32] Peiyan Hu, Rui Wang, Xiang Zheng, Tao Zhang, Haodong Feng, Ruiqi Feng, Long Wei, Yue Wang, Zhi-Ming Ma, and Tailin Wu. Wavelet diffusion neural operator. In International Conference on Learning Representations, pages 1–9, 2025.
[33] Salah A Faroughi, Nikhil M Pawar, Célio Fernandes, Maziar Raissi, Subasish Das, Nima K Kalantari, and Seyed Kourosh Mahjour. Physics-guided, physics-informed, and physics-encoded neural networks and operators in scientific computing: Fluid and solid mechanics. J. Comput. Inf. Sci. Eng., 24(4):040802, 2024.
[34] Sifan Wang, Hanwen Wang, and Paris Perdikaris. Learning the solution operator of parametric partial differential equations with physics-informed DeepONets. Sci. Adv., 7(40):eabi8605, 2021.
[35] Luis Mandl, Somdatta Goswami, Lena Lambers, and Tim Ricken. Separable physics-informed deeponet: Breaking the curse of dimensionality in physics-informed machine learning. Comput. Meth. Appl. Mech. Eng., 434:117586, 2025.
[36] Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar. Physics-informed neural operator for learning partial differential equations. ACM/JMS J. Data Sci., 1(3):1–27, 2024.
[37] Mohammad Sadegh Eshaghi, Cosmin Anitescu, Manish Thombre, Yizheng Wang, Xiaoying Zhuang, and Timon Rabczuk. Variational physics-informed neural operator (vino) for solving partial differential equations. Comput. Meth. Appl. Mech. Eng., 437:117785, 2025.
[38] Rui Zhang, Qi Meng, Rongchan Zhu, Yue Wang, Wenlei Shi, Shihua Zhang, Zhi-Ming Ma, and Tie-Yan Liu. Monte carlo neural pde solver for learning pdes via probabilistic representation. IEEE Trans. Pattern Anal. Mach. Intell., 2025.
[39] Zichao Long, Yiping Lu, and Bin Dong. Pde-net 2.0: Learning pdes from data with a numericsymbolic hybrid deep network. J. Comput. Phys., 399:108925, 2019.
[40] Chengping Rao, Pu Ren, Qi Wang, Oral Buyukozturk, Hao Sun, and Yang Liu. Encoding physics to learn reaction–diffusion processes. Nat. Mach. Intell., 5(7):765–779, 2023.
[41] Qi Wang, Pu Ren, Hao Zhou, Xin-Yang Liu, Zhiwen Deng, Yi Zhang, Ruizhi Chengze, Hongsheng Liu, Zidong Wang, Jian-Xun Wang, Ji-Rong Wen, Hao Sun, and Yang Liu. P2c2net: PDE-preserved coarse correction network for efficient prediction of spatiotemporal dynamics. In Annual Conference on Neural Information Processing Systems, 2024.
[42] Xin-Yang Liu, Min Zhu, Lu Lu, Hao Sun, and Jian-Xun Wang. Multi-resolution partial differential equations preserved learning framework for spatiotemporal dynamics. Commun. Phys., 7(1):31, 2024.
[43] Dmitrii Kochkov, Jamie A Smith, Ayya Alieva, Qing Wang, Michael P Brenner, and Stephan Hoyer. Machine learning–accelerated computational fluid dynamics. Proc. Natl. Acad. Sci., 118(21):e210178411, 2021.
[44] Zhiqing Sun, Yiming Yang, and Shinjae Yoo. A neural pde solver with temporal stencil modeling. In International Conference on Machine Learning, volume 202, pages 33135–33155. PMLR, 2023.
[45] Masanobu Horie and Naoto Mitsume. Graph neural pde solvers with conservation and similarity-equivariance. In International Conference on Machine Learning, volume 755, page 30. JMLR.org, 2024.
[46] Nicola Rares Franco, Andrea Manzoni, and Paolo Zunino. Mesh-informed neural networks for operator learning in finite element spaces. J. Sci. Comput., 97(2):35, 2023.
[47] Yuan Mi, Pu Ren, Hongteng Xu, Hongsheng Liu, Zidong Wang, Yike Guo, Ji-Rong Wen, Hao Sun, and Yang Liu. Conservation-informed graph learning for spatiotemporal dynamics prediction. In Conference on Knowledge Discovery and Data Mining, pages 1056–1067. ACM, 2025.
[48] Nick McGreivy and Ammar Hakim. Weak baselines and reporting biases lead to overoptimism in machine learning for fluid-related partial differential equations. Nat. Mach. Intell., 6(10):1256–1269, 2024.
[49] Gregory I Sivashinsky. On flame propagation under conditions of stoichiometry. SIAM J. Appl. Math., 39(1):67–82, 1980.
[50] Daniel M Michelson and Gregory I Sivashinsky. Nonlinear analysis of hydrodynamic instability in laminar flames—ii. numerical experiments. Acta Astronaut., 4(11-12):1207–1221, 1977.
[51] Giulio Guj and Fulvio Stella. A vorticity-velocity method for the numerical solution of 3d incompressible flows. J. Comput. Phys., 106(2):286–298, 1993.
[52] Geneviève Raugel and George R Sell. Navier-stokes equations on thin 3d domains. i. global attractors and global regularity of solutions. J. Am. Math. Soc., 6(3):503–568, 1993.
[53] Marten Lienen, David Lüdke, Jan Hansen-Palmus, and Stephan Günnemann. From zero to turbulence: Generative modeling for 3d flow simulation. In International Conference on Learning Representations, pages 1–9, 2024.
[54] Bogdan Raonic, Roberto Molinaro, Tim De Ryck, Tobias Rohner, Francesca Bartolucci, Rima Alaifari, Siddhartha Mishra, and Emmanuel de Bézenac. Convolutional neural operators for robust and accurate learning of pdes. In Advances in Neural Information Processing Systems, volume 36, pages 77187–77200, 2023.
[55] Simon L Cotter, Massoumeh Dashti, James Cooper Robinson, and Andrew M Stuart. Bayesian inverse problems for functions and applications to fluid mechanics. Inverse Probl., 25(11):115008, 2009.
[56] Rui Zhang, Peiyan Hu, Qi Meng, Yue Wang, Rongchan Zhu, Bingguang Chen, Zhi-Ming Ma, and Tie-Yan Liu. Drvn (deep random vortex network): A new physics-informed machine learning method for simulating and inferring incompressible fluid flows. Phy. Fluids, 34(10), 2022.
[57] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
[58] Jingrun Chen, Xurong Chi, Weinan E, and Zhouwang Yang. Bridging traditional and machine learning-based algorithms for solving pdes: The random feature method. J. Mach. Learn., 1(3):268–298, 2022.
[59] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531, 2015.
[60] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, highperformance deep learning library. In Advances in Neural Information Processing Systems, pages 8024–8035, 2019.
[61] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, pages 1–9, 2015.

Acknowledgement: The work is supported by the National Natural Science Foundation of China ( No. 62276269, No. 92270118), the Beijing Natural Science Foundation (No. 1232009), and the Strategic Priority Research Program of the Chinese Academy of Sciences (No. XDB0620103). In addition, Y.L. would like to acknowledge the support from the Fundamental Research Funds for the Central Universities (E2EG2202X2).

Author contributions: R.Z., Q.M. and H.S. contributed to the ideation and design of the research; R.Z., Q.M. and H.W. performed the research; H.S. supervised the project; all authors contributed to the research discussions, writing, and editing of the paper.

Correspondence to: Hao Sun (haosun@ruc.edu.cn).

Competing interests: The authors declare no competing interests.

Supplementary information: The supplementary information is attached.

Extended Data Figure 1: Flow statistics from PINO and OmniFluids across diverse CFD systems in the physics-only pre-training stage. a-b. Flow statistics comparison: kinetic energy spectrum Epkq and probability density function (PDF) of normalized vorticity ω (for 2D NSE) or velocity magnitude }u} (for 3D NSE).

Extended Data Figure 2: Streamtube visualizations of predicted 3D velocity magnitudes at t " 8 under ID (Re " 2000) and OOD (Re " 3000) settings. Note that only a corner of the flow field is shown for better visualization.

This supplementary document provides a detailed description of the data generation process, details on implementing OmniFluids and baseline methods, a discussion of ablation studies, and additional experimental results.

Supplementary Note A: Data generation

| Appendix Data Table 1: Computational parameters for datasets generation. | | | | |--------------------------------------------------------------------------|--|--|--| |--------------------------------------------------------------------------|--|--|--|

Dataset	Phase	Parameters	Nx	Nt	# Train/Val./Test
2D KSE	Pre-train	α P r0.1, 0.5s, β P r0.1, 0.5s	1282	5 ˆ 104	0 / 0 / 10
	Fine-tuning (ID)	α " 0.2, β " 0.5	1282	5 ˆ 104	2 / 2 / 10
	Fine-tuning (OOD)	α " 1.0, β " 1.0	1282	5 ˆ 104	2 / 2 / 10
2D NSE	Pre-train	Re P r500, 2500s, f P F	10242	1 ˆ 105	0 / 0 / 10
	Fine-tuning (ID)	Re " 2000, f P F	10242	1 ˆ 105	10 / 2 / 10
	Fine-tuning (OOD)	Re " 4000, f P F	10242	1 ˆ 105	10 / 2 / 10
3D NSE	Pre-train	Re P r500, 2500s	5123	1 ˆ 104	0 / 0 / 4
	Fine-tuning (ID)	Re " 2000	5123	1 ˆ 104	4 / 2 / 4
	Fine-tuning (OOD)	Re " 3000	5123	1 ˆ 104	4 / 2 / 4

To evaluate the performance of OmniFluids, we conduct experiments on a diverse set of PDEs that represent canonical problems in fluid dynamics, ranging from chaotic pattern formation to realistic incompressible flows. Specifically, we consider the 2D Kuramoto-Sivashinsky equations (KSEs), the 2D incompressible Navier-Stokes equations (NSEs) in velocity-vorticity form, and the 3D incompressible NSEs in velocity-pressure form. These systems are physically interrelated and collectively capture a broad spectrum of fluid phenomena–from the emergence of spatiotemporal chaos to 3D turbulence–thus offering a comprehensive evaluation framework for CFD modeling and simulation. All datasets are generated with high-resolution spectral solvers, and the detailed settings for each system are provided in Appedix Table 1.

2D Kuramoto-Sivashinsky equations. The 2D KSE is given by

$$\frac{\partial u}{\partial t} + \alpha(\Delta u + \Delta^2 u) + \frac{\beta}{2} |\nabla u|^2 = 0, \mathbf{x} \in [0, 8\pi)^2, t \in [0, 5],$$ (15)

2D incompressible vorticity-velocity NSE. The governing equation of 2D incompressible NSEs is

where ωpx, tq P R is the vorticity scalar field, u " ru, vs P R 2 is the incompressible velocity field satisfying ∇ ¨ u " 0, f and Re denotes the external forcing and Reynolds number, respectively. In our pre-training setup, we vary Re P r500, 2500s, ω⁰ P G2DNS and f P F to induce a range of vortex dynamics representative of real-world unsteady 2D flows, where G2DNS and F are 2D Gaussian random fields. To generate training data, we adapted the publicly available code in [36], accessible at https://github.com/neuraloperator/physics\_informed.

3D incompressible velocity-pressure NSE. The governing equations of 3D NSE is

where upx, tq " ru, v, ws P R 3 is the velocity field, ppx, tq P R is the pressure field, Re is the Reynolds number. The incompressibility constraint ∇ ¨ u " 0 enforces mass conservation. In our pre-training regime, we sample u⁰ P G3DNS, Re P r500, 2500s to generate a diverse ensemble of 3D unsteady flow regimes, where G3DNS is a 3D Gaussian random field. We compute the curl of a 3D Gaussian random field as the initial condition, thereby ensuring the divergence-free condition is satisfied. Data generation is based on the reference implementation provided at https://jipolanco.github.io/ PencilFFTs.jl/dev/generated/navier\_stokes/#.

Supplementary Note B: Baseline methods and implementation details

We herein introduce the baseline methods and implementation details.

B.1: Baseline methods

When transferring to specific tasks, for the 2D cases, we compare the performance of OmniFluids with various other approaches, including state-of-the-art neural operators such as PINO [36], DeepONet [16], FNO [19], FFNO [20], FactFormer [24], and CNO [54]; a large-scale data-driven pre-trained foundation model (DPOT [29]); and direct numerical simulation (DNS) based on spectral methods. For the 3D case, we use PINO [36], FNO [19], DPOT [29], and DNS as baselines for comparison. All models are trained and evaluated on NVIDIA A800 GPUs with PyTorch [60] as the deep learning framework. We fix the random seed to 0 for reproducibility.

PINO [36] is a physics-informed neural operator that incorporates PDE residuals into its loss function, enabling training without labeled data. It extends FNO with physics constraints and serves as a strong baseline for physics-aware operator learning. We use the official implementation at https://github.com/neuraloperator/physics\_informed.

DeepONet [16] is one of the earliest and most influential neural operator architectures. It uses a two-branch design: the branch net encodes input functions (e.g., initial conditions), while the trunk net encodes spatial coordinates, and their inner product yields the operator output. This structure enables strong function approximation capabilities across different input domains. We use the official implementation at https://github.com/lululxvi/deeponet.

FNO [19] is a pioneering neural operator framework based on fast Fourier transforms. It learns mappings between function spaces and is widely adopted for PDE modeling. We use the official codebase at https://github.com/neuraloperator/neuraloperator.

FFNO [20] builds on FNO by factorizing the spectral kernel into separable components, allowing for deeper architectures and reduced memory usage. We use the implementation from https: //github.com/alasdairtran/fourierflow.

FactFormer [24] is a transformer-based operator architecture with factored attention, designed for efficient and scalable modeling of spatiotemporal PDEs. It achieves top performance across various PDE tasks, and we adopt the code from https://github.com/BaratiLab/FactFormer.

CNO [54] is a convolution-based neural operator that combines resolution-invariance with locality and translation equivariance, offering a more robust architecture than Fourier-based models. Code is available at https://github.com/camlab-ethz/ConvolutionalNeuralOperator.

DPOT [29] is a large-scale data-driven pre-trained operator transformer designed as a foundation model for various physics. It is pre-trained on diverse physical systems and can be fine-tuned with minimal data for downstream tasks. We use the official codebase at https://github.com/ HaoZhongkai/DPOT.

DNS (Direct Numerical Simulation) uses spectral methods consistent with our data generation pipeline. To ensure a fair comparison with neural operator methods, we match the spatial resolution of DNS to that used in the AI-driven approaches. This setting allows us to directly compare spectral characteristics across methods. Additionally, we select the largest feasible time step that yields stable and physically meaningful rollouts.

B.2: Physics-only pre-training stage

In the physics-only pre-training stage, we include PINO [36] as a baseline to learn surrogate models for CFD systems without data supervision. PINO and OmniFluids are trained using the Adam [61] optimizer for 20,000 epochs with cosine annealing learning rate scheduling. The maximum learning rate is selected from t0.001, 0.002, 0.005u based on the lowest validation loss. For the 2D case, we set the batch size to 10 and use the maximum allowable number of spectral modes to ensure the model captures all frequency bands. For the 3D case, we use a batch size of 1. OmniFluids employs 64 modes per axis in 3D settings; however, for PINO, whose parameter complexity scales cubically with the number of modes, we use only 32 modes to maintain tractability. For the number of layers, OmniFluids uses 8 layers for 2D KSE, 12 for 2D NSE, and 6 for 3D NSE. In contrast, PINO uses a fixed 4-layer backbone, consistent with its official implementation, as increasing the depth typically degrades performance in its architecture [20]. To ensure a fair comparison, we tune the channel dimensions of each model such that the overall parameter count remains approximately equal across models.

B.3: Operator distillation

In the operator distillation stage, the architecture of the student model is kept nearly identical to that of the teacher model, except for a reduced number of spectral modes to match the lower spatial resolution of the input. Additionally, the final multi-frame decoder is replaced with a standard MLP, as we only require predictions of the physical field at the final time step. We train the student models for 20,000 epochs with cosine annealing learning rate scheduling. The maximum learning rate is selected as 0.002. The batch sizes are 10 for the 2D KSE and 2D NSE cases, and 2 for the 3D NSE cases.

B.4: Adaptation to new tasks

For OmniFluids, we adopt the distilled student model architecture for fine-tuning downstream tasks. We use fixed batch sizes of 2, 10, and 2 for the 2D KSE, 2D NSE, and 3D NSE settings, respectively. Each model is trained for 1,000 epochs with a fixed small learning rate selected from 0.00005, 0.0001, 0.0002, chosen based on validation performance.

We also use a small learning rate for adaptation for DPOT, which has been pre-trained on large-scale physics data. We employ the publicly released dpot-m model provided by the original authors, which is consistent with OmniFluids in terms of the size of parameters. The fine-tuning configuration, including batch size, number of epochs, and learning rate selection, is aligned with that used for OmniFluids.

For PINO, DeepONet, FNO, FFNO, FactFormer, and CNO, all models are trained from scratch for each target task. These baseline models are trained using the Adam optimizer for 20,000 epochs with cosine annealing learning rate scheduling. We use fixed batch sizes of 2, 10, and 2 for the 2D KSE, 2D NSE, and 3D NSE tasks, respectively. The maximum learning rate is selected from {0.001, 0.002, 0.005} based on the lowest validation loss. The key hyperparameters for PINO and FNO follow those used in the physics-only pre-training stage, with the number of Fourier modes adjusted according to the spatial resolution of each task. For DeepONet, the number of layers in branch and trunk net selected from {4, 6, 8} based on validation performance. For FFNO, we use the maximum number of spectral modes and set the number of layers to 8 for 2D KSE and 12 for 2D NSE, which aligns with OmniFluids. For FactFormer, the model depth is selected from {4, 6, 8} and the number of attention heads from {8, 16}, while other architectural settings follow the original paper. For CNO, the number of blocks is chosen from {4, 6, 8}.

Finally, we adjust the number of channels in each AI-driven method to ensure that all models have comparable parameter counts for a fair comparison.

B.5: Inverse analysis

For inverse problems, we adopt a fixed learning rate of 0.1 and optimize using the Adam optimizer for 200 epochs. The initial viscosity coefficient is set to 0.001, and the external forcing field is initialized by randomly sampling from the forcing space F.

Supplementary Note C: Ablation studies

In this section, we present a series of ablation studies to analyze the contributions of key components in OmniFluids and investigate the effects of varying the number of training data.

C.1: Ablation study on key components of OmniFluids

Appendix Data Table 2: Ablation study on key components of OmniFluids across 2D KSE task. We report the relative ℓ² error both pre-training and fine-tuning phases. Each entry shows A/B, where A is the average relative error over time, and B is the error at the final step. The best results are in bold, and the second-best results are underlined.

Model	Pre-train	Fine-tuning (ID)	Fine-tuning (OOD)
OmniFluids\Decoder	0.2572 / 0.4242	0.0318 / 0.0446	0.0624 / 0.0965
OmniFluids\MoO	0.3574 / 0.6018	0.0319 / 0.0450	0.0643 / 0.1015
OmniFluids \Multi	– / –	0.0219 / 0.0296	0.1362 / 0.1525
OmniFluids\Distillation	0.9296 / 1.5803	0.0822 / 0.0694	0.2307 / 0.2332
OmniFluids	0.1777 / 0.2831	0.0208 / 0.0276	0.0400 / 0.0590

We perform various ablation studies on the 2D KSE task to examine the impact of key components in the OmniFluids framework. Specifically, OmniFluids\Decoder replaces the final multi-frame decoder with a standard MLP; OmniFluids\MoO replaces the mixture of operators (MoO) module with standard factorized Fourier layers; OmniFluids\Multi turns off multi-task pre-training by training on a single physics configuration only, and then transfers the ID model to the OOD setting; and OmniFluids\ Distillation removes the distillation step and directly pre-trains on low-resolution spatiotemporal grids.

As shown in Appedix Table 2, all components contribute to performance improvement, particularly during the physics-only pre-training phase. Among them, operator distillation plays the most critical role, as it is infeasible to construct accurate physics-based loss functions directly on coarse spatiotemporal grids. Multi-task pre-training also significantly improves generalization, especially in OOD scenarios, as training among varied physics offers valuable inductive biases. Furthermore, both the multi-frame decoder and MoO design yield consistent improvements by enhancing expressivity and adaptability. Overall, these results demonstrate that each component is essential, enabling strong pre-trained operators that are highly transferable for downstream tasks.

C.2: The effects of varying the number of training data

Appendix Figure 1: Data-efficiency: performance of OmniFluids with varying number of training trajectories on 2D NSEs. a. ID scenario with Re " 2000. b. OOD scenario with Re " 4000.

Appedix Fig. 1 demonstrate OmniFluids' data efficiency on 2D NSEs under ID (Re " 2000) and OOD (Re " 4000) settings. The model shows robustness to fine-tuning data quantity, achieving reasonable relative ℓ² errors even in the zero-shot case. While additional training trajectories yield improved accuracy, gains saturate quickly in the ID scenario, indicating the effects of strong prior knowledge from physics-only pre-training. On the other hand, the OOD scenario exhibits a bit more sensitivity to data volume, as higher Reynolds numbers correspond to more complex flow dynamics.

C.3: Additional results

We provide extended qualitative and quantitative results to further demonstrate the robustness, accuracy, and generalization capability of OmniFluids across both 2D and 3D fluid dynamics benchmarks in Appedix Fig. 2, Appedix Fig. 3, and Appedix Fig. 4.

Appendix Figure 2: Additional predicted snapshots of baseline methods and OmniFluids on 2D KSE. a. In-distribution snapshots (t " 5) with parameters α " 0.2, β " 0.5. b. Out-of-distribution snapshots (t " 5) with parameters α " 1.0, β " 1.0.

Appendix Figure 3: Additional predicted snapshots and flow statistics of OmniFluids and baseline methods on 2D NSE. a and b. Predicted velocity (u " ru, vs) and vorticity fields (ω) at t " 10 under ID (Re " 2000) and OOD (Re " 4000) settings, respectively. c and d. Kinetic energy spectrum and normalized vorticity PDFs for ID and OOD scenarios.

Appendix Figure 4: Additional predicted snapshots and flow statistics of OmniFluids and baseline methods on 3D NSE. a and b. Predicted velocity fields (u " ru, v, ws) at t " 8 under ID (Re " 2000) and OOD (Re " 3000) settings, respectively. Note that only a corner of the flow field is shown for better visualization. c and d. Kinetic energy spectrum for ID and OOD scenarios.

zhangomnifluidsunifiedphysicspre2025.md 96 KB Riwayat Mentahan