Fourier neural operators for spatiotemporal dynamics in two-dimensional turbulence

Machine learning (ML) has improved efficiency of forecasting in weather and climate models ([1, 2]). However, long time forecasts of ML-based emulators have been found to become unphysical or unstable (see [3] and references therein). Recently, it was shown in Ref. [4]) that the reason for this failure of ML-based emulators is spectral bias, where the smaller scales are not learned and only the large-scale dynamics are captured. In this paper, we propose a hybrid emulator for isotropic two-dimensional turbulence, which combines ML predictions with the governing partial differential equations (PDE) to march forward in time. This approach when generalized can be applied for long term predictions in climate models which are notoriously complex to solve due to the high computational cost of numerical PDE solvers. The key idea is to trade speed of the pure ML-emulator with the accuracy of the hybrid framework, i.e., the hybrid framework is not as fast as the pure ML-emulator but alternating between ML-emulator and PDE solver ensures that the results of the hybrid framework remain physical.

While the idea of using neural networks for flow field reconstruction has existed for over two decades [5], the application of machine learning in fluid dynamics witnessed a massive surge a decade ago [6]. Several works have since explored data-driven turbulence models, flow field reconstruction near walls, flow pattern recognition, machine learning accelerated simulations, and super-resolution (see [7, 8, 9, 10, 11, 12, 13] and references therein). Although, there exist several ML-based approaches for steady-state Reynolds-averaged Navier-Stokes simulations and the interpolation of flow fields in space and time, extrapolation problems have remained hard to tackle. This is particularly due to the nonlinear, multiscale, and chaotic nature of turbulence which remains a computational challenge for several important real-world applications of Navier-Stokes equations. The large memory requirement and scarcity of data have so far prohibited learning a generalized solution operator of the Navier-Stokes equations. Nevertheless, neural operators and embedding physical constraints in machine learning have shown promise in accelerating computational fluid dynamics.

In this paper, we evaluate the merits and discuss the pitfalls of a certain ML-based model for predicting spatiotemporal fluid dynamics. We choose to study decaying turbulence in two dimensions (2D) as it is an important problem and can be extended to forced turbulence or three dimensions. In Sec. II we briefly introduce the ML model. We adopt a purely data-driven approach and couple it with the numerical PDE solver to contain the growth of errors. It is unfair to expect a purely data-driven approach trained on a limited amount of data to accurately predict long term dynamics of a chaotic system. Thus, in Sec. III we discuss the data set and in Sec. IV assess the time over which a data-driven solver can be expected to make accurate predictions by calculating the Lyapunov exponents and computing evolution statistics of global quantities. The key question is how much data (samples and time refinement) is needed for an ML model to become a viable alternate to numerical solvers. Purely data-driven approaches are unlikely to replace the physical law-based numerical solvers [14, 15] as they lack the knowledge to model all relevant physics of a complex multiscale flow. Nevertheless, one needs to discuss its advantages and limitations by putting it on the same footing as numerical PDE solvers. To that end, in Sec. VI we identify an ML model, quantify the time refinement and number of samples (with different initial flow fields) required to train an ML model, and couple the ML model with PDE solver to make accurate long time predictions of the flow field at the physical conditions and Reynolds number relevant to atmospheric physics [16].

A neural operator is a neural network architecture that is designed to approximate a solution operator of resolution-independent PDEs. They learn mappings between infinite dimensional function spaces using a finite set of input-output data [17, 18, 19, 20, 21]. For example, the Fourier neural operator (FNO) operates in the frequency domain by learning mappings between Fourier coefficients of high-dimensional data [18, 20]. The deep network operator (DeepONet) [19] utilizes two subnetworks to encode solution operators to several deterministic and stochastic differential equations. The Laplace neural operator has been shown to work well with non-periodic functions [22]. The Markov neural operators are used to learn chaotic dynamics of non-ergodic dissipative systems [23]. The Clifford neural operators employ correlations between fields to learn multivector fields [24]. The Riemann operator network learns solution operators to hyperbolic PDEs with discontinuous solutions [25]. Additionally, there are several variations of each neural operator with their own special characteristics.

