Learned Simulators for Turbulence

1 INTRODUCTION . Turbulent fluid dynamics are ubiquitous throughout the natural world , and many scientific and engineering disciplines depend on high quality simulations of turbulence . Examples range from design problems in aeronautics ( Rhie , 1983 ) and medicine ( Sallam & Hwang , 1984 ) to scientific problems of understanding molecular dynamics ( Smith , 2015 ) , weather ( Beljaars , 2003 ) , and vast galactic systems ( Canuto & Christensen-Dalsgaard , 1998 ) . Predicting turbulent dynamics is challenging , and finding a general solution to the governing equations of fluid dynamics , the Navier-Stokes Equation , is a famously open problem in Mathematical Physics ( Jaffe , 2006 ) . The challenge is due in large part to the fact that turbulence is chaotic : small-scale changes in the initial conditions can lead to a large difference in the outcome . Nonetheless , over the past several decades , high-quality numerical solvers have been engineered that integrate the governing partial differential equations ( PDEs ) , which can maintain accuracy over long integration periods even for complex dynamics . These solvers often must operate over high-resolution grids and therefore require substantial computational resources . Otherwise , high-frequency information is lost due to hard-to-model “ numerical viscosity ” , and the simulated dynamics can further diverge from the underlying equations due to the chaotic nature of turbulence . Because of the relevance of turbulence simulation to deep mathematical questions and real-world problems , efficient turbulence simulation is likely to be a central problem for Machine Learning in the years to come . Learned simulation represents a promising avenue for efficient fluid simulation because learned models can potentially adapt to capture the large-scale dynamics on coarse grids . Learned simulators vary in the extent to which they incorporate components from classical solvers , like those that learn closures and subgrid discretization ( Duraisamy et al. , 2019 ; Freund et al. , 2019 ; Pathak et al. , 2020 ; Um et al. , 2021 ; Kochkov et al. , 2021 ) , versus adopting a pure ML approach ( Kim et al. , 2018 ; Lusch et al. , 2018 ; Sanchez-Gonzalez et al. , 2020 ; Wang et al. , 2020 ; Li et al. , 2020c ; a ; Pfaff et al. , 2021 ) . Advantages of the pure ML route are that the same model can be used without specialized knowledge of the domain and that they do not involve the challenge of interfacing with PDE solvers . However , fully learned simulators often work well only when conditions are similar to the training distribution , which can limit their stability over time and their generalization capabilities . Our primary contribution is to demonstrate that fully learned simulators can learn to predict turbulent dynamics on coarse grids more accurately than classical solvers . This is true over a range of performance metrics , but where the learned simulators particularly excel in their ability to preserve high frequency information . This shows that neural networks can adapt to capture aspects of the large scale dynamics that classical solvers disproportionately miss . While learned simulators can be prohibitively unstable , we show that the simple practice of tuning training noise and the temporal downsampling factor can reliably produce stable rollouts . We introduce a simple model that combines neural networks architectures for grids ( Dilated CNNs ( Yu & Koltun , 2016 ) , ResNets ( He et al. , 2015 ) ) with recent advancements on general purpose learned simulation on graphs ( encode-process-decode models ( Sanchez-Gonzalez et al. , 2020 ; Pfaff et al. , 2021 ) trained with noise ( Sanchez-Gonzalez et al. , 2018 ; 2020 ; Pfaff et al. , 2021 ) ) . We implement this as well as a set of learned models from the literature that have been proposed for modeling turbulence , and show that Dilated ResNet outperforms the other , often more complicated models . Rather than looking at a single problem , we evaluate our models over a variety of challenging turbulent and chaotic domains , showing that Dilated ResNet is generally able to capture a variety of turbulent dynamics out-of-the-box . This includes two environments generated by the new Athena++ solver ( Stone et al. , 2020 ) , a state of the art ( magneto- ) hydrodynamics simulator used for astrophysics , where extremely high-resolution and large-scale simulations are often required . Finally , we evaluate the ability of the learned models to generalize to initial conditions , rollout durations , and environment sizes outside of the training distribution . Generalizing out of distribution is a known challenge for learned models , and we document where our models fail and what model and training choices can improve generalization . 2 RELATED WORK . Turbulent fluid dynamics are ubiquitous across the universe , and turbulence simulation is a crucial step in problems across disciplines . These include forecasting problems , like predicting the spread of a wildfire ( Leonardi et al. , 2019 ) or the evolution of a hurricane ( Zhang et al. , 2020 ) , engineering challenges , like designing more aerodynamic machinery ( Rhie , 1983 ) , and scientific questions , like understanding the physics of insect flight ( Dickinson et al. , 1999 ) or dynamics of celestial objects ( Fielding et al. , 2018 ) . Across applications , a variety of statistical measures are relevant for quantifying physical variables under turbulence ( Velocity Autocorrelation , variable histograms , Cooling , Energy spectra ) . The primary challenge for simulating turbulence is accurately capturing small- and large-scale flow dynamics , and their interplay , simultaneously . Traditional approaches to computational fluid dynamics integrate the governing PDEs describing the fluid flow , given initial and boundary conditions . State-of-the-art turbulence simulators , such as Athena++ , are built upon efficient algorithms that prioritize the preservation of high-order statistics . Such simulators are used frequently in scientific applications to numerically compare model predictions with experimental observations , but can be extremely–often prohibitively–expensive to run . Thus , there is significant practical interest in developing more efficient simulators . Deep learning models for predicting physical dynamics have traditionally been used for modeling low-dimensional dynamical systems with graph-like ( Battaglia et al. , 2016 ; Sanchez-Gonzalez et al. , 2018 ) and grid-like specifications ( Raissi et al. , 2017 ; Chen et al. , 2018 ) , but there is a growing body of work for fine-grained modeling of complex dynamics , as in particle-based liquids and soft-body physics ( Li et al. , 2019 ; Ummenhofer et al. , 2020 ; Sanchez-Gonzalez et al. , 2020 ) . Learned models are often used in computer graphics to speed up ( Wiewel et al. , 2019 ; Um et al. , 2018 ; Ladický et al. , 2015 ; Holden et al. , 2019 ) or super-resolve ( Xie et al. , 2018 ) simulations . A number of papers predict the steady-state flow around 2D airfoil wings , using CNN ( Thuerey et al. , 2020 ; Bhatnagar et al. , 2019 ) or GCN ( Belbute-Peres et al. , 2020 ) architectures , or the transient dynamics of various low Reynolds number flows using mesh representations ( Pfaff et al. , 2021 ) . In recent years , a number of models have been introduced specifically for learning turbulent dynamics . Many of these use neural networks to model some part of an otherwise classical , theoretically motivated numerical integration setup ( Duraisamy et al. , 2019 ; Freund et al. , 2019 ; Pathak et al. , 2020 ; Um et al. , 2021 ; Kochkov et al. , 2021 ; Tompson et al. , 2016 ) , PDE formulation ( Raissi et al. , 2017 ; Champion et al. , 2019 ) , or embedding space ( Lusch et al. , 2018 ) . Other models incorporate some information about turbulence simulation into the architecture , but like ours , are learned endto-end in a supervised way . Wang et al . ( Wang et al. , 2020 ) implemented a fully learned , neural network “ Turbulent Flow Network ” ( TF-Net ) that applies physically inspired decomposition with learned parameters and processes the components with U-Nets for fluid simulation . Li and colleagues introduced “ Fourier Neural Operators ” for turbulence simulation ( Li et al. , 2020a ) , which apply neural networks that convolution parameters in the Fourier domain . This model is part of a broader family of work ( Portwood et al. , 2019 ; Li et al. , 2020b ) which combines basis transforms ( Fourier , Multipole , etc ) with neural operators that process the transformed input . We implement models from these papers as baselines to compare to our learned simulator . 3 MODEL . 3.1 LEARNED SIMULATION FRAMEWORK ON GRIDS . We train a simulation model s : X → X that maps Xt to Xt+∆t , where X is a space of n+1 dimensional tensors which represent the state variables for each point on a grid with n spatial dimensions and one feature dimension at time t. The true physical dynamics applied over K time steps yields a trajectory of states which we denote ( Xt0 , ... , XtK ) . The “ rollout ” trajectory produced by iteratively applying s for K steps is denoted ( X̃t0 , X̃t1 , ... , X̃tK ) , where X̃t0 = Xt0 represents initial conditions given as input . Our learnable simulators ( Fig . 1a ) implement s using neural networks denoted NN ( with weights θ ) , and predict the next state as the difference from the current state , X̃t+∆t = s ( X̃t ; θ ) = X̃t + NN ( X̃t ; θ ) . 3.2 MODELS . In this work we study “ fully learned “ models , which learn the physical dynamics end-to-end without any hard-coded solver components . We implemented a novel Dilated ResNet model , as well as several baseline models . Detailed model parameters can be found in the Appendix . Dilated ResNet Encode-Process-Decode ( Dil-ResNet ) We designed this architecture to combine the encode-process-decode paradigm ( Sanchez-Gonzalez et al. , 2018 ; 2020 ) with dilated convolutional networks ( Fig . 1a ) . The encoder and decoder each consist of single linear convolutional layer . The processor consists of N = 4 dilated CNN blocks ( dCNNn ) connected in series with residual connections ( Yu & Koltun , 2016 ; He et al. , 2015 ) . Each block consists of 7 dilated CNN layers with dilation rates of ( 1 , 2 , 4 , 8 , 4 , 2 , 1 ) . A dilated rate of N , indicates each pixel is convolved with pixels that are multiples of N pixels away , such that N=1 reduces to regular convolution . This model is not specialized for turbulence modeling per se , and its components are general-purpose tools for improving performance in CNNs , but we demonstrate that it is effective for it . Residual connections help avoid vanishing gradients , and dilations allow long-range communication while preserving local structure , making it a good fit for modeling PDEs . All individual CNNs use a kernel size of 3 and have 48 output channels ( except the decoder , which has an output channel for each feature ) . Each individual CNN layer in the processor is immediately followed by a rectified linear unit ( ReLU ) activation function . The Encoder CNN and the Decoder CNNs do not use activations . Encode-U-Net-Decode ( U-Net ) We designed this architecture as a simplification of the TF-Net baseline explained below . It uses a simple linear CNN encoder and decoder ( same as Dil-ResNet ) and a standard U-Net block ( Ronneberger et al. , 2015 ) . As in ( Ronneberger et al. , 2015 ) , each CNN stack is made of 3 CNNs , followed by activations . Other hyperparameters of the U-Net also matched those used in ( Ronneberger et al. , 2015 ) . Turbulent Flow Net ( TF-Net ) ( Wang et al. , 2020 ) TF-Net uses a domain-specific variation of UNets to model turbulence , along with other architectural domain-specific elements inspired by RANSLES Hybrid Coupling in the encoder . Hyperparameters also matched those used in ( Ronneberger et al. , 2015 ; Wang et al. , 2020 ) . Constrained Turbulent Flow Net ( Con-TF-Net ) ( Wang et al. , 2020 ) Wang et al . also explored a constraint term in the loss which penalizes the divergence , which was constrained to be zero everywhere for the PDE in their incompressible turbulence experiments . We implemented a similar modetal , adapting the constrained function to match constraints or preserved quantities in our datasets . For KS-1D , the mean is constrained to be 0 ; for INCOMP-2D , the divergence is constrained to be 0 , and for COMPDECAY-3D the total energy is constrained to be the same as input ’ s total energy . Constrained Dilated ResNet ( Con-Dil-ResNet ) This model is Dil-ResNet ( described above ) with the additional constraint terms from ( Con-TF-Net ) added to the loss function . This allows us to examine the effect of the loss term separately from the other aspects of the TF-Net architecture . Fourier Neural Operator ( FNO ) ( Li et al. , 2020a ) This model uses the neural operator formalism introduced in earlier work ( Li et al. , 2020c ; b ) , applied in the Fourier domain . We implemented FNO for 1D and 2D spatial domains as in ( Li et al. , 2020a ) Hyperparameters are from ( Li et al. , 2020a ) . CNN padding We use periodic padding for spatial axes in the case of periodic boundary conditions . For the COOLMIXLAYER-3D which has a fixed boundary condition along the vertical axis , we mask the value at the boundary out from the loss at train time and set it to the ground truth at each step of the rollout . We also augment the input state with a binary feature to distinguish boundary pixels .

Summary: This paper present a  neural network-based model by training learned simulators at low spatial and temporal resolutions to capture turbulent dynamics generated at high resolution. It also shows that the proposed model can simulate turbulent dynamics more accurately than classical numerical solvers at the same low resolutions across various scientifically relevant metrics. Finally, a series of numerical examples including one-dimensional, two-dimensional, and three-dimensional nonlinear equations are used to demonstrate the feasibility of the algorithm. I reviewed this paper in NeurIPS 2021, which was titled “Learned Simulators for Astrophysical Turbulence”. The authors revised the paper based on previous reviews. Thanks for the efforts the authors have made.