# Machine learning–accelerated computational fluid dynamics

See allHide authors and affiliations

Edited by Andrea L. Bertozzi, University of California, Los Angeles, CA, and approved March 25, 2021 (received for review January 29, 2021)

## Significance

Accurate simulation of fluids is important for many science and engineering problems but is very computationally demanding. In contrast, machine-learning models can approximate physics very quickly but at the cost of accuracy. Here we show that using machine learning inside traditional fluid simulations can improve both accuracy and speed, even on examples very different from the training data. Our approach opens the door to applying machine learning to large-scale physical modeling tasks like airplane design and climate prediction.

## Abstract

Numerical simulation of fluids plays an essential role in modeling many physical phenomena, such as weather, climate, aerodynamics, and plasma physics. Fluids are well described by the Navier–Stokes equations, but solving these equations at scale remains daunting, limited by the computational cost of resolving the smallest spatiotemporal features. This leads to unfavorable trade-offs between accuracy and tractability. Here we use end-to-end deep learning to improve approximations inside computational fluid dynamics for modeling two-dimensional turbulent flows. For both direct numerical simulation of turbulence and large-eddy simulation, our results are as accurate as baseline solvers with 8 to 10× finer resolution in each spatial dimension, resulting in 40- to 80-fold computational speedups. Our method remains stable during long simulations and generalizes to forcing functions and Reynolds numbers outside of the flows where it is trained, in contrast to black-box machine-learning approaches. Our approach exemplifies how scientific computing can leverage machine learning and hardware accelerators to improve simulations without sacrificing accuracy or generalization.

Simulation of complex physical systems described by nonlinear partial differential equations (PDEs) is central to engineering and physical science, with applications ranging from weather (1, 2) and climate (3, 4) and engineering design of vehicles or engines (5) to wildfires (6) and plasma physics (7). Despite a direct link between the equations of motion and the basic laws of physics, it is impossible to carry out direct numerical simulations at the scale required for these important problems. This fundamental issue has stymied progress in scientific computation for decades and arises from the fact that an accurate simulation must resolve the smallest spatiotemporal scales.

A paradigmatic example is turbulent fluid flow (8), underlying simulations of weather, climate, and aerodynamics. The size of the smallest eddy is tiny: For an airplane with chord length of 2 m, the smallest length scale (the Kolomogorov scale) (9) is

Here, we introduce a method for calculating the accurate time evolution of solutions to nonlinear PDEs, while using an order-of-magnitude coarser grid than is traditionally required for the same accuracy. This is a type of numerical solver that does not average unresolved degrees of freedom but instead uses discrete equations that give pointwise accurate solutions on an unresolved grid. We discover these algorithms using machine learning (ML), by replacing the components of traditional solvers most affected by the loss of resolution with learned alternatives. As shown in Fig. 1*A*, for a two-dimensional DNS of a turbulent flow our algorithm maintains accuracy while using *B*). We also apply the method to a high-resolution LES simulation of a turbulent flow and show similar performance enhancements, maintaining pointwise accuracy on

There has been a flurry of recent work using ML to improve turbulence modeling. One major family of approaches uses ML to fit closures to classical turbulence models based on agreement with high-resolution DNSs (21⇓⇓–24). While potentially more accurate than traditional turbulence models, these new models have not achieved reduced computational expense. Another major thrust uses “pure” ML, aiming to replace the entire Navier–Stokes simulation with approximations based on deep neural networks (25⇓⇓⇓⇓–30). A pure ML approach can be extremely efficient, avoiding the severe time-step constraints required for stability with traditional approaches. Because these models do not include the underlying physics, they often cannot enforce hard constraints, such as conservation of momentum and incompressibility. While these models often perform well on data from the training distribution, they often struggle with generalization. For example, they perform worse when exposed to novel forcing terms. We believe “hybrid” approaches that combine the best of ML and traditional numerical methods are more promising. For example, ML can replace (31) or accelerate (32) iterative solves used inside some simulation methods without reducing accuracy. Here we focus on hybrid models that use ML to correct errors in cheap, underresolved simulations (33⇓–35). These models borrow strength from the coarse-grained simulations and are potentially much faster than pure numerical simulations due to the reduced grid size.

