Machine learning accelerated computational fluid dynamics

Dmitrii Kochkov,1, ∗ Jamie A. Smith,1, ∗ Ayya Alieva,1 Qing Wang,1 Michael P. Brenner,1, 2 and Stephan Hoyer1, †

1Google Research

2School of Engineering and Applied Sciences, Harvard University, Cambridge, MA 02138

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 theNavier-Stokes equations, but solving these equations at scale remains daunting, limited by thecomputational cost of resolving the smallest spatiotemporal features. This leads to unfavorabletrade-offs between accuracy and tractability. Here we use end-to-end deep learning to improveapproximations inside computational fluid dynamics for modeling two-dimensional turbulent flows.For both direct numerical simulation of turbulence and large eddy simulation, our results are asaccurate as baseline solvers with 8-10x finer resolution in each spatial dimension, resulting in 40-80xfold computational speedups. Our method remains stable during long simulations, and generalizes toforcing functions and Reynolds numbers outside of the flows where it is trained, in contrast to blackbox machine learning approaches. Our approach exemplifies how scientific computing can leveragemachine learning and hardware accelerators to improve simulations without sacrificing accuracy orgeneralization.

I. INTRODUCTION

Simulation of complex physical systems described bynonlinear partial differential equations are central to en-gineering and physical science, with applications rang-ing from weather [1, 2] and climate [3, 4] , engineer-ing design of vehicles or engines [5], to wildfires [6] andplasma physics [7]. Despite a direct link between theequations of motion and the basic laws of physics, it isimpossible to carry out direct numerical simulations atthe scale required for these important problems. Thisfundamental issue has stymied progress in scientific com-putation for decades, and arises from the fact that anaccurate simulation must resolve the smallest spatiotem-poral scales.

A paradigmatic example is turbulent fluid flow [8], under-lying simulations of weather, climate, and aerodynam-ics. The size of the smallest eddy is tiny: for an air-plane with chord length of 2 meters, the smallest lengthscale (the Kolomogorov scale) [9] is O(10−6)m. Classicalmethods for computational fluid dynamics (CFD), suchas finite differences, finite volumes, finite elements andpseudo-spectral methods, are only accurate if fields varysmoothly on the mesh, and hence meshes must resolvethe smallest features to guarantee convergence. For aturbulent fluid flow, the requirement to resolve the small-est flow features implies a computational cost scalinglike Re3, where Re = UL/ν, with U and L the typi-cal velocity and length scales and ν the kinematic vis-cosity. A tenfold increase in Re leads to a thousandfoldincrease in the computational cost. Consequently, di-rect numerical simulation (DNS) for e.g. climate and

∗ D.K. and J.A. S. contributed equally to this work.† To whom correspondence should be addressed. E-mail:dkochkov@google.com, jamieas@google.com, shoyer@google.com

weather are impossible. Instead, it is traditional to usesmoothed versions of the Navier Stokes equations [10, 11]that allow coarser grids while sacrificing accuracy, suchas Reynolds Averaged Navier-Stokes (RANS) [12, 13],and Large-Eddy Simulation (LES) [14, 15]. For exam-ple, current state-of-art LES with mesh sizes of O(10)to O(100) million has been used in the design of inter-nal combustion engines [16], gas turbine engines [17, 18],and turbo-machinery [19]. Despite promising progress inLES over the last two decades, there are severe limits towhat can be accurately simulated. This is mainly due tothe first-order dependence of LES on the sub-grid scale(SGS) model, especially for flows whose rate controllingscale is unresolved [20].

Here, we introduce a method for calculating the accuratetime evolution of solutions to nonlinear partial differen-tial equations, while using an order of magnitude coarsergrid than is traditionally required for the same accuracy.This is a novel type of numerical solver that does notaverage unresolved degrees of freedom, but instead usesdiscrete equations that give pointwise accurate solutionson an unresolved grid. We discover these algorithms us-ing machine learning, by replacing the components oftraditional solvers most affected by the loss of resolu-tion with learned alternatives. As shown in Fig. 1(a),for a two dimensional direct numerical simulation of aturbulent flow, our algorithm maintains accuracy whileusing 10× coarser resolution in each dimension, result-ing in a ∼ 80 fold improvement in computational timewith respect to an advanced numerical method of simi-lar accuracy. The model learns how to interpolate localfeatures of solutions and hence can accurately generalizeto different flow conditions such as different forcings andeven different Reynolds numbers [Fig. 1(b)]. We also ap-ply the method to a high resolution LES simulation ofa turbulent flow and show similar performance enhance-ments, maintaining pointwise accuracy on Re = 100, 000

arX

iv:2

102.

0101

0v1

[ph

ysic

s.fl

u-dy

n] 2

8 Ja

n 20

21

2

Training dataset

Forced turbulenceMore turbulentDecayingLarger domain

Generalization tests

(a)

(b)

Time until correlation < 0.95

Run

time

per

time

unit

(s)

512× 512

Direct simulation

Learned interpolation

86x speedup

1024× 1024

2048× 2048

4096× 4096

8192× 8192

1024× 1024

512× 512

256× 256

(c)

Convective flux

interpolation

interpolation

External forcing

Div

erge

nce

Exp

licit

times

tep

Pre

ssur

e pr

ojec

tion

Con

volu

tiona

lne

ural

net

wor

k

New

vel

ocity

Filt

erco

nstr

aint

s

Old

vel

ocity

FIG. 1. Overview of our approach and results. (a) Accu-racy versus computational cost with our baseline (direct sim-ulation) and ML accelerated (learned interpolation) solvers.The x axis corresponds to pointwise accuracy, showing howlong the simulation is highly correlated with the ground truth,whereas the y-axis shows the computational time needed tocarry out one simulation time-unit on a single TPU core. Eachpoint is annotated by the size of the corresponding spatialgrid: for details see the appendix. (b) Illustrative trainingand validation examples, showing the strong generalizationcapabilities of our model. (c) Structure of a single time stepfor our ”learned interpolation” model, with a convolutionalneural net controlling learned approximations inside the con-vection calculation of a standard numerical solver. ψ and urefer to advected and advecting velocity components. For dspatial dimensions there are d2 replicates of the convectiveflux module, corresponding to the flux of each velocity com-ponent in each spatial direction.

LES simulations using 8× times coarser grids with ∼ 40fold computational speedup.

There has been a flurry of recent work using machinelearning to improve turbulence modeling. One majorfamily of approaches uses ML to fit closures to classicalturbulence models based on agreement with high resolu-tion direct numerical simulations (DNS) [21–24]. Whilepotentially more accurate than traditional turbulencemodels, these new models have not achieved reducedcomputational expense. Another major thrust uses“pure” ML, aiming to replace the entire Navier Stokessimulation with approximations based on deep neuralnetworks [25–30]. A pure ML approach can be extremelyefficient, avoiding the severe time-step constraints re-

