1 Introduction↩︎

Modeling turbulent convection is challenging and has applications from atmospheric flows [1] to industrial processes such as silicon wafer production [2]. A standard benchmark is Rayleigh–Bénard convection (RBC), where a fluid heated from below develops an overturning circulation [3] with the strength of turbulence governed by the Rayleigh number (\(Ra\)). At high Ra the flow becomes strongly turbulent, making high-resolution numerical simulations computationally expensive. Recent machine learning approaches, particularly Fourier Neural Operators (FNOs) [4], provide an alternative to mesh-based solvers by learning solution operators that, unlike PINNs [5], generalize across resolutions. In this work, we apply FNOs to 2D RBC using Dedalus [6] as ground truth, focusing on stable small-step predictions that preserve turbulent statistics and align with time-stepping solvers. We propose a lean FNO design suitable for use in iterative methods such as Parareal [7], [8].

1.0.0.1 Related Work.

Hammoud et al. [9] show that PINNs reconstruct previous RBC states with relative \(L_2\)-errors on the order of \(10^{-1}\) at \(\text{Ra}=10^7\), while Almeida et al. [10] achieve less than 10 % reconstruction error using an encoder–decoder DeepONet-based approach [11]. Wu et al. [12] introduce COAST, a causal operator framework with adaptive time stepping, reporting a VRMSE of 0.23 for \(Ra=10^7\) over a 5-step rollout. Straat et al. [13] apply FNO to 2D RBC at multiple Rayleigh numbers and obtain an average relative \(L_2\)-error of \(0.021\) on a \(96 \times 64\) grid. In contrast, our 2D lean FNO model learns to predict increments rather than solutions and achieves significantly higher accuracy when predicting future solution states. It has fewer parameters and faster inference, which is beneficial when coupled with iterative numerical algorithms. However, it exhibits faster error accumulation under longer auto-regressive rollouts or larger time steps compared to 3D FNOs.

1.0.0.2 Contributions

We propose a lean FNO architecture that learns the increment from the state at time \(t\) to a state at time \(t + \Delta t\), similar to a time-stepping algorithm in a mesh-based simulation. We perform an ablation study which reveals that using multi-layer 1D convolutional scaling operators delivers better accuracy than linear layers. The performance of our lean FNO architecture in terms of accuracy, memory footprint and inference times is compared against the model by [13]. Finally, we analyze how FNO generalizes across spatial and temporal resolutions and show that it interpolates to new meshes, but that accuracy remains limited by the resolution of the training data.

2 Rayleigh-Bénard Convection↩︎

Rayleigh-Bénard convection occurs when a fluid within a confined channel is subjected to a temperature difference. In our setup, the lower plate is heated to a temperature \(\theta_{H}\) while the upper plate is cooled to \(\theta_{C}\) with a vertical distance (d) between the plates. This causes the cooler, denser fluid to sink and the warmer, lighter fluid to rise, generating convective rolls whose patterns are influenced by the aspect ratio of the channel and the physical properties of the fluid. The non-dimensional momentum conservation equations governing the dynamics of RBC are \[\label{eq:u95t} \frac{\partial \mathbf{u}}{\partial t} + (\mathbf{u} \cdot \nabla) \mathbf{u} = -\nabla p + \theta \mathbf{\hat{z}} + \sqrt{\frac{\text{Pr}}{\text{Ra}}} \nabla^2 \mathbf{u}.\tag{1}\] Here, \(\mathbf{u} = (u, w)\) denotes the velocity vector with horizontal component \(u\) and vertical component \(w\), \(p\) is the pressure, \(\theta\) is buoyancy, and the Prandtl number \(\text{Pr}\) is the ratio of kinematic viscosity \(\mu\) and thermal diffusivity \(\kappa\). The Rayleigh number \(\text{Ra}\), which characterizes the convective motion, is defined as \[\text{Ra} = \frac{\alpha \text{g} (\theta_{H}-\theta_{C})d^3}{\mu \kappa}\] with \(\alpha\) being the thermal expansion coefficient and \(\text{g}\) the acceleration due to gravity. The continuity equation for incompressible flow that ensures mass conservation is \[\label{eq:divu} \nabla \cdot \mathbf{u} = 0\tag{2}\] Conservation of energy provides the following equation for the temperature \[\label{eq:theta95t} \frac{\partial \theta}{\partial t} + (\mathbf{u} \cdot \nabla) \theta = \frac{1}{\sqrt{\text{Ra Pr}}}\nabla^2 \theta.\tag{3}\]

