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: The paper proposed CoAE-MLSim to learn with relatively fewer samples of PDE solutions and solve PDEs. CoAE-MLSim uses the idea of domain decomposition: first learn the solution on local subdomains using autoencoder, and then couple these sub-solutions by an iterative algorithm. Numerical experiments are performed to test the efficiency of the method.