PEST: Physics-Enhanced Swin Transformer for 3D Turbulence Simulation

Accurate simulation of turbulent flows is fundamental to scientific and engineering applications. Direct numerical simulation (DNS) offers the highest fidelity but is computationally prohibitive, while existing data-driven alternatives struggle with stable long-horizon rollouts, physical consistency, and faithful simulation of small-scale structures. These challenges are particularly acute in three-dimensional (3D) settings, where the cubic growth of spatial degrees of freedom dramatically amplifies computational cost, memory demand, and the difficulty of capturing multi-scale interactions. To address these challenges, we propose a Physics-Enhanced Swin Transformer (PEST) for 3D turbulence simulation. PEST leverages a window-based self-attention mechanism to effectively model localized PDE interactions while maintaining computational efficiency. We introduce a frequency-domain adaptive loss that explicitly emphasizes small-scale structures, enabling more faithful simulation of high-frequency dynamics. To improve physical consistency, we incorporate Navier–Stokes residual constraints and divergence-free regularization directly into the learning objective. Extensive experiments on two representative turbulent flow configurations demonstrate that PEST achieves accurate, physically consistent, and stable autoregressive long-term simulations, outperforming existing data-driven baselines.

💡 Research Summary

The paper introduces PEST (Physics‑Enhanced Swin Transformer), a novel data‑driven framework for high‑resolution three‑dimensional turbulence simulation. While direct numerical simulation (DNS) offers unparalleled fidelity, its computational cost grows dramatically with Reynolds number, making it infeasible for many engineering problems. Existing machine‑learning approaches either sacrifice long‑term stability, ignore physical constraints, or fail to capture small‑scale structures that are crucial for the turbulent energy cascade. PEST addresses these three challenges simultaneously.

Architecture. PEST adopts a 3D Swin Transformer as its backbone. The input sequence of five consecutive flow states is first embedded by a 3D convolutional patch projector, then processed through a hierarchy of window‑based multi‑head self‑attention (W‑MSA) blocks. Regular and shifted windows alternate, ensuring that information can cross window boundaries while preserving the locality inherent to differential operators in PDEs. This design reduces the quadratic attention cost to linear‑ish O(N·M³) where N is the number of tokens and M is the window size, making it tractable for large 3D grids. A temporal attention module aggregates information across the input frames, and a set of learnable query tokens produces the next five time steps via cross‑attention and temporal self‑attention. A final transposed 3D convolution restores the original spatial resolution. To further suppress discontinuities at window borders, a gradient‑matching L1 loss (L_grad) is added, encouraging smooth gradients across the whole domain.

Frequency‑adaptive spectral loss. Recognizing that the standard L2 loss heavily weights the high‑energy low‑frequency components, the authors invoke Parseval’s theorem to rewrite the spatial L2 error as an exact sum of per‑wavenumber errors in Fourier space. This equivalence permits selective re‑weighting of frequency bands without breaking compatibility with physical‑space operators needed for the physics‑based constraints. A curriculum‑guided weighting schedule starts with larger emphasis on low frequencies and gradually increases the weight on high frequencies as training progresses, thereby forcing the network to learn the sub‑dominant small‑scale vortical structures that are essential for accurate dissipation and long‑term stability.

Physics‑informed regularization. Two physics‑based penalties are incorporated directly into the loss: (1) the Navier‑Stokes residual computed by automatic differentiation of the predicted velocity and pressure fields, and (2) a divergence‑free term enforcing ∇·u=0 for incompressible flow. These terms anchor the learned dynamics to the governing equations, reducing drift from physical laws during autoregressive rollout.

Uncertainty‑based multi‑loss balancing. Because data‑driven and physics‑driven losses differ in scale and learning dynamics, the paper adopts an uncertainty‑based weighting scheme. At each iteration, the variance of each loss component is estimated, and the inverse of this variance is used to adaptively scale the corresponding weight. This dynamic balancing prevents any single objective from dominating optimization and mitigates mode collapse.

Experiments. The method is evaluated on two benchmark turbulent flows with distinct characteristics: a transitional shear flow and a rotating turbulence case. PEST is compared against nine state‑of‑the‑art baselines, including Fourier Neural Operators, TF‑Net, OFormer, Transolver, and various physics‑informed networks. Metrics cover pointwise mean‑squared error, spectral energy spectra, kinetic‑energy decay, and divergence error. Across all metrics, PEST achieves lower errors and more faithful spectral shapes, especially in the high‑wavenumber range. Autoregressive rollouts over hundreds of time steps remain stable, whereas many baselines exhibit error blow‑up or loss of small‑scale features. Visualizations confirm that PEST preserves both large‑scale coherent structures and fine‑scale vortices, with minimal artifacts at window boundaries. Ablation studies demonstrate the contribution of each component: the spectral loss improves high‑frequency fidelity, the physics residuals reduce divergence and enforce energy consistency, and the uncertainty weighting stabilizes training.

Conclusion and outlook. PEST successfully unifies efficient local attention, frequency‑aware loss design, and explicit physics constraints to deliver accurate, physically consistent, and long‑term stable 3D turbulence simulations. The work opens avenues for extending the framework to more complex multiphysics problems (e.g., reacting flows, heat transfer) and for exploring model compression techniques to enable real‑time or on‑device turbulent flow prediction.