3 Fourier Neural Operator↩︎

Fourier Neural Operators (FNOs) solve PDEs by learning mappings between function spaces using spectral representations [4]. The architecture consists of an input lifting layer \(\mathcal{P}\) that maps the input function \(a\) from channel dimension \(d_a\) to a higher-dimensional representation \(v\) of dimension \(d_v\), followed by multiple Fourier layers and a final projection \(\mathcal{Q}\) that maps to the output \(u\) with channel dimension \(d_u\). Each Fourier layer applies a Fourier transform \(\mathcal{F}\), performs spectral convolution on a fixed set of low-frequency modes \(R\), and transforms back via the inverse transform \(\mathcal{F}^{-1}\). A local linear term \(Wv\) and a nonlinear activation \(\sigma\) are added to capture local and higher-frequency features beyond the truncated spectrum.

3.1 Generation of training data↩︎

Training and evaluation data were generated using Dedalus [6] on a \(256 \times 64\) grid with a \(RK443\) time stepper and \(\Delta t = 10^{-3}\), with \(\theta_H=1\) and \(\theta_C=0\). Simulations start from random perturbations, run to \(T_{init}=100\) to reach a pseudo-steady turbulent state. Data is collected from \(T_{init}=100\) to \(T_{sim}=200\) for \(\text{Pr}=1\) and \(\text{Ra}=10^7\), with spectral analysis confirming adequate resolution. Ten simulations with different random seeds are used to record \(\mathcal{U}(t)=[\mathbf{u}(t),\theta(t),p(t)]\) at each time step, forming 200 input–output pairs \((\mathcal{U}(t),\mathcal{U}(t+\Delta t))\) over \([100,200]\). From 2000 samples, 80 % are used for training and 20 % for validation.

3.2 Learning Objective↩︎

FNOs are typically trained to predict the full solution of a PDE, but learning updates to the state can be more effective. We therefore compare two objectives: a solution objective, where the model directly predicts \(\mathcal{U}(t+\Delta t)\) from \(\mathcal{U}(t)\), and an increment objective, where it predicts the scaled update \(\Delta \mathcal{U} = \Delta t^{-1}(\mathcal{U}(t+\Delta t)-\mathcal{U}(t))\). In the latter case, the solution is recovered via \(\mathcal{U}(t+\Delta t)=\mathcal{U}(t)+\Delta t\,\mathcal{O}(\mathcal{U}(t))\), where \(\mathcal{O}(\mathcal{U}(t))\) is the model increment prediction, making the FNO analogous to a one-step time integrator.

Both objectives are trained using a relative \(L_2\) loss, applied either to the solution or the increment. For evaluation, we measure the relative error in the reconstructed solution. Note that loss and error coincide only for the solution objective. Since small \(\Delta t\) leads to nearly identical states, trivial predictions can yield low loss; therefore, we compare against an identity baseline that propagates \(\mathcal{U}(t)\) unchanged (or predicts zero increment) to ensure meaningful learning beyond this naive predictor.

4 Results↩︎

Training was conducted on a single NVIDIA A100 (40 GB) GPU of the JURECA-DC supercomputer, using a custom codebase [14] and data¹. Our ablation study on \(\mathcal{P}\) and \(\mathcal{Q}\) (linear vs. 1D convolutional layers with \(d_v=16\)) shows that increasing both depth and width improves performance and using cosine LR scheduler leads to better learning. The improved model is further evaluated to test the temporal and spatial generalization, and extended to compare against existing work.

4.1 Training↩︎

The baseline model is trained using both the solution and increment objectives for \(\Delta t = 10^{-3}\). 1 compares the relative error for both objectives across the four components of the solution, that is velocity in the horizontal (\(u\)) and vertical (\(w\)) directions, buoyancy(\(\theta\)), and pressure (\(p\)), averaged over 200 samples. Furthermore, the training loss is lower than the identity loss for increment objective but not for the solution objective.

The average relative error for the increment objective is two orders of magnitude smaller than for the solution objective. The noise introduced by the FNO when predicting the solution leads to an increase in error by about a factor of ten compared to when the FNO predicts the increment for buoyancy contours. Predicting increments substantially improves the quality of the generated solutions.

4.2 Improved model↩︎

