A composable autoencoder-based algorithm for accelerating numerical simulations

1 INTRODUCTION . Numerical solutions to partial differential equations ( PDEs ) are dependent on PDE conditions such as , geometry of the computational domain , boundary conditions , initial conditions and source terms . Commercial PDE solvers have shown a tremendous success in accurately modeling PDEs for a wide range of applications . These solvers generalize across different PDE conditions but can be computationally slow . Moreover , their solutions are not reusable and need to be solved from scratch every time the PDE conditions are changed . The idea of using Machine Learning ( ML ) with PDEs has been explored for several decades ( Crutchfield & McNamara , 1987 ; Kevrekidis et al. , 2003 ) but with recent developments in computing hardware and ML techniques , these efforts have grown immensely . Although ML approaches are computationally fast , they fall short of traditional PDE solvers with respect to accuracy and generalization to a wide range of PDE conditions . Most data-driven and physics-constrained approaches employ static-inferencing strategies , where a mapping function is learnt between PDE solutions and corresponding conditions . In many cases , PDE conditions are sparse and high-dimensional , and hence , difficult to generalize . Additionally , current ML approaches do not make use of the key ideas from traditional PDE solvers such as , domain decomposition , solver methods , numerical discretization , constraint equations , symmetry evaluations and tighter non statistical evaluation metrics , which were established over several decades of research and development . In this work , we propose a novel ML approach that is motivated from such ideas and relies on dynamic inferencing strategies that afford the possibility of seamless coupling with traditional PDE solvers , when necessary . The proposed ML-approach , CoAE-MLSim ( Composable AutoEncoder Machine Learning Simulation ) , operates at the level of local subdomains , which consist of a group of pixels in 2D or voxels in 3D ( for example 8 or 16 in each spatial direction ) . As shown in Figure 1 , the CoAE-MLSim has 3 main components : A ) learn solutions on local subdomains , B ) learn the rules of how a group of local subdomains connect together to yield locally consistent solutions & C ) deploy an iterative algorithm to establish local consistency across all groups of subdomains in the entire computational domain . The solutions on subdomains are learnt into low-dimensional representations using solution autoencoders , while the rules between groups of subdomains , corresponding to flux conservation between PDE solution is learnt using flux conservation autoencoders . The autoencoders are trained a priori on local subdomains using very few PDE solution samples . During inference , the pretrained autoencoders are combined with iterative solution algorithms to determine PDE solutions for different PDEs over a wide range of sparse , high-dimensional PDE conditions . The solution strategy in the CoAE-MLSim approach is very similar to traditional PDE solvers , moreover the iterative inferencing strategy allows for coupling with traditional PDE solvers to accelerate convergence and improve accuracy and generalizability . Significant contributions of this work : 1 . The CoAE-MLSim approach combines traditional PDE solver strategies with ML techniques to accurately model numerical simulations . 2 . Our approach operates on local subdomains and solves PDEs in a low-dimensional space . This enables generalization to arbitrary PDE conditions within a high-dimensional distribution . 3 . The CoAE-MLSim approach training corresponds to training the various autoencoders on local subdomains and hence , it contains very few parameters and requires less data . 4 . The iterative inferencing algorithm is completely unsupervised and allows for coupling with traditional PDE solvers . 5 . Finally , we show generalization of our model for wide variations in sparse , high dimensional PDE conditions using evaluation metrics in tune with commercial solvers . 2 RELATED WORKS . The use of ML for solving PDEs has gained tremendous traction in the past few years . Most of the literature has revolved around improving neural network architectures and optimization techniques to enable generalizable and accurate learning of PDE solutions . More recently , there has been a lot of focus on learning mesh independent , infinite dimensional operators with neural networks ( NNs ) Bhattacharya et al . ( 2020 ) ; Anandkumar et al . ( 2020 ) ; Li et al . ( 2020b ) ; Patel et al . ( 2021 ) ; Lu et al . ( 2021b ) ; Li et al . ( 2020a ) . The neural operators are trained from high-fidelity solutions generated by traditional PDE solvers on a mesh of specific resolution and do not require any knowledge of the PDE . From the perspective of mesh independent learning , Battaglia et al . ( 2018 ) introduced a graph network architecture , which has proven to be effective in solving dynamical systems directly on unstructured computational meshes . Sanchez-Gonzalez et al . ( 2020 ) and Pfaff et al . ( 2020 ) use the graph network for robustly solving transient dynamics on arbitrary meshes and to accurately capture the transient solution trajectories . For dynamical systems , the neural operators and mesh based learning strategies have shown reasonable prediction capabilities , however , the generalizability of these methods to different PDE conditions remains to be seen . Another direction of research falls under the category of training neural networks with physics constrained optimization . Research in this space involves constraining neural networks with additional physics-based losses introduced through loss re-engineering . PDE constrained optimization have shown to improve interpretability and accuracy of results . Raissi et al . ( 2019 ) and Raissi & Karniadakis ( 2018 ) introduced the framework of physics-informed neural network ( PINN ) to constrain neural networks with PDE derivatives computed using Automatic Differentiation ( AD ) ( Baydin et al. , 2018 ) . In the past couple of years , the PINN framework has been extended to solve complicated PDEs representing complex physics ( Jin et al. , 2021 ; Mao et al. , 2020 ; Rao et al. , 2020 ; Wu et al. , 2018 ; Qian et al. , 2020 ; Dwivedi et al. , 2021 ; Haghighat et al. , 2021 ; Haghighat & Juanes , 2021 ; Nabian et al. , 2021 ; Kharazmi et al. , 2021 ; Cai et al. , 2021a ; b ; Bode et al. , 2021 ; Taghizadeh et al. , 2021 ; Lu et al. , 2021c ; Shukla et al. , 2021 ; Hennigh et al. , 2020 ; Li et al. , 2021 ) . More recently , alternate approaches that use discretization techniques using higher order derivatives and specialize numerical schemes to compute derivatives have shown to provide better regularization for faster convergence ( Ranade et al. , 2021b ; Gao et al. , 2021 ; Wandel et al. , 2020 ; He & Pathak , 2020 ) . However , the use of optimization techniques to solve PDEs , although accurate , has proved to be extremely slow as compared to traditional solvers . There is a large body of work related to training neural networks in tandem with differentiable PDE solvers in the loop to improve long range time predictions . The use of differentiable solvers to provide accurate estimates of adjoints has shown to improve learning and to provide better control of PDE solutions and transient system dynamics ( Amos & Kolter , 2017 ; Um et al. , 2020 ; de Avila BelbutePeres et al. , 2018 ; Toussaint et al. , 2018 ; Wang et al. , 2020 ; Holl et al. , 2020 ; Portwood et al. , 2019 ) . These frameworks are useful in providing rapid feedback to neural networks to improve convergence stability and enable efficient exploration of solution space . Bar-Sinai et al . ( 2019 ) , Xue et al . ( 2020 ) and Kochkov et al . ( 2021 ) train neural networks in conjunction with a differentiable solver to learn high-fidelity PDE discretization on coarse grids . Singh et al . ( 2017 ) and Holland et al . ( 2019 ) use differentiable solvers to tune model parameters in a live simulation . However , these methods require a differentiable solver in the training loop and commercial solvers that solve industrial scale problems may not be differentiable . All of the methods discussed above try to learn the learn the dynamics of the PDE in the solution space on entire computational domain , which can be a challenging task as domains can be high-dimensional with complicated physics . For many problems , local learning from smaller restricted domains have proved to accelerate learning of neural networks and provide better accuracy and generalization . Lu et al . ( 2021a ) and Wang et al . ( 2021 ) learn on localized domains but infer on larger computational domains using a stitching process . Bar-Sinai et al . ( 2019 ) and Kochkov et al . ( 2021 ) learn coefficients of numerical discretization schemes from high fidelity data , which is sub-sampled on coarse grids . Beatson et al . ( 2020 ) learns surrogate models for smaller components to allow for cheaper simulations . Other methods compress PDE solutions on to lower-dimensional manifolds . This has shown to improve accuracy and generalization capability of neural networks ( Wiewel et al. , 2020 ; Maulik et al. , 2020 ; Kim et al. , 2019 ; Murata et al. , 2020 ; Fukami et al. , 2020 ; Ranade et al. , 2021a ) . The ideas proposed in this work share the same goal with the aforementioned literature to accelerate numerical simulations without compromising on accuracy and generalizability . Our work draws inspiration from techniques such as local and latent space learning to solve PDEs in both transient and steady-state settings . More importantly , the approach proposed here relies heavily on unsupervised techniques such as autoencoders to solve PDEs Ranade et al . ( 2021a ; b ) ; Maleki et al . ( 2021 ) and mainly focus on leveraging knowledge from traditional solvers to develop a robust , stable , accurate and generalizable machine learning system . 3 COAE-MLSIM MODEL DETAILS . 3.1 SIMILARITIES WITH TRADITIONAL PDE SOLVERS . Consider a set of coupled PDEs with n solution variables . For the sake of notation simplicity , we take n = 2 , such that u ( x , y , z , t ) and v ( x , y , z , t ) are defined on a computational domain Ω with boundary conditions specified on the boundary of the computational domain , Ωb . It should be noted that extension to more solution variables is trivial . The coupled PDEs are defined as follows : L1 ( u , v ) − F1 = 0 ; L2 ( u , v ) − F2 = 0 ( 1 ) where , L1 , L2 denote PDE operators and F1 , F2 represent PDE source terms . The PDE operators can vary for different PDEs . For example , in a non-linear PDE such as the unsteady , incompressible Navier-Stokes equation the operator , L = ∂∂t + ~a.~∇− ~∇.~∇ Traditional PDE solvers solve PDEs given in Eq . 1 by representing solutions variables , u , v , and their linear and non-linear spatio-temporal derivatives on a discretized computational domain . The numerical approximations on discrete domains are computed on a finite number of computational elements known as a mesh , using techniques such as Finite Difference Method ( FDM ) , Finite Volume Method ( FVM ) and Finite Element Method ( FEM ) . These solvers use iterative solution algorithms to conserve fluxes between neighboring computational elements and determine consistent PDE solutions over the entire domain at convergence . The CoAE-MLSim approach is designed to perform similar operations but at the level of subdomains with assistance from ML techniques .

Summary: This paper proposes a new ML approach called CoAE-MLSim that is a faster alternative to PDE solvers. Compared to previous ML work on this problem, it aims to be more accurate and generalize better across PDE conditions. They also aim to require fewer PDE solutions to train the model.