In this work, we select the Fourier neural operator (FNO) as it is a promising tool for learning discretization-agnostic approximations to the solution operator of differential equations. FNO truncates the high-frequency modes that are considered to be less important. Informally, if $\mathcal{A}$ and $\mathcal{U}$ are two Banach spaces of functions defined, respectively, on bounded domains $\Omega_{d_{a}}\subset\mathbb{R}^{d_{a}}$ and $\Omega_{d_{u}}\subset\mathbb{R}^{d_{u}}$, and $\mathcal{G}$ is an operator that maps functions $a\in\mathcal{A}$ to functions $u\in\mathcal{U}$, i.e., $\mathcal{G}:\mathcal{A}\rightarrow\mathcal{U}$, a FNO parameterizes the operator $\mathcal{G}$ by $\mathcal{G}_{\theta}:\mathcal{A}\times\theta\rightarrow\mathcal{U}$, such that for some $\theta^{*}\in\mathbb{R}^{p}$, we have the approximation $\mathcal{G}_{\theta^{*}}\ \approx\ \mathcal{G}$. Thus $\mathcal{G}_{\theta^{*}}$ acts as a surrogate to the operator $\mathcal{G}$. In the case of PDEs, the space $\mathcal{A}$ comprises “input” functions that the PDEs depend on, such as initial or boundary conditions, while the operator $\mathcal{G}$ maps these inputs to the solution $u\in\mathcal{U}$ satisfying the PDEs with the given inputs. In the present study, the input space $\mathcal{A}$ comprises initial conditions.

The data set consists of vorticity and velocity fields from 5000 simulations of 2D decaying turbulence. These simulations differ from one another in that they are initialized with different uniformly distributed random numbers but the same transport coefficient. The initial condition generates several opposite vortices that diffuse as time progresses and convect around the periodic domain depending on their initial locations. The Reynolds number is in the range of 7000-8000 for different samples depending on the initial condition. We solve the Navier-Stokes equations for the velocity field using the entropic lattice Boltzmann method [26, 27] on a grid of $256\times 256$ points. The lattice Boltzmann method is an alternate methodology for computational fluid dynamics based on discrete space-time kinetic theory (see [28] for details). Each initial condition consumes 263 seconds on a single core of AMD EPYC 7402 CPU. The flow field is first allowed to evolve for $0.5\,t_{c}$ (where $t_{c}=L/U_{0}$ is the convection time and $U_{0}$ is the characteristic velocity) so that the initial sharp discontinuities vanish. Thereafter, time is reset to zero and the velocity (${\bf u}$) and vorticity ($\omega_{z}$) are sampled from time $t=0$ to $t=t_{c}$ in steps of $0.005\,t_{c}$. Note that the vorticity field $\omega_{z}(x,y)$ is calculated as the curl of the velocity using $\omega_{z}(x,y)=\nabla\times{\bf u}(x,y)$, where $\nabla$ is the nabla operator and the velocity vector ${\bf u}$ has components $u_{1},u_{2}$ for two-dimensional fluid flow. Figure 8 visualizes the two dimensional vorticity field of one such sample where several inversely rotating vortices are visible. In general, in an incompressible viscous flow the vortices stretch (absent for two dimensions) and turn under the influence of each other’s field and diffuse as the time advances.

In general, turbulent dynamics are chaotic. Chaotic systems are characterized by sensitivity to initial conditions, i.e., two nearby initial conditions separate exponentially even when the underlying dynamical equations are deterministic ([29], [30]). For a given chaotic system, the Lyapunov time gives a time scale over which dynamics are predictable. Data-driven methods need to be mindful of this time scale to ensure accuracy and confidence when extrapolating in time. In this section, we first analyze the data set of 2D turbulence to understand the spatiotemporal characteristic of the flow field by computing the evolutions of the correlation coefficient and relative $L_{2}$ errors. The analysis ensures that there has been a meaningful evolution or separation from initial condition so that the prediction holds value, as, one pitfall is to make extremely short time predictions when the fields have evolved by such a tiny amount that even the initial condition would be an acceptable prediction. This is followed by estimating the Lyapunov time scale.