quired for stability with traditional approaches. Becausethese models do not include the underlying physics, theyoften struggle to enforce constraints, such as conservationof momentum and incompressibility. While these modelsoften perform well on data from the training distribu-tion, they often struggle with generalization. For exam-ple, they perform worse when exposed to novel forcingterms. A third approach, which we build upon in thiswork, uses ML to correct errors in cheap, under-resolvedsimulations [31–33]. These models borrow strength fromthe coarse-grained simulations.

In this work we design algorithms that accurately solvethe equations on coarser grids by replacing the compo-nents most affected by the resolution loss with betterperforming learned alternatives. We use data driven dis-cretizations [34, 35] to interpolate differential operatorsonto a coarse mesh with high accuracy [Fig. 1(c)]. Wetrain the solver inside a standard numerical method forsolving the underlying PDEs as a differentiable program,with the neural networks and the numerical method writ-ten in a framework (JAX [36]) supporting reverse-modeautomatic differentiation. This allows for end-to-end gra-dient based optimization of the entire algorithm, similarto prior work on density functional theory [37], moleculardynamics [38] and fluids [31, 32]. The methods we deriveare equation specific, and require training a coarse resolu-tion solver with high resolution ground truth simulations.Since the dynamics of a partial differential equation arelocal, the high resolution simulations can be carried outon a small domain. The models remains stable duringlong simulations and has robust and predictable gener-alization properties, with models trained on small do-mains producing accurate simulations on larger domains,with different forcing functions and even with differentReynolds number. Comparison to pure ML baselinesshows that generalization arises from the physical con-straints inherent in the formulation of the method.

II. BACKGROUND

A. Navier-Stokes

Incompressible fluids are modeled by the Navier-Stokesequations:

∂u

∂t=−∇ · (u⊗ u) +

1

Re∇2u− 1

ρ∇p+ f (1a)

∇ · u =0 (1b)

where u is the velocity field, f the external forcing, and⊗ denotes a tensor product. The density ρ is a constant,and the pressure p is a Lagrange multiplier used to en-force (1b). The Reynolds number Re dictates the bal-ance between the convection (first) or diffusion (second)terms in the right hand side of (1a). Higher Reynoldsnumber flows dominated by convection are more complex

3

and thus generally harder to model; flows are considered“turbulent” if Re � 1.

Direct numerical simulation (DNS) solves (1) directly,whereas large eddy simulation (LES) solves a spatiallyfiltered version. In the equations of LES, u is replaced bya filtered velocity u and an sub-grid term −∇·τ is addedto the right side of (1a), with the sub-grid stress definedas τ = u⊗ u − u ⊗ u. Because u⊗ u is un-modeled,solving LES also requires a choice of closure model for τas a function of u. Numerical simulation of both DNSand LES further requires a discretization step to approx-imate the continuous equations on a grid. Traditionaldiscretization methods (e.g., finite differences) convergeto an exact solution as the grid spacing becomes small,with LES converging faster because it models a smootherquantity. Together, discretization and closure models arethe two principle sources of error when simulating fluidson coarse grids [31, 39].

B. Learned solvers

Our principle aim is to accelerate DNS without compro-mising accuracy or generalization. To that end, we con-sider ML modeling approaches that enhance a standardCFD solver when run on inexpensive to simulate coarsegrids. We expect that ML models can improve the accu-racy of the numerical solver via effective super-resolutionof missing details [34]. Because we want to train neuralnetworks for approximation inside our solver, we wrotea new CFD code in JAX [36], which allows us to effi-ciently calculate gradients via automatic differentiation.Our CFD code is a standard implementation of a finitevolume method on a regular staggered mesh, with first-order explicit time-stepping for convection, diffusion andforcing, and implicit treatment of pressure; for details seethe appendix.

The algorithm works as follows: in each time-step, theneural network generates a latent vector at each grid lo-cation based on the current velocity field, which is thenused by the sub-components of the solver to account forlocal solution structure. Our neural networks are con-volutional, which enforces translation invariance and al-lows them to be local in space. We then use componentsfrom standard standard numerical methods to enforce in-ductive biases corresponding to the physics of the NavierStokes equations, as illustrated by the light gray boxes inFig. 1(c): the convective flux model improves the approx-imation of the discretized convection operator; the diver-gence operator enforces local conservation of momentumaccording to a finite volume method; the pressure projec-tion enforces incompressibility and the explicit time stepoperator forces the dynamics to be continuous in time,allowing for the incorporation of additional time varyingforces. “DNS on a coarse grid” blurs the boundaries oftraditional DNS and LES modeling, and thus invites avariety of data-driven approaches. In this work we focus

on two types of ML components: learned interpolationand learned correction. Both focus on the convectionterm, the key term in (1) for turbulent flows.

1. Learned interpolation (LI)

In a finite volume method, u denotes a vector field ofvolume averages over unit cells, and the cell-averaged di-vergence can be calculated via Gauss’ theorem by sum-ming the surface flux over the each face. This suggeststhat our only required approximation is calculating theconvective flux u ⊗ u on each face, which requires in-terpolating u from where it is defined. Rather than us-ing typical polynomial interpolation, which is suitablefor interpolation without prior knowledge, here we usean approach that we call learned interpolation based ondata driven discretizations [34]. We use the outputs ofthe neural network to generate interpolation coefficientsbased on local features of the flow, similar to the fixedcoefficients of polynomial interpolation. This allows us toincorporate two important priors: (1) the equation main-tains the same symmetries and scaling properties (e.g.,rescaling coordinates x→ λx) as the original equations,and (2) as the mesh spacing vanishes, the interpolationmodule retains polynomial accuracy so that our modelperforms well in the regime where traditional numericalmethods excel.

2. Learned correction (LC)

An alternative approach, closer in spirit to LES modeling,is to simply model a residual correction to the discretizedNavier-Stokes equations [(1)] on a coarse-grid [31, 32].Such an approach generalizes traditional closure modelsfor LES, but in principle can also account for discretiza-tion error. We consider learned correction models of theform ut = u∗t + LC(u∗t ), where LC is a neural networkand u∗t is the uncorrected velocity field from the numer-ical solver on a coarse grid. Modeling the residual isappropriate both from the perspective of a temporallydiscretized closure model, and pragmatically because therelative error between ut and u∗t in a single time stepis small. LC models have fewer inductive biases thanLI models, but they are simpler to implement and po-tentially more flexible. We also explored LC models re-stricted to take the form of classical closure models (e.g.,flow-dependent effective tensor viscosity models), but therestrictions hurt model performance and stability.

3. Training

The training procedure tunes the machine learning com-ponents of the solver to minimize the discrepancy be-tween an expensive high resolution simulation and a sim-ulation produced by the model on a coarse grid. We

4

DS 2

048×

2048

Time step = 0 Time step = 500 Time step = 1000Time step = 1500LI

64×

64DS

64×

64

10

5

0

5

10

Vorti

city