In this work we design algorithms that accurately solve the equations on coarser grids by replacing the components most affected by the resolution loss with better-performing learned alternatives. We use data-driven discretizations (36, 37) to interpolate differential operators onto a coarse mesh with high accuracy (Fig. 1*C*). We train the model inside a standard numerical method for solving the underlying PDEs as a differentiable program, with the neural networks and the numerical method written in a framework [JAX (38)] supporting reverse-mode automatic differentiation. This allows for end-to-end gradient-based optimization of the entire algorithm, similar to prior work on density functional theory (39), molecular dynamics (40), and fluids (33, 34). The methods we derive are equation-specific and require high-resolution ground-truth simulations for training data. Since the dynamics of a PDE are local, the high-resolution simulations can be carried out on a small domain. The models remain stable during long simulations and have robust and predictable generalization properties, with models trained on small domains producing accurate simulations on larger domains, with different forcing functions and even with different Reynolds number. Comparison to pure ML baselines shows that generalization arises from the physical constraints inherent in the formulation of the method.

## Background

### Navier–Stokes.

Incompressible fluids are modeled by the Navier–Stokes equations:**1b**. The Reynolds number **1a**. Higher Reynolds number flows dominated by convection are more complex and thus generally harder to model; flows are considered “turbulent” if

DNS solves Eq. **1** directly, whereas LES solves a spatially filtered version. In the equations of LES, u is replaced by a filtered velocity **1a**, with the subgrid stress defined as

## Methods

### Learned Solvers.

Our principal aim is to accelerate DNS without compromising accuracy or generalization. To that end, we consider ML modeling approaches that enhance a standard CFD solver when run on inexpensive-to-simulate coarse grids. We expect that ML models can improve the accuracy of the numerical solver via effective superresolution of missing details. Unlike traditional numerical methods, our learned solvers are optimized to fit the observed manifold of solutions to the equations they solve, rather than arbitrary polynomials. Empirically, this can significantly improve accuracy over high-order numerical methods (36), although we currently lack a theoretical explanation. Because we want to train neural networks for approximation inside our solver, we wrote a new CFD code in JAX (38), which allows us to efficiently calculate gradients via automatic differentiation. Our base CFD code is a standard implementation of a finite volume method on a regular staggered mesh, with first-order explicit time stepping for convection, diffusion, and forcing and implicit treatment of pressure; for details see *SI Appendix*.

The algorithm works as follows. In each time step, the neural network generates a latent vector at each grid location based on the current velocity field, which is then used by the subcomponents of the solver to account for local solution structure. Our neural networks are convolutional, which enforces translation invariance and allows them to be local in space. We then use components from standard numerical methods to enforce inductive biases corresponding to the physics of the Navier–Stokes equations, as illustrated by the light gray boxes in Fig. 1*C*; the convective flux model improves the approximation of the discretized convection operator, the divergence operator enforces local conservation of momentum according to a finite volume method, and the pressure projection enforces incompressibility and the explicit time-step operator forces the dynamics to be continuous in time, allowing for the incorporation of additional time-varying forces. “DNS on a coarse grid” blurs the boundaries of traditional DNS and LES modeling and thus invites a variety of data-driven approaches. In this work we focus on two types of ML components: learned interpolation and learned correction. Both center on convection, the key term in Eq. **1** for turbulent flows.

### Learned Interpolation.

In a finite volume method, u denotes a vector field of volume averages over unit cells, and the cell-averaged divergence can be calculated via Gauss’ theorem by summing the surface flux over each face. This suggests that our only required approximation is calculating the convective flux *SI Appendix*, Fig. S2 and our prior work (36, 37) for examples.

### Learned Correction.