The evolution of volume mean of vorticity, standard deviation of vorticity, and global enstrophy (defined as the sum of square of vorticity fluctuation over the domain) with time for a subset of the full dataset has been shown in Fig. 1. The left column describes the behaviour of a data set as is, while the right column consists of the normalized datasets where the vorticity at each point in a sample has been scaled with the vorticity field’s mean and standard deviation at the initial instant. The mean remains constant at zero due to incompressibility (thus vorticity fluctuation is the same as vorticity), while the standard deviation decays with time as is expected of decaying 2D turbulence. The global enstrophy computed from normalized vorticity also decays as small scale structures are dissipated due to viscosity. In order to emphasize the turbulent nature of the flow and determine time scales associated with chaotic dynamics, we plot in Figs. 2 and 3 two measures of time evolution for ten samples from the full dataset. Fig. 2 shows the separation between the vorticity fields for each dataset at time $t$ from its initial value at time $t=0$, scaled with their respective initial values. In Fig. 3, we plot the correlation coefficient of the vorticity field after time $t$ and the initial vorticity field for the data sets scaled with their respective initial values. As expected, this correlation coefficient between the vorticity fields decays with increasing time.

For a dynamical system $\dot{\bm{x}}=\bm{f(x)}$, in practice, we select two nearby initial conditions with initial separation denoted by $\delta x_{0}=||\delta\bm{x}(t=0)||$, with a suitable norm chosen as measure of distance in the state space. As these trajectories evolve in time, we track their separation $\delta x(t)=||\delta\bm{x}(t)||$ as a function of time. Sensitivity to initial conditions dictates that nearby initial conditions separate exponentially and can be expressed mathematically as $\delta x(t)\approx(\delta x_{0})e^{\lambda t}$ . In this work, we define Lyapunov time as $T_{L}=1/\Lambda$, where $\Lambda$ is the maximum Lyapunov exponent defined as $\Lambda=\textrm{max}<\lambda>$, where

with $\lambda=({1}/{t})\ln{({\delta x(t)}/{\delta x_{0}})}$. To calculate $\Lambda$, we started with two initial conditions $A$ and $B$ such that $\delta x_{0}=||u^{(A)}_{1}(t=0)-u^{(B)}_{1}(t=0)||=10^{-2}$, with $||\cdot||$ being the $L_{2}$ norm. We then trace the evolution of $u^{(A)}_{1},u^{(A)}_{2},u^{(B)}_{1},u^{(B)}_{2}$ and calculate the quantity $\lambda_{i}$ at each time-step $t_{i}$ for $u_{1}$ and $u_{2}$. Since the chaotic system has a finite maximum extent, the quantity $\lambda_{i}$ grows as trajectories separate exponentially under chaotic dynamics until the separation between them becomes the size of the chaotic attractor itself in the state space. The larger of the two $\Lambda$’s is $\approx 2.15$ while their average is $\approx 1.7$. Thus, a conservative estimate of the Lyapunov time is $T_{L}=\Lambda^{-1}\approx 0.45$ convective time–units. This is consistent with Fig. 3 where vorticity correlation coefficients are seen to decay until $T_{L}$ after which the trajectories become independent. We will therefore restrict the predictions from FNO to a time horizon smaller than $T_{L}$.

In this section, we briefly introduce the methods for spatiotemporal dynamics of decaying turbulence. The goal is to find the best methodology for extrapolating the flow field in time from a few inherently chronological snapshots. The three methods that we consider in this paper are:

• •

  2D FNO with temporal channels: Here, the two dimensions of FNO model spatial features. This is augmented with channels (similar to the concept of RGB channels in a convolutional neural network for image data [31]) to include the time dimension where the channels are chronologically ordered. The number of input channels and output channels are parameters equal to number of input and output time snapshots respectively.

• •

  3D FNO: In 3D FNO, two dimensions model the spatial features and one dimension models the temporal features. This does not differentiate between the spatial and temporal dimensions.