0 500 1000 1500 2000Time step

0.0

0.2

0.4

0.6

0.8

1.0

Vorti

city

corre

latio

n

DS 64 × 64DS 128 × 128DS 256 × 256DS 512 × 512DS 1024 × 1024DS 2048 × 2048LI 64 × 64

1 10 30Wavenumber k

106

107

108

109

Scal

ed e

nerg

y sp

ectru

m E

(k)k

5

0 2 4 6 8 10 12Simulation time(a) (b) (c)

FIG. 2. Learned interpolation achieves accuracy of direct simulation at ∼ 10× higher resolution. (a) Evolution of predictedvorticity fields for reference (DS 2048× 2048), learned (LI 64× 64) and baseline (DS 64× 64) solvers, starting from the sameinitial velocities. The yellow box traces the evolution of a single vortex. (b) Comparison of the vorticity correlation betweenpredicted flows and the reference solution for our model and DNS solvers. (c) Energy spectrum scaled by k5 averaged betweentime-steps 10000 and 20000, when all solutions have decorrelated with the reference solution.

accomplish this via supervised training where we use acumulative point-wise error between the predicted andground truth velocities as the loss function

L(x, y) =

tT∑ti

MSE(u(ti), u(ti)).

The ground truth trajectories are obtained by usinga high resolution simulation that is then coarsened tothe simulation grid. Including the numerical solver inthe training loss ensures fully “model consistent” train-ing where the model sees its own outputs as inputs[23, 32, 39], unlike typical a priori training where simu-lation is only performed offline. As an example, for theKolmogorov flow simulations below with Reynolds num-ber 1000, our ground truth simulation had a resolution of2048 cells along each spatial dimension. We subsamplethese ground truth trajectories along each dimension andtime by a factor of 32. For training we use 32 trajectoriesof 4800 sequential time steps each, starting from differ-ent random initial conditions. To evaluate the model,we generate much longer trajectories (tens of thousandsof time steps) to verify that models remain stable andproduce plausible outputs. We unroll the model for 32time steps when calculating the loss, which we find im-proves model performance for long time trajectories [32],in some cases using gradient checkpoints at each modelstep to reduce memory usage [40].

III. RESULTS

We take a utilitarian perspective on model evaluation:simulation methods are good insofar as they demonstrateaccuracy, computational efficiency and generalizability.In this case, accuracy and computational efficiency re-quire the method to be faster than the DNS baseline,

while maintaining accuracy for long term predictions;generalization means that although the model is trainedon specific flows, it must be able to readily generalize wellto new simulation settings, including to different forc-ings and different Reynolds numbers. In what follows,we first compare the accuracy and generalizability of ourmethod to both direct numerical simulation and severalexisting ML-based approaches for simulations of two di-mensional turbulence flow. In particular, we first con-sider Kolmogorov flow [41], a parametric family of forcedturbulent flows obeying the Navier-Stokes equation [(1)],with forcing f = sin(4y)x− 0.1u, where the second termis a velocity dependent drag preventing accumulation ofenergy at large scales [42]. Kolmogorov flow producesa statistically stationary turbulent flow, with flow com-plexity controlled by a single parameter, the Reynoldsnumber Re.

A. Accelerating DNS

The accuracy of a direct numerical simulation quicklydegrades once the grid resolution cannot capture thesmallest details of the solution. In contrast, our ML-based approach strongly mitigates this effect. Figure 2shows the results of training and evaluating our modelon Kolmogorov flows at Reynolds number Re = 1000.All datasets were generated using high resolution DNS,followed by a coarsening step.

1. Accuracy

The scalar vorticity field ω = ∂xuy − ∂yux is a conve-nient way to describe a two-dimensional incompressibleflows [42]. Accuracy can be quantified by correlating vor-

5

DS 2

048×

2048

Time step = 0 Time step = 400 Time step = 800 Time step = 1200LI

64×

64DS

64×

64

10

5

0

5

10

Vorti

city

0 500 1000 1500Time step

0.0

0.2

0.4

0.6

0.8

1.0

Vorti

city

corre

latio

n

DS 64 × 64DS 128 × 128DS 256 × 256DS 512 × 512DS 1024 × 1024DS 2048 × 2048LI 64 × 64

1 10 30Wavenumber k

107

108

Scal

ed e

nerg

y sp

ectru

m E

(k)k

5

0 2 4 6 8 10Simulation time(a) (b) (c)

FIG. 3. Learned interpolation generalizes well to decaying turbulence. (a) Evolution of predicted vorticity fields as a function oftime. (b) Vorticity correlation between predicted flows and the reference solution. (c) Energy spectrum scaled by k5 averagedbetween time-steps 2000 and 2500, when all solutions have decorrelated with the reference solution.

ticity fields,[43] C(ω, ω) between the ground truth solu-tion ω and the predicted state ω. Fig. 2 compares thelearned interpolation model (64 × 64) to fully resolvedDNS of Kolmogorov flow (2048 × 2048) using an initialcondition that was not included in the training set. Strik-ingly, the learned discretization model matches the point-wise accuracy of DNS with a ∼ 10 times finer grid. Theeventual loss of correlation with the reference solution isexpected due to the chaotic nature of turbulent flows;this is marked by a vertical grey line in Fig. 2(b), indi-cating the first three Lyapunov times. Fig. 2 (a) showsthe time evolution of the vorticity field for three differentmodels: the learned interpolation matches the groundtruth (2048× 2048) more accurately than the 512× 512baseline, whereas it greatly outperforms a baseline solverat the same resolution as the model (64× 64).

The learned interpolation model also produces a simi-lar energy spectrum E(k) = 1

2 |u(k)|2 to DNS. With de-creasing resolution, DNS cannot capture high frequencyfeatures, resulting in an energy spectrum that “tails off”for higher values of k. Fig. 2 (c) compares the energyspectrum for learned interpolation and direct simulationat different resolutions after 104 time steps. The learnedinterpolation model accurately captures the energy dis-tribution across the spectrum.

2. Computational efficiency

The ability to match DNS with a ∼ 10 times coarsergrid makes the learned interpolation solver much faster.We benchmark our solver on a single core of Google’sCloud TPU v4, a hardware accelerator designed for accel-erating machine learning models that is also suitable formany scientific computing use-cases [44–46]. The TPUis designed for high throughput vectorized operations,and extremely high throughput matrix-matrix multipli-cation in low precision (bfloat16). On sufficiently large

grid sizes (256 × 256 and larger), our neural net makesgood use of matrix-multiplication unit, achieving 12.5xhigher throughput in floating point operations per sec-ond than our baseline CFD solver. Thus despite using150 times more arithmetic operations, the ML solver isonly about 12 times slower than the traditional solver atthe same resolution. The 10× gain in effective resolutionin three dimensions (two space dimensions and time, dueto the Courant condition) thus corresponds to a speedupof 103/12 ≈ 80.

3. Generalization