An alternative approach, closer in spirit to LES modeling, is to simply model a residual correction to the discretized Navier–Stokes equations (Eq. **1**) on a coarse grid (33, 34). Such an approach generalizes traditional closure models for LES but in principle can also account for discretization error. We consider learned correction models of the form

### Training.

The training procedure tunes the ML components of the solver to minimize the discrepancy between an expensive high-resolution simulation and a simulation produced by the model on a coarse grid. We accomplish this via supervised training where we use a cumulative pointwise error between the predicted and ground truth velocities as the loss function:

## Results

We take a utilitarian perspective on model evaluation: Simulation methods are good insofar as they demonstrate accuracy, computational efficiency, and generalization. DNS excels at accuracy and generalization but is not efficient. Useful ML methods for fluids should be faster than standard baselines (e.g., DNS) with the same accuracy. Although trained on specific flows, they must readily generalize to new simulation settings, such as different domains, forcings, and Reynolds numbers. In what follows, we first compare the accuracy and generalization of our method to both DNS and several existing ML-based approaches for simulations of two-dimensional turbulence flow. In particular, we first consider Kolmogorov flow (43), a parametric family of forced two-dimensional turbulent flows obeying the Navier–Stokes equation (Eq. **1**), with periodic boundary conditions and forcing

### Accelerating DNS.

The accuracy of a DNS quickly degrades once the grid resolution cannot capture the smallest details of the solution. In contrast, our ML-based approach strongly mitigates this effect. Fig. 2 shows the results of training and evaluating our model on Kolmogorov flows at Reynolds number

#### Accuracy.

The scalar vorticity field *B*, indicating the first three Lyapunov times. Fig. 2*A* shows the time evolution of the vorticity field for three different models: the learned interpolation matches the ground truth (2,048 × 2,048) more accurately than the

The learned interpolation model also produces an energy spectrum *C* compares the energy spectrum for learned interpolation and direct simulation at different resolutions after

#### Computational efficiency.

The ability to match DNS with a

#### Generalization.

In order to be useful, a learned model must accurately simulate flows outside of the training distribution. We expect our models to generalize well because they learn local operators: Interpolated values and corrections at a given point depend only on the flow within a small neighborhood around it. As a result, these operators can be applied to any flow that features similar local structures to those seen during training. We consider three different types of generalization tests: 1) larger domain size, 2) unforced decaying turbulent flow, and 3) Kolmogorov flow at a larger Reynolds number.

First, we test generalization to larger domain sizes with the same forcing. Our ML models have essentially the exact same performance as on the training domain, because they only rely upon local features of the flows (*SI Appendix*, Fig. S5 and see Fig. 5).

Second, we apply our model trained on Kolmogorov flow to decaying turbulence, by starting with a random initial condition with high wavenumber components and letting the turbulence evolve in time without forcing. Over time, the small scales coalesce to form large-scale structures, so that both the scale of the eddies and the Reynolds number vary. Fig. 3 shows that a learned discretization model trained on Kolmogorov flows with

Our final generalization test is harder: Can the models generalize to higher Reynolds number where the flows are more complex? The universality of the turbulent cascade (1, 48, 49) implies that at the size of the smallest eddies (the Kolmogorov length scale), flows “look the same” regardless of Reynolds number when suitably rescaled. This suggests that we can apply the model trained at one Reynolds number to a flow at another Reynolds number by simply rescaling the model to match the new smallest length scale. To test this we construct a new validation dataset for Kolmogorov flow with *A* shows that with this scaling our model achieves the accuracy of DNS running at 6× finer resolution. This degree of generalization is remarkable, given that we are now testing the model with a flow of substantially greater complexity. Fig. 4*B* visualizes the vorticity, showing that higher complexity is captured correctly, as is further verified by the energy spectrum shown in Fig. 4*C*.

### Comparison to Other ML Models.