• •

  Hybrid FNO-PDE: This method alternates between FNO and the PDE solver by using the output of FNO as input to PDE solver and vice versa. This could use either of the two FNO models listed above. It combines the speed of machine learning with the accurate physics of PDE solver. We leverage torchScript to execute this workflow [32].

In the above methods, we use PyTorch’s neural operator library [33] for FNO and particle-resolved direct numerical simulation [16] as the C++ PDE solver.

In this section, we study the errors of different FNO models for the spatiotemporal dynamics of decaying 2D turbulence. It should be noted here that the traditional PDE solvers have well understood time-stepping thresholds rooted in their stability analyses and requirements to capture the smallest scales of turbulence. However, in absence of an analogous thresholds for data-driven models one needs to rely on numerical experiments to decide the time resolution required for an accurate solution. Our previous study has demonstrated that a minimum of ten snapshots with a time resolution of $0.05\,t_{c}$ is necessary for making predictions at Reynolds number of 1000 [34]. In this paper, we increase the Reynolds number by approximately $10$ times and refine the time resolution by a factor of $10$. We will evaluate the models in terms of their number of parameters, training time, and error accumulated in long time roll-outs. In all the following studies, unless otherwise noted, the errors are reported as an average of 500 samples (all with different initial conditions) that were not a part of the training data.

The numerical experiments in this paper lead to a viable hybrid methodology for accelerating direct numerical simulation by using FNO. We have demonstrated that the FNO when trained on a sufficiently large data set leads to a pre-trained model which when coupled with the numerical PDE solver can be used to make stable long-term predictions. The 3D FNO models were found to have a longer training time and yield larger errors for decaying 2D turbulence because the periodic spatial dimension cannot be treated at the same footing as non-periodic temporal dimensions. Thus, we conclude that stacking the temporal dimension across channels leads to smaller errors and requires a shorter training time.

It should also be noted that the FNO model generalizes well as it was trained on data set generated using lattice Boltzmann but coupled with finite difference based Navier-Stokes solver. This was possible as the physical conditions for both the numerical methods were identical by design, however, a word of caution is in order as the generalizability of pre-trained machine learning models has been discussed recently in the context of “foundational” models. The FNO models discussed in this work have been trained on the data of decaying 2D turbulence for a specific range of Reynolds number. In order to generalize it further as a solution operator for Navier-Stokes equation one needs to embed more physics in the training. Foundational models to be trained on sufficiently diverse data set and should at the minimum replicate canonical test cases of fluid dynamics. Their training also needs to incorporate the knowledge of governing equations as explored in many physics-informed machine learning works [35, 36]. An extension of the present framework to 3D should be straightforward with 3D FNO for spatial and channels for temporal dimensions.

High-fidelity direct numerical simulation of turbulent flows for most real-world applications remains an outstanding challenge due to its computational cost. Although, the advantage of machine learning inference manifests in the form of reduced time-to-solution, there exist several costs that are not frequently discussed such as – (i) the computational costs of training a model, (ii) procuring a GPU, and (iii) energy cost, some of which are amortized over several inference steps. In our hybrid PDE-FNO model presented here, the PDE solver takes 20 secs on AMD EPYC 7413 24-Core Processor for evolving flow field over $0.025\,t_{c}$. The machine learning component takes 0.1 secs for host to device data transfer and 0.3 secs for FNO inference on Nvidia A6000 GPU. However, the degree of optimization of the software libraries as well as scalability of the methodologies vary. Most high performance computing (HPC) codes are written in C++ while machine learning frameworks are Python based. While coupling the two one also needs to account for the cost of transforming the data from C++ dynamic memory arrays into libTorch tensors. Several components discussed above often require optimization for real-world HPC applications which incur costs in terms of software redevelopment efforts. Thus, a fair comparison of the computational cost of different methodologies accounting for a wide range of metrics needs more research.

The authors gratefully acknowledge financial support from the Laboratory Directed Research and Development program at Brookhaven National Laboratory, which is sponsored by the US Department of Energy, Office of Science, under Contract Number DE-SC0012704. This research used resources of the National Energy Research Scientific Computing Center (NERSC), a Department of Energy Office of Science User Facility using NERSC award ASCR-ERCAP0027399.