In order to be useful, a learned model must accuratelysimulate flows outside of the training distribution. Weexpect our models to generalize well because they learnlocal operators: interpolated values and corrections ata given point depend only on the flow within a smallneighborhood around it. As a result, these operatorscan be applied to any flow that features similar localstructures as those seen during training. We considerthree different types of generalization tests: (1) largerdomain size, (2) unforced decaying turbulent flow, and(3) Kolmogorov flow at a larger Reynolds number.

First, we test generalization to larger domain sizes withthe same forcing. Our ML models have essentially theexact same performance as on the training domain, be-cause they only rely upon local features of the flows.(seeAppendix E and Fig. 5).

Second, we apply our model trained on Kolmogorov flowto decaying turbulence, by starting with a random initialcondition with high wavenumber components, and lettingthe turbulence evolve in time without forcing. Over time,the small scales coalesce to form large scale structures, sothat both the scale of the eddies and the Reynolds num-ber vary. Figure 3 shows that a learned discretization

6

DS 4

096×

4096

Time step = 0 Time step = 800 Time step = 1600Time step = 2400LI

128

×12

8DS

128

×12

8

10

5

0

5

10

Vorti

city

0 1000 2000 3000Time step

0.0

0.2

0.4

0.6

0.8

1.0

Vorti

city

corre

latio

n

DS 128 × 128DS 256 × 256DS 512 × 512DS 1024 × 1024DS 2048 × 2048DS 4096 × 4096LI 128 × 128

1 10 60Wavenumber k

106

107

108

109

1010

Scal

ed e

nerg

y sp

ectru

m E

(k)k

5

0 2 4 6 8 10 12Simulation time(a) (b) (c)

FIG. 4. Learned interpolations can be scaled to simulate higher Reynolds numbers without retraining. (a) Evolution ofpredicted vorticity fields as a function of time for Kolmogorov flow at Re = 4000. (b) Vorticity correlation between predictedflows and the reference solution. (c) Energy spectrum scaled by k5 averaged between time-steps 6000 and 12000, when allsolutions have decorrelated with the reference solution.

model trained on Kolmogorov flows Re = 1000 can matchthe accuracy of DNS running at ∼ 7 times finer resolu-tion. A standard numerical method at the same resolu-tion as the learned discretization model is corrupted bynumerical diffusion, degrading the energy spectrum aswell as pointwise accuracy.

Our final generalization test is harder: can the modelsgeneralize to higher Reynolds number where the flows aremore complex? The universality of the turbulent cascade[1, 47, 48] 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 oneReynolds number to a flow at another Reynolds numberby simply rescaling the model to match the new small-est length scale. To test this we construct a new datasetfor a Kolmogorov flow with Re = 4000. The theory oftwo-dimensional turbulence [49] implies that the smallest

eddy size decreases as 1/√

Re, implying that the small-est eddies in this flow are 1/2 that for original flow withRe = 1000. We therefore can use a trained Re = 1000model at Re = 4000 by simply halving the grid spac-ing. Fig. 4 (a) shows that with this scaling, our modelachieves the accuracy of DNS running at 7 times finerresolution. This degree of generalization is remarkable,given that we are now testing the model with a flow ofsubstantially greater complexity. Fig. 4 (b) visualizesthe vorticity, showing that higher complexity is capturedcorrectly, as is further verified by the energy spectrumshown in Fig. 4 (c).

B. Comparison to other ML models

Finally, we compare the performance of learned interpo-lation to alternative ML-based methods. We considerthree popular ML methods: ResNet (RN) [50], Encoder-

Processor-Decoder (EPD) [51, 52] architectures and thelearned correction (LC) model introduced earlier. Thesemodels all perform explicit time-stepping without anyadditional latent state beyond the velocity field, whichallows them to be evaluated with arbitrary forcings andboundary conditions, and to use the time-step based onthe CFL condition. By construction, these models areinivariant to translation in space and time, and havesimilar runtime for inference (varying within a factor oftwo). To evaluate training consistency, each model istrained 9 times with different random initializations onthe same Kolmogorov Re = 1000 dataset described pre-viously. Hyperparameters for each model were chosen asdetailed in Appendix G, and the models are evaluatedon the same generalization tasks. We compare their per-formance using several metrics: time until vorticity cor-relation falls below 0.95 to measure pointwise accuracyfor the flow over short time windows, the absolute errorof the energy spectrum scaled by k5 to measure statisti-cal accuracy for the flow over long time windows, and thefraction of simulated velocity values that does not exceedthe range of the training data to measure stability.

Fig. 5 compares results across all considered configura-tions. Overall, we find that learned interpolation (LI)performs best, although learned correction (LC) is notfar behind. We were impressed by the performance of theLC model, despite its weaker inductive biases. The dif-ference in effective resolution for pointwise accuracy (8×vs 10× upscaling) corresponds to about a factor of twoin run-time. There are a few isolated exceptions wherepure black box methods outperform the others, but notconsistently. A particular strength of the learned inter-polation and correction models is their consistent perfor-mance and generalization to other flows, as seen from thenarrow spread of model performance for different randominitialization and their consistent dominance over othermodels in the generalization tests. Note that even a mod-

7

est 4× effective coarse-graining in resolution still corre-sponds to a 5× computational speed-up. In contrast, theblack box ML methods exhibit high sensitivity to ran-dom initialization and do not generalize well, with muchless consistent statistical accuracy and stability.

C. Acceleration of LES

Finally, up until now we have illustrated our method forDNS of the Navier Stokes equations. Our approach isquite general and could be applied to any nonlinear par-tial differential equation. To demonstrate this, we ap-ply the method to accelerate LES, the industry standardmethod for large scale simulations where DNS is not fea-sible.

Here we treat the LES at high resolution as the groundtruth simulation and train an interpolation model on acoarser grid for Kolmogorov flows with Reynolds num-ber Re = 105 according to the Smagorinsky-Lilly modelSGS model [8]. Our training procedure follows the ex-act same approach we used for modeling DNS. Note inparticular that we do not attempt to model the param-eterized viscosity in the learned LES model, but ratherlet learned interpolation model this implicitly. Fig. 6shows that learned interpolation for LES still achievesan effective 8× upscaling, corresponding to roughly 40×speedup.

IV. DISCUSSION

In this work we present a data driven numerical methodthat achieves the same accuracy as traditional finite dif-ference/finite volume methods but with much coarser res-olution. The method learns accurate local operators forconvective fluxes and residual terms, and matches theaccuracy of an advanced numerical solver running at 8–10× finer resolution, while performing the computation40–80× faster. The method uses machine learning to in-terpolate better at a coarse scale, within the frameworkof the traditional numerical discretizations. As such, themethod inherently contains the scaling and symmetryproperties of the original governing Navier Stokes equa-tions. For that reason, the methods generalize much bet-ter than pure black-box machine learned methods, notonly to different forcing functions but also to differentparameter regimes (Reynolds numbers).

