Learning continuous-time PDEs from sparse data with graph neural networks

1 INTRODUCTION . We consider continuous dynamical systems with a state u ( x , t ) ∈ R that evolves over time t ∈ R+ and spatial locations x ∈ Ω ⊂ RD of a bounded domain Ω . We assume the system is governed by an unknown partial differential equation ( PDE ) u̇ ( x , t ) : = du ( x , t ) dt = F ( x , u , ∇xu , ∇2xu , . . . ) , ( 1 ) where the temporal evolution u̇ of the system depends on the current state u and its spatial first and higher-order partial derivatives w.r.t . the coordinates x . Such PDE models are the cornerstone of natural sciences , and are widely applicable to modelling of propagative systems , such as behavior of sound waves , fluid dynamics , heat dissipation , weather patterns , disease progression or cellular kinetics ( Courant & Hilbert , 2008 ) . Our objective is to learn the differential F from data . There is a long history of manually deriving mechanistic PDE equations for specific systems ( Cajori , 1928 ) , such as the Navier-Stokes fluid dynamics or the Schrödinger ’ s quantum equations , and approximating their solution forward in time numerically ( Ames , 2014 ) . These efforts are complemented by data-driven approaches to infer any unknown or latent coefficients in the otherwise known equations ( Isakov , 2006 ; Berg & Nyström , 2017 ; Santo et al. , 2019 ) , or in partially known equations ( Freund et al. , 2019 ; Seo & Liu , 2019b ; Seo et al. , 2020 ) . A series of methods have studied neural proxies of known PDEs for solution acceleration ( Lagaris et al. , 1998 ; Raissi et al. , 2017 ; Weinan & Yu , 2018 ; Sirignano & Spiliopoulos , 2018 ) or for uncertainty quantification ( Khoo et al. , 2017 ) . Related work . Recently the pioneering work of Long et al . ( 2017 ) proposed a fully non-mechanistic method PDE-Net , where the governing equation F is learned from system snapshot observations as a convolutional neural network ( CNN ) over the input domain discretised into a spatio-temporal grid . Further works have extended the approach with residual CNNs ( Ruthotto & Haber , 2019 ) , symbolic neural networks ( Long et al. , 2019 ) , high-order autoregressive networks ( Geneva & Zabaras , 2020 ) , and feed-forward networks ( Xu et al. , 2019 ) . These models are fundamentally limited to discretizing the input domain with a sample-inefficient grid , while they also do not support continuous evolution over time , rendering them unable to handle temporally or spatially sparse or non-uniform observations commonly encountered in realistic applications . Models such as ( Battaglia et al. , 2016 ; Chang et al. , 2016 ; Sanchez-Gonzalez et al. , 2018 ) are related to the interaction networks where object ’ s state evolves as a function of its neighboring objects , which forms dynamic relational graphs instead of grids . In contrast to the dense solution fields of PDEs , these models apply message-passing between small number of moving and interacting objects , which deviates from PDEs that are strictly differential functions . In Poli et al . ( 2019 ) graph neural ordinary differential equations ( GNODE ) were proposed as a framework for modeling continuous-time signals on graphs . The main limitations of this framework in application to learning PDEs are the lack of spatial information about physical node locations and lack of motivation for why this type of model could be suitable . Our work can be viewed as connecting graph-based continuous-time models with data-driven learning of PDEs in spatial domain through a classical PDE solution technique . Contributions . In this paper we propose to learn free-form , continuous-time , a priori fully unknown PDE model F from sparse data measured on arbitrary timepoints and locations of the coordinate domain Ω with graph neural networks ( GNN ) . Our contributions are : • We introduce continuous-time representation and learning of the dynamics of PDE-driven systems • We propose efficient graph representation of the domain structure using the method of lines with message passing neural networks • We achieve state-of-the-art learning performance on realistic PDE systems with irregular data , and our model is highly robust to data sparsity Scripts and data for reproducing the experiments can be found in this github repository . 2 METHODS . We apply the method of lines ( MOL ) ( Schiesser , 2012 ) to numerically solve Equation 1 . The MOL consists of selecting N nodes in Ω and discretizing spatial derivatives in F at these nodes . We place the nodes to the observation locations ( x1 , . . . , xN ) . The discretization leads to F being approximated by F̂ and produces the following system of ordinary differential equations ( ODEs ) whose solution asymptotically approximates the solution of Equation 1 u̇ ( t ) = u̇1 ( t ) ... u̇N ( t ) = du ( x1 , t ) dt ... du ( xN , t ) dt ≈ F̂ ( x1 , xN ( 1 ) , u1 , uN ( 1 ) ) ... F̂ ( xN , xN ( N ) , uN , uN ( N ) ) ∈ RN . ( 2 ) As the discretized F̂ inherits its unknown nature from the true PDE function F , we approximate F̂ by a learnable neural surrogate function . The system ’ s state at xi is defined as ui , while N ( i ) is a set of indices of neighboring nodes other than i that are required to evaluate F̂ at xi , and xN ( i ) with uN ( i ) are positions and states of nodes N ( i ) . This shows that the temporal derivative u̇i of ui depends not only on the location and state at the node i , but also on locations and states of neighboring nodes , resulting in a locally coupled system of ODEs . Each ODE in the system follows the solution at a fixed location xi . Numerous ODE solvers have been proposed ( such as Euler and Runge-Kutta solvers ) to solve the full system u ( t ) = u ( 0 ) + ∫ t 0 u̇ ( τ ) dτ , ( 3 ) where 0 ≤ τ ≤ t is a cumulative intermediate time variable . Solving equation 3 forward in time scales linearly both with respect to the number of nodes N and the number of evaluated time points M , while saturating the input space Ω requires a large number of nodes . In practice , PDEs are often applied for two- and three-dimensional spatial systems where the method is efficient . 2.1 POSITION-INVARIANT GRAPH NEURAL NETWORK DIFFERENTIAL . After introducing Equation 2 , we transition from learning F to learning F̂ . The value of F̂ at a node i must depend only on the nodes i and N ( i ) . Furthermore , the number of arguments and their order in F̂ is not known in advance and might be different for each node . This means that our model F̂ must be able to work with an arbitrary number of arguments and must be invariant to permutations of their order . Graph neural networks ( GNNs ) ( Wu et al. , 2020 ) satisfy these requirements . In a more restricted setting , where the number of neighbors and their order is known , ( e.g . if the grid is uniform ) other types of models such as multilayer perceptrons and convolutional neural networks can be used as well . We consider a type of GNNs called message passing neural networks ( MPNNs ) ( Gilmer et al. , 2017 ) to represent F̂ as F̂θ ( xN ( i ) − xi , ui , uN ( i ) ) , ( 4 ) where xN ( i ) − xi = { xj − xi : j ∈ N ( i ) } and θ denote parameters of the MPNN . This formulation assumes the absence of position-dependent quantities in F̂ , but models based on this formulation are invariant to translations and rotations of Ω , which makes generalization to systems with different node positions feasible , and prevents overfitting by memorizing position-specific dynamics . We use MPNNs , which is a type of spatial-based GNNs , due to their flexibility and computational efficiency . The main alternative – spectral-based GNNs – have relatively poor scaling with the number of nodes and learn global , or domain-dependent , filters due to the need to perform eigenvalue decomposition of the Laplacian matrix . 2.2 MESSAGE PASSING NEURAL NETWORKS . Let a graph G = ( V , E ) contain nodes V = { xi } Ni=1 , defined by the measurement positions , and undirected edges E = { eij } , and assume each node and edge are associated with a node feature vi and an edge feature eij , respectively . We use the node neighborhood N ( i ) to define edges . Neighbors for each node were selected by applying Delaunay triangulation to the measurement positions . Two nodes were considered to be neighbors if they lie on the same edge of at least one triangle ( Figure 1 ) . Delaunay triangulation has such useful properties as maximizing the minimum angle within each triangle in the triangulation and containing the nearest neighbor of each node which helps to obtain a good quality discretization of Ω . In message passing graph neural networks we propagate a latent state for K ≥ 1 graph layers , where each layer k consists of first aggregating messages m ( k ) i for each node i , and then updating the corresponding node states h ( k ) i , m ( k+1 ) i = ⊕ j∈N ( i ) φ ( k ) ( h ( k ) i , h ( k ) j , eij ) , ( 5 ) h ( k+1 ) i = γ ( k ) ( h ( k ) i , m ( k+1 ) i ) , ( 6 ) where ⊕ denotes a permutation invariant aggregation function ( e.g . sum , mean , max ) , and φ ( k ) , γ ( k ) are differentiable functions parameterized by deep neural networks . At any time τ , we initialise the latent states h ( 0 ) i = vi = ui ( τ ) and node features to the current state ui ( τ ) of the system . We define edge features eij : = xj − xi as location differences . Finally , we use the node states at the last graph layer of the MPNN to evaluate the PDE surrogate dû ( xi , t ) dt = F̂θ ( xN ( i ) − xi , ui , uN ( i ) ) = h ( K ) i , ( 7 ) which is used to solve Equation 3 for the estimated states û ( t ) = ( û ( x1 , t ) , . . . , û ( xN , t ) ) . 2.3 ADJOINT METHOD FOR LEARNING CONTINUOUS-TIME MPNN SURROGATES . Parameters of F̂θ are defined by θ which is the union of parameters of functions φ ( k ) , γ ( k ) , k = 1 , . . . , K in the MPNN . We fit θ by minimizing the mean squared error between the observed states ( y ( t0 ) , . . . , y ( tM ) ) and the estimated states ( û ( t0 ) , . . . , û ( tM ) ) , L ( θ ) = ∫ tM t0 ` ( t , û ) dt = ∫ tM t0 1 M + 1 M∑ i=0 ||û ( ti ) − y ( ti ) ||22δ ( t− ti ) dt ( 8 ) = 1 M + 1 M∑ i=1 ||û ( ti ) − y ( ti ) ||22 . ( 9 ) While discrete-time neural PDE models evaluate the system state only at measurement time points , more accurate continuous-time solution for the estimated state generally requires many more evaluations of the system state . If an adaptive solver is used to obtain the estimated states , the number of time steps performed by the solver might be significantly larger than M . The amount of memory required to evaluate the gradient of L ( θ ) by backpropagation scales linearly with the number of solver time steps . This typically makes backpropagation infeasible due to large memory requirements . We use an alternative approach , which allows computing the gradient for memory cost , which is independent from the number of the solver time steps . The approach was presented in Chen et al . ( 2018 ) for neural ODEs and is based on the adjoint method ( Pontryagin , 2018 ) . The adjoint method consists of a single forward ODE pass 3 until state û ( tM ) at the final time tM , and subsequent backward ODE pass solving the gradients . The backward pass is performed by first solving the adjoint equation λ̇ ( t ) T = ∂ ` ∂û ( t ) − λ ( t ) T ∂F̂ ∂û ( t ) . ( 10 ) for the adjoint variables λ from t = tM until t = 0 with λ ( tM ) = 0 , and then computing dL dθ = − ∫ T 0 λ ( t ) T ∂F̂ ∂θ dt ( 11 ) to obtain the final gradient .

Summary: This submission proposes extensions of PDE-net that relax some constraints that could help extend the range of applications of this approach. First, rather than fixing a spatial discretization in the form of a grid, the authors use a Delaunay triangulation to represent the domain. The updates to the nodes of this triangulation are performed using a message-passing GNN framework which couples neighboring nodes. Secondly, the authors use a classical adjoint method to allow for arbitrary time-discretizations (though this may be much more expensive in practice).