Learning Time-dependent PDE Solver using Message Passing Graph Neural Networks

1 INTRODUCTION . Physical phenomena are generally modeled through partial differential equations ( PDEs ) that govern the dynamic evolution or static solution of a physical system . Numerically solving partial differential equations is an important aspect of scientific and mathematical modeling in a broad range of fields including physics , biology , material science , and finance . There have been many efforts to develop efficient and accurate numerical solvers for PDEs using different techniques including finite difference ( LeVeque , 2007 ; Shashkov & Steinberg , 2018 ) , finite volume ( Eymard et al. , 2000 ; Brenner et al. , 2008 ) , and finite element schemes ( Reddy , 2014 ; Wriggers , 2008 ) . While these methods have been successful in producing accurate solutions , major challenges for accelerating and reducing computational cost when the governing PDE is known , and also determining the governing PDE when the physical system is unknown , remains to be addressed , problems such as those in climate modeling , turbulent flow , contact problems , or plastic material deformation . With the recent developments in deep learning , faster algorithms have been proposed to evaluate the response of a physical system , using only observational data . While deep learning approaches , such as multi-layer perceptron ( MLP ) or convolutional neural networks ( CNNs ) , are powerful in learning PDE solutions , they are , however , restricted to a specific discretization of the physical domain in which they are trained . As a result , the learned model is limited to a specific domain and can not be generalized to solve on different domains or for different discretizations , although the underlying physics remains to be the same . New training is required for any change in the physical domain or discretization . Here we propose a discretization and domain invariant neural network time-dependent PDE solver based on message passing graph neural nets ( MPGNN ) which is trained on a sample domain with different discretizations . The trained MPGNN can then be used to solve for different discretization or even on other domains as long as the underlying physics remains the same . We further show that a recurrent version of MPGNN can be used to find a temporal sequence of solutions to a PDE . 2 RELATED WORKS . One class of neural-net-based PDE solvers focuses on using neural networks as proxies of PDEs and aims at finding the solution by minimizing a loss that corresponds to the solution satisfying the governing equations and the boundary conditions ( Raissi et al. , 2017a ; b ; Lagaris et al. , 1998 ; Weinan & Yu , 2018 ; Sirignano & Spiliopoulos , 2018 ; Khoo et al. , 2021 ) . Although such an approach helps to find the one-time solution of a PDE with an instance of parameters , a slight modification to the PDE parameters , boundary conditions , or the domain requires re-training of the network . Another approach to solving PDEs is to use convolutional neural nets and snapshots of observations over the discretized input domain and to learn the dynamic evolution of a PDE ( Long et al. , 2018 ; Shi et al. , 2020 ) . Further modifications such as using residual connections ( Ruthotto & Haber , 2020 ) , or autoregressive dense encoder-decoder ( Geneva & Zabaras , 2020 ) , or symbolic multi-layer neural network ( Long et al. , 2019 ) in addition to the CNN can be used to improve the results . While these models do not require prior knowledge of the PDE , they are limited to domain discretization ( as a result can not be generalized to arbitrary domains ) and are limited to certain time discretization ( as a result unable to handle temporally and spatially sparse or non-uniform observations ) . Inspired by the discretization techniques in solving PDEs , a class of methods uses observational data to learn the discretization approximation required for the updates in classical computational PDE solver methods ( Bar-Sinai et al. , 2019 ; Kochkov et al. , 2021 ; Zhuang et al. , 2021 ; Han et al. , 2018 ) . In this approach , a neural network is used for better interpolation at coarse scale to be used in the framework of traditional numerical discretization . These methods are used in conjunction with classical numerical methods and can improve the accuracy and accelerate the solutions of the traditional numerical schemes ( Kochkov et al. , 2021 ) . Although these methods have been shown to generalize to new parameter regimes of a PDE , they are still bounded to the initially trained discretization and can not be used for arbitrary domains without re-training . Lastly , a class of neural PDE solvers focus on graph representation of the discretized mesh datastructure to approximate the PDE solution ( Li et al. , 2020a ; b ; Iakovlev et al. , 2020 ; Belbute-Peres et al. , 2020 ) . The numerical solution of a PDE is an approximation of the solution on discrete locations comprising a discretized mesh of continuous space . Each node represents a region in the continuous space and the approximate solution of the PDE in that region is assigned to the representative node . The discretized mesh forms a graph where each node is used to model the state of the system and forms a connectivity graph connecting to the neighboring nodes . This method has successfully been used to solve time-independent PDEs with different mesh sizes on the same physical domain ( Li et al. , 2020a ) . The connectivity and the location of the nodes can further be optimized to learn the solution with different levels of precision ( Alet et al. , 2019 ) . If the PDE includes long-range interactions , which happens mostly in time-independent PDEs , a multi-level graph neural network framework to encapsulate long-range interactions can be used to improve the results ( Li et al. , 2020b ) . In contrast to time-independent PDEs , in the case of time-dependent PDEs , it has been shown that a continuous-time model similar to physics informed neural nets but with a graph neural network can be used to recover system ’ s dynamics with sparse observational data recorded at irregular times ( Iakovlev et al. , 2020 ; Poli et al. , 2019 ) . Recently , it have been shown that message passing graph neural networks can be used to implement powerful physical simulation engines ( Pfaff et al. , 2020 ; Sanchez-Gonzalez et al. , 2020 ) . The state of a physical system can be expressed using a particle-based method as a reduced order model . The particles are then expressed as nodes in a graph and the message passing neural network learns to compute the dynamics of the particles ( Sanchez-Gonzalez et al. , 2020 ) . In addition to particle-based methods , mesh-based methods have been shown to be successful in physical simulations ( Pfaff et al. , 2020 ) . Such graph-based models , first encodes the input data into a latent space and then process it in the latent space ( reduced model ) , and to obtain the physical results decode the data back to the physical space . Here , we first show why graph neural networks can generalize to learn fast PDE solvers inspired by finite difference schemes . We introduce domain invariant features and boundary conditions inspired by classical PDE solvers to improve the generalization of the learned PDE solver operator . With the introduced features , we show that message passing graph neural network architecture efficiently fits the classical PDE solvers and can learn time-stepping/solver operators for linear and nonlinear PDEs with different boundary conditions . We further demonstrate that our trained graph neural network solver can be generalized to solve PDEs on physical domains different from the domain that it is trained on . This is beneficial to train GNN on a sample of small domains for even with unknown dynamics , and further , explore the dynamic behavior on different larger physical domains . Lastly , we show that a recurrent version of our MPGNN can be used to predict the temporal sequence of solutions to a PDE . 3 TIME-DEPENDENT PDES . We consider continuous dynamical system u ( x , t ) ∈ R evolving over time t ∈ R+ and spatial coordinate x ∈ Ω ⊂ Rd where Ω is a bounded d-dimensional domain . We assume the system is governed by a partial differential equation of the form ut = N [ u ; λ ] ( 1 ) where N [ · ; λ ] denotes the linear/nonlinear differential operator ( s ) parameterized by the vector λ . In the above general form , the temporal evolution of the current state ut depends on the differential operator N [ · ; λ ] which may include various spatial derivatives of the state including∇u , ∇2u , etc . Depending on the differential operator N , appropriate boundary conditions on ∂Ω is required for a well-posed PDE with a unique solution . Such PDE model is the cornerstone of mathematical models and is widely used to model various systems , from fluid dynamics , thermal sciences , to acoustics , and quantum mechanics . As an example , ut = N [ u ; λ ] = λ1∇2u + λ2 · ∇u constitutes a convection diffusion equation for u ∈ R as variable of interest , where λ = λ1 , λ2 are the diffusitivity and the velocity field vector with which the quantity is moving with . The state of the system at each time can be obtained using its initial state and time integration as u ( x , t ) = u ( x , 0 ) + ∫ t 0 N ( u ; λ ) dt . Numerous numerical techniques such as Finite Elements , Spectral Methods , Finite Difference , or Finite Volume techniques have been developed to efficiently approximate the differential operator N ( · ; λ ) and solve for a dynamical system over time . In all numerical schemes , the domain is first discretized with a mesh , the differential operator is approximated locally using neighboring points , and the solution is calculated over small time steps using a time integrator such as Euler ’ s scheme , i.e. , un+1 ( xi ) = u n ( xi ) + δtF ( un ( xi ) , ∇un ( xi ) , ∇2un ( xi ) , · · · ; λ ( xi ) ) ( 2 ) where the superscript n shows the solution over discretised time tn , and the differential operator N ( u ; λ ) = F ( u , ∇u , ∇2u , · · · ; λ ) shows that it contains information about local spatial derivatives . As an example , consider solving heat equation , ut = D∇2u , where D is the diffusion constant and ∇2u = ∂2u/∂x2 + ∂2u/∂y2 , on an structured grid shown in Fig . 1a . Let uni , j be the discretized solution at time t = nδt and spatial location x = iδx and y = jδy where δt , δx , and δy are the time , horizontal , and vertical spatial discretization respectively . The time and spatial derivatives in the heat equation can be expanded using Taylor series at each discretized point , where ∂uni , j/∂t = ( u n+1 i , j − uni , j ) /δt , ∂2uni , j/∂x 2 = ( uni+1 , j − 2uni , j + uni+1 , j ) /δx2 , and etc . Re-writing the equation for an arbitrary discretized point , we find un+1i , j = u n i , j + δtF where F = α ( ui , j+1 − ui , j ) + α ( ui , j−1 − ui , j ) + β ( ui+1 , j − ui , j ) + β ( ui−1 , j − ui , j ) , where α = Dδt/δx2 and β = Dδt/δy2 . Solving for this equation for all the points along with the boundary conditions the updates for the discretized points can be achieved . Note that here the update rule can be seen as the summation of updates that only depend on neighboring points . Although in this example with a simple linear equation and a structured grid it was easy to find the update rule F , given an arbitrary domain that requires a triangular mesh for discretization ( see Fig . 1b ) and a nonlinear governing equation the update rule is not straight forward to be worked out . Our objective here is to learn the approximation of the differential operator using a graph representation of the domain with message passing neural networks . Since in a general PDE the differential operator F ( u , ∇u , ∇2u , · · · ; λ ) contains local spatial derivatives and only local neighboring values at a point are relevant to approximate the differential operator at that point . As a result , a graph neural network is a promising framework to approximate the above right-hand side for the next value predictions .

Summary: A message passing graph neural network is used to predict the time-dependent solution to a family of PDEs. The PDE parameters, initial conditions, and boundary conditions are input to the network, an it outputs the solution 1 time-step into the future. The authors show how to train such a system on the Heat equation, Navier-Stokes, and advection-diffusion equations. The flexibility of the message passing graph neural network permits its deployment on a diverse set of geometries and boundary conditions.