What outlook do our results suggest for speeding up 3Dturbulence? In general, the runtime T for efficient MLaugmented simulation of time-dependent PDEs shouldscale like

T ∼ (CML + Cphysics)

(N

K

)d+1

, (2)

where CML is the cost of ML inference per grid point,Cphysics is the cost of baseline numerical method, N

0 2 4 6

LILC

EPDRN

Forc

ed tu

rbul

ence 16x2x 4x 8x1x

Pointwise accuracy

108109

16x2x 4x 8x1xStatistical accuracy

0.0 0.5 1.0

Stability

0 2 4 6

LILC

EPDRNLa

rger

dom

ain 8x1x 16x2x 4x

107108109

8x1x 16x2x4x

0.0 0.5 1.0

0 2 4 6

LILC

EPDRNDe

cayi

ng

16x2x 4x 8x1x

108

16x2x 4x 8x1x

0.0 0.5 1.0

0 2 4 6Time until correl. < 0.95

LILC

EPDRNM

ore

turb

ulen

t 8x1x 16x2x 4x

108109

Spect. error (decreasing)

8x1x 16x2x4x

0.0 0.5 1.0Proportion stable at t=50

FIG. 5. Learned discretizations outperform a wide range ofbaseline methods in terms of accuracy, stability and general-ization. Each row within a subplot shows performance metricsfor nine model replicates with the same architecture, but dif-ferent randomly initialized weights. The models are trainedon forced turbulence, with larger domain, decaying and moreturbulent flows as generalization tests. Vertical lines indicateperformance of non-learned baseline models at different reso-lutions (all baseline models are perfectly stable). The point-wise accuracy test measures error at the start of time integra-tion, whereas the statistical accuracy and stability tests areboth performed on simulations at time 50 after about 7000time integration steps (twice the number of steps for more tur-bulent flow). The points indicated by a left-pointing trianglein the statistical accuracy tests are clipped at a maximumerror.

is the number of grid points along each dimension ofthe resolved grid, d is the number of spatial dimensionsand K is the effective coarse graining factor. Currently,CML/Cphysics ≈ 12, but we expect that much more effi-cient machine learning models are possible, e.g., by shar-ing work between time-steps with recurrent neural nets.We expect the 10× decrease in effective resolution discov-ered here to generalize to 3D and more complex problems.This suggests that speed-ups in the range of 103–104 maybe possible for 3D simulations. Further speed-ups, as re-quired to capture the full range of turbulent flows, will re-quire either more efficient representations for flows (e.g.,based on solution manifolds rather than a grid) or beingsatisfied with statistical rather than pointwise accuracy(e.g., as done in LES modeling).

In summary, our approach expands the Pareto frontier ofefficient simulation in computational fluid dynamics, asillustrated in Fig. 1(a). With ML accelerated CFD, usersmay either solve expensive simulations much faster, or in-

8

DS 2

048×

2048

Time step = 0 Time step = 400 Time step = 800 Time step = 1200LI

64×

64DS

64×

64

10

5

0

5

10

Vorti

city

0 500 1000 1500Time step

0.0

0.2

0.4

0.6

0.8

1.0

Vorti

city

corre

latio

n

DS 64 × 64DS 128 × 128DS 256 × 256DS 512 × 512DS 1024 × 1024DS 2048 × 2048LI 64 × 64

1 10 30Wavenumber k

106

107

108

109

1010

Scal

ed e

nerg

y sp

ectru

m E

(k)k

5

0 2 4 6 8 10Simulation time(a) (b) (c)

FIG. 6. Learned discretizations achieve accuracy of LES simulation running on 8× times finer resolution. (a) Evolution ofpredicted vorticity fields as a function of time. (b) Vorticity correlation between predicted flows and the reference solution.(c) Energy spectrum scaled by k5 averaged between time-steps 3800 and 4800, when all solutions have decorrelated with thereference solution.

crease accuracy without additional costs. To put theseresults in context, if applied to numerical weather predic-tion, increasing the duration of accurate predictions from4 to 7 time-units would correspond to approximately 30years of progress [53]. These improvements are possi-ble due to the combined effect of two technologies stillundergoing rapid improvements: modern deep learningmodels, which allow for accurate simulation with muchmore compact representations, and modern acceleratorhardware, which allows for evaluating said models witha remarkably small increase in computational cost. Weexpect both trends to continue for the foreseeable future,and to eventually impact all areas of computationallylimited science.

ACKNOWLEDGEMENT

We thank John Platt and Rif A. Saurous for encourag-ing and supporting this work and for important conver-sations, and Yohai bar Sinai, Anton Geraschenko, Yi-fan Chen and Jiawei Zhuang for important conversa-tions.

Appendix A: Direct numerical simulation

Here we describe the details of the numerical solver thatwe use for data generation, model comparison and thestarting point of our machine learning models. Our solveruses a staggered-square mesh [54] in a finite volume for-mulation: the computational domain is broken into com-putational cells where the velocity field is placed on theedges, while the pressure is solved at the cell centers.Our choice of real-space formulation of the Navier-Stokesequations, rather than a spectral method is motivated bypractical considerations: real space simulations are much

more versatile when dealing with boundary conditionsand non-rectangular geometries. We now describe theimplementation details of each component.

Convection and diffusion

We implement convection and diffusion operators basedon finite-difference approximations. The Laplace opera-tor in the diffusion is approximated using a second ordercentral difference approximation. The convection term issolved by advecting all velocity components simultane-ously, using a high order scheme based on Van-Leer fluxlimiter [55]. For the results presented in the paper weused explicit time integration using Euler discretization.This choice is motivated by performance considerations:for the simulation parameters used in the paper (highReynold number) implicit diffusion is not required forstable time-stepping, and is approximately twice as slow,due to the additional linear solves required for each veloc-ity component. For diffusion dominated simulations im-plicit diffusion would result in faster simulations.

Pressure

To account for pressure we use a projection method,where at each step we solve the corresponding Pois-son equation. The solution is obtained using either afast diagonalization approach with explicit matrix mul-tiplication [56] or a real-valued fast Fourier transform(FFT). The former is well suited for small simulation do-mains as it has better accelerator utilization, while FFThas best computational complexity and performs best inlarge simulations. For wall-clock evaluations we choosebetween the fast diagonalization and FFT approach bychoosing the fastest for a given grid.

9

Forcing and closure terms

We incorporate forcing terms together with the acceler-ations due to convective and diffusive processes. In anLES setting the baseline and ground truth simulationsadditionally include a subgrid scale model that is alsotreated explicitly. We use the Smagorinsky-Lilly model((A1)) where ∆ is the grid spacing and Cs = 0.2:

τij = −2 (Cs∆)2|S|Sij (A1)

Sij =1

2(∂iuj + ∂jui) |S| = 2

√∑i,j

SijSij

Appendix B: Datasets and simulation parameters