Based on the results above, we propose the model configuration described in 2 as improvement over the baseline FNO. We train both for 11500 epochs to predict increments \(\Delta t = 10^{-3}\) using data generated as mentioned in 3.1 which takes 11h on a single NVIDIA A100 GPU. 2 also shows the relative error in the four variables in 1 and 3 against the Dedalus reference on \(256 \times 64\) grid for 200 samples from the validation dataset when predicting a single increment \(\Delta t = 10^{-3}\).

4.3 Impact of model time step and autoregressive FNO application↩︎

To evaluate the impact of different time steps, the model is retrained for \(\Delta t = 1\), \(10^{-1}\), \(10^{-2}\), in addition to \(10^{-3}\) using the same architecture (2) on 2000 samples for 11500 epochs. Models are evaluated over time horizons of \(1\), \(0.1\), \(0.01\), and \(0.001\) starting from \(t=100\), using autoregressive rollout when the model time step is smaller than the target horizon. As shown in 3, shorter horizons yield lower validation errors due to reduced solution change, while autoregressive inference introduces mild error accumulation, leading to slightly higher errors than single-step predictions at matching time steps.

4.5pt

4.4 Resolution invariance↩︎

To evaluate generalization across mesh resolutions, the FNO trained on \(\Delta t = 10^{-3}\) increments is tested on a \(512 \times 128\) grid, twice the training resolution of \(256 \times 64\) used for \(Ra = 10^7\). Errors are computed against reference Dedalus solutions and are found to be of the same order as those on the \(256 \times 64\) grid (4). Errors change very little if the reference solution on a \(512 \times 128\) is restricted to a \(256 \times 64\) mesh and this data used as input to the FNO. Evaluating the FNO on the trained resolution (\(256 \times 64\)) and interpolating the output to the finer grid via FFT delivers very similar errors as evaluating the FNO directly on the finer mesh. This demonstrates that the FNO shows mesh invariance by interpolating to unseen resolutions without significant additional error. However, unlike numerical solvers, increasing inference resolution does not improve accuracy, which is instead determined by the resolution of the training data.

4.5 Comparison with Straat et al. ↩︎

To evaluate the performance of our improved FNO architecture, we compare it against the model of Straat et al. [13] for \(Ra = 5 \times 10^6\) with \(Pr = 0.7\) on a \(96 \times 64\) grid and boundary temperatures \(\theta_H = 2\) and \(\theta_C = 1\). We adopt their setup but use an architecture with 8 Fourier layers, 64 projection channels, and 16 Fourier modes to better capture the turbulent dynamics and larger time step. The model is trained on 44940 samples from 15 Dedalus simulations and validated on 14980 samples from 5 simulations, each run for 300 s with data sampled at \(\Delta t = 0.1\) after a 200 s warm-up, and trained on JURECA-DC using 8 NVIDIA A100 GPUs for 500 epochs in 3.5 hours, reaching a training and validation loss of 0.03.

For evaluation, predictions are generated over a 30 s window via 60 recursive steps with a model time step of 0.5 s. Averaged over 50 samples from 10 random starting points across 5 validation runs, our model achieves an average loss of \(0.11\) (Table 5), compared to \(0.04\) reported by Straat et al. Direct component comparison is not possible due to lack of data in Straat et al. While their FNO3D accumulates error more slowly, our model attains comparable initial accuracy despite its much smaller architecture, requiring only 132 MB (33 million, FP32 parameters) versus 3037 MB, and achieves faster inference times of 10 ms (batch size 10) and 30 ms (batch size 50) per prediction window, compared to 0.45 s on an NVIDIA A40 for Straat et al.

5 Conclusion↩︎

In this work, we systematically evaluate Fourier Neural Operators for two-dimensional Rayleigh–Bénard convection in the turbulent regime. We find that predicting increments rather than full states significantly improves accuracy and stability, reducing errors by up to two orders of magnitude for small time steps and mitigating noise, effectively turning the FNO into a data-driven one-step integrator. Architectural choices such as deeper 1D convolutional lifting/projection layers and a cosine learning rate scheduler further improve accuracy, yielding a compact and efficient model.

Our results show that FNOs are resolution-invariant in the sense that they interpolate to unseen meshes, but their accuracy is ultimately limited by the training data resolution rather than improved by finer inference grids. Compared to the much larger 3D FNO by Straat et al, our lean model achieves a trade-off between accuracy, memory, and inference speed, with substantially smaller footprint and faster runtime but stronger long-horizon error accumulation. These properties highlight the potential of FNOs as efficient surrogate models and as coarse propagators in space–time parallel methods such as Parareal.

References↩︎