Finally, we compare the performance of learned interpolation to alternative ML-based methods. We consider three popular ML methods: ResNet (51), encoder–processor–decoder (52, 53) architectures, and the learned correction model introduced earlier. These models all perform explicit time stepping without any additional latent state beyond the velocity field, which allows them to be evaluated with arbitrary forcings and boundary conditions and to use the time step based on the Courant–Friedrichs–Lewy condition. By construction, these models are invariant to translation in space and time and have similar runtime for inference (varying within a factor of 2). To evaluate training consistency, each model is trained nine times with different random initializations on the same Kolmogorov *SI Appendix*, and the models are evaluated on the same generalization tasks. We compare their performance using several metrics: time until vorticity correlation falls below 0.95 to measure pointwise accuracy for the flow over short time windows, the absolute error of the energy spectrum scaled by

Fig. 5 compares results across all considered configurations. Overall, we find that learned interpolation performs best, although learned correction is not far behind. We were impressed by the performance of the learned correction model, despite its weaker inductive biases. The difference in effective resolution for pointwise accuracy (

### Accelerating LES.

Finally, up until now we have illustrated our method for DNS of the Navier–Stokes equations. Our approach is quite general and could be applied to any nonlinear PDE. To demonstrate this, we apply the method to accelerate LES, the industry standard method for large-scale simulations where DNS is not feasible.

Here we treat the LES at high resolution as the ground-truth simulation and train an interpolation model on a coarser grid for Kolmogorov flows with Reynolds number

## Discussion

In this work we present a data-driven numerical method that achieves the same accuracy as traditional finite difference/finite volume methods but with much coarser resolution. The method learns accurate local operators for convective fluxes and residual terms and matches the accuracy of an advanced numerical solver running at 8 to

What outlook do our results suggest for speeding up three-dimensional turbulence? In general, the runtime T for efficient ML augmented simulation of time-dependent PDEs should scale like

In summary, our approach expands the Pareto frontier of efficient simulation in CFD, as illustrated in Fig. 1*A*. With ML-accelerated CFD, users may either solve expensive simulations much faster or increase accuracy without additional costs. To put these results in context, if applied to numerical weather prediction, increasing the duration of accurate predictions from 4 to 7 time units would correspond to approximately 30 y of progress (56). These improvements are possible due to the combined effect of two technologies still undergoing rapid improvements: modern deep learning models, which allow for accurate simulation with much more compact representations, and modern accelerator hardware, which allows for evaluating said models with a remarkably small increase in computational cost. We expect both trends to continue for the foreseeable future and to eventually impact all areas of computationally limited science.

## Data Availability

Source code for our models, including learned components, and training and evaluation datasets are available at GitHub (https://github.com/google/jax-cfd).

## Acknowledgments

We thank John Platt and Rif A. Saurous for encouraging and supporting this work and for important conversations and Yohai bar Sinai, Anton Geraschenko, Yi-fan Chen, and Jiawei Zhuang for important conversations.

## Footnotes

↵

^{1}D.K. and J.A.S. contributed equally to this work.- ↵
^{2}To whom correspondence may be addressed. Email: dkochkov{at}google.com, jamieas{at}google.com, brenner{at}seas.harvard.edu, or shoyer{at}google.com.

Author contributions: D.K., J.A.S., M.P.B., and S.H. designed research; D.K., J.A.S., A.A., Q.W., M.P.B., and S.H. performed research; D.K., J.A.S., A.A., Q.W., M.P.B., and S.H. analyzed data; and D.K., J.A.S., M.P.B., and S.H. wrote the paper.

Competing interest statement: The authors are employees of Google, which sells hardware and software for machine learning. A patent filing has been submitted based on this work.

This article is a PNAS Direct Submission.

↵*In our case the Pearson correlation reduces to a cosine distance because the flows considered here have mean velocity of 0.

This article contains supporting information online at https://www.pnas.org/lookup/suppl/doi:10.1073/pnas.2101784118/-/DCSupplemental.

- Copyright © 2021 the Author(s). Published by PNAS.

This open access article is distributed under Creative Commons Attribution-NonCommercial-NoDerivatives License 4.0 (CC BY-NC-ND).

## References

- ↵
- L. F. Richardson

- ↵
- ↵
- T. Schneider et al.

- ↵
- P. Neumann et al.

- ↵
- J. D. Anderson

- ↵
- A. Bakhshaii,
- E. A. Johnson

- ↵
- W. M. Tang,
- V. S. Chan

- ↵
- S. B. Pope

- ↵
- U. Frisch

- ↵
- R. D. Moser,
- S. W. Haering,
- R. Y. Gopal

- ↵
- ↵
- J. Boussinesq

- ↵
- G. Alfonsi

- ↵
- J. Smagorinsky

- ↵
- ↵
- Q. Malé et al.

- ↵
- ↵
- L. Esclapez et al.

- ↵
- C. P. Arroyo,
- P. Kholodov,
- M. Sanjosé,
- S. Moreau

- ↵
- ↵
- J. Ling,
- A. Kurzawski,
- J. Templeton

- ↵
- ↵
- R. Maulik,
- O. San,
- R. Adil,
- V. Prakash

- ↵
- A. Beck,
- D. Flad,
- C.-D. Munz

- ↵
- B. Kim et al.

- ↵
- Z. Li et al.

- ↵
- K. Bhattacharya,
- B. Hosseini,
- N. B. Kovachki,
- A. M. Stuart

- ↵
- R. Wang,
- K. Kashinath,
- M. Mustafa,
- A. Albert,
- R. Yu

- ↵
- ↵
- N. B. Erichson,
- M. Muehlebach,
- M. W. Mahoney

- ↵
- J. Tompson,
- K. Schlachter,
- P. Sprechmann,
- K. Perlin

*34th International Conference on Machine Learning*(2017). http://proceedings.mlr.press/v70/tompson17a/tompson17a.pdf. Accessed 10 December 2020. - ↵
- O. Obiols-Sales,
- A. Vishnu,
- N. Malaya,
- A. Chandramowlishwaran

- ↵
- J. Sirignano,
- J. F. MacArt,
- J. B. Freund

- ↵
- H. Larochelle,
- M. Ranzato,
- R. Hadsell,
- M. F. Balcan, and
- H. Lin

- K. Um,
- R. Brand,
- Y. R. Fei,
- P. Holl,
- N. Thuerey

- ↵
- J. Pathak et al.

- ↵
- Y. Bar-Sinai,
- S. Hoyer,
- J. Hickey,
- M. P. Brenner

- ↵
- J. Zhuang et al.

- ↵
- J. Bradbury et al.

- ↵
- L. Li et al.

- ↵
- H. Larochelle,
- M. Ranzato,
- R. Hadsell,
- M. F. Balcan, and
- H. Lin

- S. S. Schoenholz,
- E. D. Cubuk

- ↵
- K. Duraisamy

- ↵
- A. Griewank

- ↵
- G. J. Chandler,
- R. R. Kerswell

- ↵
- ↵
- Z. Ben-Haim et al.

- ↵
- K. Yang,
- Y.-F. Chen,
- G. Roumpos,
- C. Colby,
- J. Anderson

- ↵
- T. Lu,
- Y.-F. Chen,
- B. Hechtman,
- T. Wang,
- J. Anderson

- ↵
- A. N. Kolmogorov

- ↵
- ↵
- ↵
- K. He,
- X. Zhang,
- S. Ren,
- J. Sun

- ↵
- P. W. Battaglia et al.

- ↵
- A. Sanchez-Gonzalez et al.

- ↵
- H. Larochelle,
- M. Ranzato,
- R. Hadsell,
- M. F. Balcan, and
- H. Lin

- T. Nguyen,
- R. Baraniuk,
- A. Bertozzi,
- S. Osher,
- B. Wang

- ↵
- P.-J. Hoedt et al.

- ↵

## Citation Manager Formats

## Article Classifications

- Physical Sciences
- Applied Mathematics