In the main text we introduced five datasets: two Kol-mogorov flows at Re = 1000 and Re = 4000, Kolmogorovflow with Re = 1000 on a 2× larger domain, decaying tur-bulence and an LES dataset with Reynolds number 105.Dataset generation consisted of three steps: (1) burn-insimulation from a random initial condition; (2) simula-tion for a fixed duration using high resolution solver; (3)downsampling of the solution to a lower grid for trainingand evaluation.

The burn-in stage is fully characterized by the burn-intime and initialization wavenumber which represent thediscarded initial transient and the peak wavenumber ofthe log-normal distribution from which random initialconditions were sampled from. The maximum amplitudeof the initial velocity field was set to 7 for forced sim-ulations and 4.2 in the decaying turbulence, which wasselected to minimize the burn-in time, as well as maintainstandard deviation of the velocity field close to 1.0. Theinitialization wavenumber was set to 4. Simulation pa-rameters include simulation resolution along each spatialdimension, forcing and Reynolds number. The resultingvelocity trajectories are then downsampled to the saveresolution. Note that besides the spatial downsamplingwe also perform downsampling in time to maintain theCourant–Friedrichs–Lewy (CFL) factor fixed at 0.5, stan-dard for numerical simulations. Parameters specifying allfive datasets are shown in Table A1. We varied the size ofour datasets from 12200 time slices at the downsampledresolution for decaying turbulence to 34770 slices in theforced turbulence settings. Such extended trajectoriesallow us to analyze stability of models when performinglong-term simulations. We note that when performingsimulations at higher Reynolds numbers we used rescaledgrids to maintain fixed time stepping. This is motivatedto potentially allow methods to specialize on a discretetime advancements. When reporting the results, spatialand time dimensions are scaled back to a fixed value toenable direct comparisons across the experiments.

For comparison of our models with numerical methods

Dataset Resolution Re Burn-in timeKolmogorov Re = 1000 2048→ 64 1000 40Kolmogorov Re = 4000 4096→ 128 4000 40Kolmogorov Re = 1000 4096→ 128 1000 40Decaying turbulence 2048→ 64 1000 (t = 0) 4LES Re = 105 2048→ 64 105 40

TABLE A1. Simulation parameters used for generation ofthe datasets in this manuscript. For resolution, the notationX → Y indicates that simulation was performed at resolutionX and the result was downsampled to resolution Y .

we use the corresponding simulation parameters whilechanging only the resolution of the underlying grid. Asmentioned in the Appendix A, when measuring the per-formance of all solvers on a given resolution we choosesolver components to maximize efficiency.

Appendix C: Learned interpolations

To improve numerical accuracy of the convective pro-cess our models use psuedo-linear models for interpola-tion [34, 35]. The process of interpolating ui to u(x) isbroken into two steps:

1. Computation of local stencils for interpolation tar-get x.

2. Computation of a weighted sum∑n

i=1 aiui over thestencil.

Rather using a fixed set of interpolating coefficients ai(e.g., as done for typical polynomial interpolation), wechoose ai from the output of a neural network addition-ally constrained to satisfy

∑ni=1 ai = 1, which guaran-

tees that the interpolation is at least first order accurate.We do so by generating interpolation coefficients with anaffine transformation a = Ax + b on the unconstraintedoutputs x of the neural network, where A is the null-space of the constraint matrix (a 1 × n matrix of ones)of shape n × (n − 1) and b is an arbitrary valid set ofcoefficients (we use linear interpolation from the nearestsource term locations). This is a special case of the pro-cedure for arbitrary polynomial accuracy constraints onfinite difference coefficients described in [34, 35]. In thiswork, we use a 4 × 4 patch centered over the top-rightcorner of each unit-cell, which means we need 15 uncon-strained neural network outputs to generate each set ofinterpolation coefficients.

Appendix D: Neural network architectures

All of the ML based models used in this work are basedon fully convolutional architectures. Fig. A1 depictsour three modeling approaches (pure ML, learned in-terpolation and learned correction) and three architec-turs for neural network sub-components (Basic ConvNet,

10

Neural network

Relu

Conv2D 64 3x3

Conv2D 64 3x3

Conv2D 64 5x5

Relu

Conv2D 2 5x5

Conv2D 64 5x5

Physics

Relu

Conv2D 64 3x3

Conv2D 3x3

Neural network

PhysicsNeural network

Pure ML Learned CorrectionLearned Interpolation

Conv2D 2 3x3

Relu

Conv2D 64 3x3

Encoder-Process-Decoder ResNetBasic ConvNet

Conv2D 2 3x3

Relu

Conv2D 64 3x3

FIG. A1. Modeling approaches and neural network ar-chitecutres used in this paper. Rescaling of inputs and out-puts from neural network sub-modules with fixed constantsis not depicted. Dashed lines surround components that arerepeatedly applied the indicated number of times.

Encoder-Process-Decoder and ResNet). Outputs and in-puts for each neural network layer are linearly scaled suchthat appropriate values are in the range (−1, 1).

For our physics augmented solvers (learned interpola-tion and learned correction), we used the basic ConvNetarchicture, with Nout = 120 (8 interpolations that need15 inputs each) for learned interpolation and Nout = 2for learned correction. Our experiments found accuracywas slightly improved by using larger neural networks,but not sufficiently to justify the increased computationalcost.

For pure ML solvers, we used EPD and ResNet mod-els that do not build impose physical priors beyond timecontinuity of the governing equations. In both cases aneural network is used to predict the acceleration due tophysical processes that is then summed with the forcingfunction to compute the state at the next time. BothEPD and ResNet models consist of a sequence of CNNblocks with skip-connections. The main difference is that

the EPD model has an arbitrarily large hidden state (inour case, size 64), whereas the ResNet model has a fixedsize hidden state equal to 2, the number of velocity com-ponents.

Appendix E: Details of accuracy measurements

In the main text we have presented accuracy results basedon correlation between the predicted flow configurationsand the reference solution. We reach the same conclu-sions based on other common choices, such as squaredand absolute errors as shown in Fig. A2 comparinglearned discretizations to DNS method on Kolmogorovflow with Reynolds number 1000.

As mentioned in the main text, we additionally evaluatedour models on enlarged domains while keeping the flowcomplexity fixed. Because our models are local, this isthe simplest generalization test. As shown in Fig. A3(and Fig. 5 in the main text), the improvement for 2×larger domains is identical to that found on a smallerdomain.

Appendix F: Details of overview figure

The Pareto frontier of model performance shown inFig. 1(a) is based on extrapolated accuracy to an 8×larger domain. Due to computational limits, we wereunable to run the simulations on the 16384× 16384 gridfor measuring accuracy, so the time until correlation goesbelow 0.95 was instead taken from the 1× domain size,which as described in the preceding section matched per-formance on the 2× domain.

The 512×512 learned interpolation model corresponds tothe depicted results in Fig. 2, whereas the 256× 256 and1024× 1024 models were retrained on the same trainingdataset for 2× coarser or 2× finer coarse-graining.

Appendix G: Hyperparameter tuning

All models were trained using Adam [57] optimizer. Ourinitial experiments showed that none of the models hada consistent dependence on the optimization parameters.The set that worked well for all models included learningrate of 10−3, b1 = 0.9 and b2 = 0.99. For each modelwe performed a hyperparameter sweep over the lengthof the training trajectory tT used to compute the loss,model capacity such as size and number of hidden layers,as well as a sweep over different random initializations toassess reliability and consistency of training.

[1] L. F. Richardson, Weather prediction by numerical pro-cess (Cambridge university press, 2007).

[2] P. Bauer, A. Thorpe, and G. Brunet, The quiet revolutionof numerical weather prediction, Nature 525, 47 (2015).

11

0.0

0.2

0.4

0.6

0.8

1.0Vo

rticit

y co

rrela

tion

DNS 64 x 64DNS 128 x 128DNS 256 x 256DNS 512 x 512DNS 1024 x 1024LD 64 x 64

0 1000 2000 3000Time step

0.0

0.2

0.4

0.6

0.8

1.0

Abso

lute

vel

ocity

erro

r

0 1000 2000 3000Time step

0.0

0.2

0.4

0.6

0.8

1.0Sq

uare

d ve

locit

y er

ror

(a)

(b) (c)

FIG. A2. Comparison of ML + CFD model to DNS runningat different resolutions using varying metrics. (a) Vorticitycorrelation as presented in the main text. (b) Mean absoluteerror. (c) Absolute error of the kinetic energy.

[3] T. Schneider, J. Teixeira, C. S. Bretherton, F. Brient,K. G. Pressel, C. Schar, and A. P. Siebesma, Climategoals and computing the future of clouds, Nature ClimateChange 7, 3 (2017).

[4] P. Neumann, P. Duben, P. Adamidis, P. Bauer,M. Bruck, L. Kornblueh, D. Klocke, B. Stevens, N. Wedi,and J. Biercamp, Assessing the scales in numericalweather and climate predictions: will exascale be the res-cue?, Philosophical Transactions of the Royal Society A377, 20180148 (2019).

[5] J. Anderson, Basic philosophy of cfd, in ComputationalFluid Dynamics (Springer, 2009) pp. 3–14.

[6] A. Bakhshaii and E. A. Johnson, A review of a new gener-ation of wildfire–atmosphere modeling, Canadian Journalof Forest Research 49, 565 (2019).

[7] W. M. Tang and V. Chan, Advances and challenges incomputational plasma science, Plasma physics and con-trolled fusion 47, R1 (2005).

[8] S. B. Pope, Turbulent flows (Cambridge University Press,2000).

[9] U. Frisch, Turbulence: the legacy of AN Kolmogorov(Cambridge University Press, 1995).

[10] R. D. Moser, S. W. Haering, and G. R. Yalla, Statisti-cal properties of subgrid-scale turbulence models, AnnualReview of Fluid Mechanics 53 (2020).

[11] C. Meneveau and J. Katz, Scale-invariance and turbu-lence models for large-eddy simulation, Annual Reviewof Fluid Mechanics 32, 1 (2000).

[12] J. Boussinesq, Theorie de l’ecoulement tourbillant,Memoires de l’Academie des Sciences 23, 46 (1877).

[13] G. Alfonsi, Reynolds-Averaged Navier–Stokes equationsfor turbulence modeling, Appl. Mech. Rev. 62 (2009).

[14] J. Smagorinsky, General circulation experiments withthe primitive equations: I. the basic experiment, Mon.Weather Rev. 91, 99 (1963).

[15] M. Lesieur and O. Metais, New trends in Large-Eddysimulations of turbulence, Annu. Rev. Fluid Mech. 28,45 (1996).

[16] Q. Male, G. Staffelbach, O. Vermorel, A. Misdariis,F. Ravet, and T. Poinsot, Large eddy simulation of Pre-Chamber ignition in an internal combustion engine, FlowTurbul. Combust. 103, 465 (2019).

[17] P. Wolf, G. Staffelbach, L. Y. M. Gicquel, J.-D. Muller,and T. Poinsot, Acoustic and large eddy simulation stud-ies of azimuthal modes in annular combustion chambers,Combust. Flame 159, 3398 (2012).

[18] L. Esclapez, P. C. Ma, E. Mayhew, R. Xu, S. Stouffer,T. Lee, H. Wang, and M. Ihme, Fuel effects on lean blow-out in a realistic gas turbine combustor, Combust. Flame181, 82 (2017).

[19] C. P. Arroyo, P. Kholodov, M. Sanjose, and S. Moreau,CFD modeling of a realistic turbofan blade for noise pre-diction. part 1: Aerodynamics, in Proceedings of GlobalPower and Propulsion Society (2019).

[20] S. B. Pope, Ten questions concerning the large-eddy sim-ulation of turbulent flows, New Journal of Physics 6(2004).

[21] J. Ling, A. Kurzawski, and J. Templeton, Reynoldsaveraged turbulence modelling using deep neural net-works with embedded invariance, J. Fluid Mech. 807,155 (2016).

[22] K. Duraisamy, G. Iaccarino, and H. Xiao, Turbulencemodeling in the age of data, Annual Review of FluidMechanics 51, 357 (2019).

[23] R. Maulik, O. San, A. Rasheed, and P. Vedula, Subgridmodelling for two-dimensional turbulence using neuralnetworks, Journal of Fluid Mechanics 858, 122 (2019).

[24] A. Beck, D. Flad, and C.-D. Munz, Deep neural networksfor data-driven les closure models, Journal of Computa-tional Physics 398, 108910 (2019).

[25] B. Kim, V. C. Azevedo, N. Thuerey, T. Kim, M. Gross,and B. Solenthaler, Deep fluids: A generative network forparameterized fluid simulations, in Computer GraphicsForum, Vol. 38 (Wiley Online Library, 2019) pp. 59–70.

[26] Z. Li, N. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhat-tacharya, A. Stuart, and A. Anandkumar, Neural oper-ator: Graph kernel network for partial differential equa-tions, arXiv preprint arXiv:2003.03485 (2020).

[27] K. Bhattacharya, B. Hosseini, N. B. Kovachki, and A. M.Stuart, Model reduction and neural networks for para-metric pdes, arXiv preprint arXiv:2005.03180 (2020).

[28] R. Wang, K. Kashinath, M. Mustafa, A. Albert, andR. Yu, Towards physics-informed deep learning for tur-bulent flow prediction, in Proceedings of the 26th ACMSIGKDD International Conference on Knowledge Dis-covery & Data Mining (2020) pp. 1457–1466.

[29] B. Lusch, J. N. Kutz, and S. L. Brunton, Deep learningfor universal linear embeddings of nonlinear dynamics,Nature communications 9, 1 (2018).

[30] N. B. Erichson, M. Muehlebach, and M. W. Mahoney,Physics-informed autoencoders for lyapunov-stable fluidflow prediction, arXiv preprint arXiv:1905.10866 (2019).

[31] J. Sirignano, J. F. MacArt, and J. B. Freund, Dpm: Adeep learning pde augmentation method with applica-tion to large-eddy simulation, Journal of ComputationalPhysics 423, 109811 (2020).

[32] K. Um, P. Holl, R. Brand, N. Thuerey, et al., Solver-in-the-loop: Learning from differentiable physics to interactwith iterative pde-solvers, in 34th Conference on Neural

12

DNS

4096

×40

96Time step = 0 Time step = 500 Time step = 1000Time step = 1500

ML+

CFD

128×

128

DNS

128×

128

10.0 7.5 5.0 2.5 0.0 2.5 5.0 7.5 10.0Vorticity

0 500 1000 1500 2000Time step

0.0

0.2

0.4

0.6

0.8

1.0

Vorti

city

corre

latio

n

DNS 128 × 128DNS 256 × 256DNS 512 × 512DNS 1024 × 1024DNS 2048 × 2048DNS 4096 × 4096ML+CFD 128 × 128

1 10 30Wavenumber k

105

106

107

108

109

1010

Scal

ed e

nerg

y sp

ectru

m E

(k)k

5

0 2 4 6 8 10 12Simulation time(a) (b) (c)

FIG. A3. Comparison of ML + CFD model to DNS running at different resolutions when evaluated on a 2× larger domainwith characteristic length-scales matching those of the training data. (a) Visualization of the vorticity at different times. (b)Vorticity correlation as a function of time. (c) Scaled energy spectrum E(k) ∗ k5.

Information Processing Systems (NeurIPS 2020) (2020)p. in press.

[33] J. Pathak, M. Mustafa, K. Kashinath, E. Motheau,T. Kurth, and M. Day, Using machine learning to aug-ment coarse-grid computational fluid dynamics simula-tions, arXiv preprint arXiv:2010.00072 (2020).

[34] Y. Bar-Sinai, S. Hoyer, J. Hickey, and M. P. Brenner,Learning data-driven discretizations for partial differen-tial equations, Proceedings of the National Academy ofSciences 116, 15344 (2019).

[35] J. Zhuang, D. Kochkov, Y. Bar-Sinai, M. P. Bren-ner, and S. Hoyer, Learned discretizations for passivescalar advection in a 2-d turbulent flow, arXiv preprintarXiv:2004.05477 (2020).

[36] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson,C. Leary, D. Maclaurin, and S. Wanderman-Milne, Jax:composable transformations of python+ numpy pro-grams, 2018, URL http://github. com/google/jax , 18(2020).

[37] L. Li, S. Hoyer, R. Pederson, R. Sun, E. D. Cubuk, P. Ri-ley, and K. Burke, Kohn-Sham equations as regularizer:building prior knowledge into machine-learned physics,Phys. Rev. Lett. , in press (2020).

[38] S. S. Schoenholz and E. D. Cubuk, JAX, M.D.: A frame-work for differentiable physics, in 34th Conference onNeural Information Processing Systems (NeurIPS 2020)(2020) p. in press.

[39] K. Duraisamy, Machine learning-augmented reynolds-averaged and large eddy simulation models of turbulence,(2020), arXiv:2009.10675 [physics.flu-dyn].

[40] A. Griewank, Achieving logarithmic growth of temporaland spatial complexity in reverse automatic differentia-tion, Optim. Methods Softw. 1 (1994).

[41] G. J. Chandler and R. R. Kerswell, Invariant recurrentsolutions embedded in a turbulent two-dimensional kol-mogorov flow, J. Fluid Mech. 722, 554 (2013).

[42] G. Boffetta and R. E. Ecke, Two-Dimensional turbulence,Annu. Rev. Fluid Mech. 44, 427 (2012).

[43] In our case the Pearson correlation reduces to a cosinedistance because the flows considered here have mean

velocity of 0.[44] Z. Ben-Haim, V. Anisimov, A. Yonas, V. Gulshan,

Y. Shafi, S. Hoyer, and S. Nevo, Inundation modelingin data scarce regions, (2019), arXiv:1910.05006 [cs.LG].

[45] K. Yang, Y.-F. Chen, G. Roumpos, C. Colby, and J. An-derson, High performance monte carlo simulation of isingmodel on TPU clusters, in Proceedings of the Interna-tional Conference for High Performance Computing, Net-working, Storage and Analysis, SC ’19 No. Article 83(Association for Computing Machinery, New York, NY,USA, 2019) pp. 1–15.

[46] T. Lu, Y.-F. Chen, B. Hechtman, T. Wang, and J. An-derson, Large-Scale discrete fourier transform on TPUs,(2020), arXiv:2002.03260 [cs.MS].

[47] A. N. Kolmogorov, The local structure of turbulence inincompressible viscous fluid for very large reynolds num-bers, Cr Acad. Sci. URSS 30, 301 (1941).

[48] R. H. Kraichnan and D. Montgomery, Two-dimensionalturbulence, Reports on Progress in Physics 43, 547(1980).

[49] G. Boffetta and R. E. Ecke, Two-dimensional turbulence,Annual Review of Fluid Mechanics 44, 427 (2012).

[50] K. He, X. Zhang, S. Ren, and J. Sun, Deep residual learn-ing for image recognition, in Proceedings of the IEEE con-ference on computer vision and pattern recognition (2016)pp. 770–778.

[51] P. W. Battaglia, J. B. Hamrick, V. Bapst, A. Sanchez-Gonzalez, V. Zambaldi, M. Malinowski, A. Tacchetti,D. Raposo, A. Santoro, R. Faulkner, et al., Relational in-ductive biases, deep learning, and graph networks, arXivpreprint arXiv:1806.01261 (2018).

[52] A. Sanchez-Gonzalez, J. Godwin, T. Pfaff, R. Ying,J. Leskovec, and P. W. Battaglia, Learning to simu-late complex physics with graph networks, arXiv preprintarXiv:2002.09405 (2020).

[53] P. Bauer, A. Thorpe, and G. Brunet, The quiet revolutionof numerical weather prediction, Nature 525, 47 (2015).

[54] J. M. McDonough, Lectures in computational fluid dy-namics of incompressible flow: Mathematics, algorithmsand implementations, (2007).

13

[55] P. K. Sweby, High resolution schemes using flux limitersfor hyperbolic conservation laws, SIAM journal on nu-merical analysis 21, 995 (1984).

[56] R. E. Lynch, J. R. Rice, and D. H. Thomas, Direct so-lution of partial difference equations by tensor productmethods, Numer. Math. 6, 185 (1964).

[57] D. P. Kingma and J. Ba, Adam: A method for stochasticoptimization, arXiv preprint arXiv:1412.6980 (2014).