Learning the Dynamics of Physical Systems from Sparse Observations with Finite Element Networks

1 INTRODUCTION The laws driving the physical world are often best described by partial differential equations ( PDEs ) that relate how a magnitude of interest changes in time with its change in space . They describe how the atmosphere and oceans circulate and interact , how structures deform under load and how electromagnetic waves propagate ( Courant & Hilbert , 2008 ) . Knowledge of these equations lets us predict the weather ( Coiffier , 2011 ) , build sturdier structures , and communicate wirelessly . Yet , in many cases we only know the PDEs governing a system partially ( Isakov , 2006 ) or not at all , or solving them is too computationally costly to be practical ( Ames , 2014 ) . Machine learning researchers try to fill in these gaps with models trained on collected data . For example , neural networks have been trained for weather forecasts ( Shi et al. , 2015 ) and fluid flow simulations ( Belbute-Peres et al. , 2020 ) , both of which are traditionally outcomes of PDE solvers . Even the dynamics of discrete dynamical systems such as traffic ( Li et al. , 2018 ) and crowds ( Zhang et al. , 2017 ) have been learned from data . A challenge facing these models is the high cost of acquiring training data , so the data is usually only available sparsely distributed in space . Since graphs are a natural way to structure sparse data , models incorporating graph neural networks ( GNNs ) have been particularly successful for spatio-temporal forecasting ( Yu et al. , 2018 ; Wu et al. , 2019 ) . In the domain of physical processes we can reasonably assume that the observed system follows a PDE . There are mainly two ways to incorporate this assumption as a-priori knowledge into a model . First , we can encode a known PDE into a loss function that encourages the model to fulfill the equation ( Raissi et al. , 2019 ) . Another way to go about this is to derive the model structure itself from known laws such as the convection-diffusion equation ( de Bézenac et al. , 2018 ) . In this paper we will follow the second approach . Consider a dynamical system on a bounded domain Ω ⊂ Rd that is governed by the PDE ∂tu = F ( t , x , u , ∂xu , ∂ 2 xu , ... ) ( 1 ) on functions u : [ 0 , T ] × Ω → Rm . If we have a dense measurement u0 : Ω → Rm of the current state of the system and a solution u that satisfies Eq . ( 1 ) for all t ∈ [ 0 , T ] and also fulfills the initial condition u ( 0 , x ) = u0 ( x ) at all points x ∈ Ω , we can use u as a forecast for the state of the system until time T . From a spatio-temporal forecasting perspective , this means that we can forecast the evolution of the system if we have a continuous measurement of the state , know the dynamics F , and can find solutions of Eq . ( 1 ) efficiently . Unfortunately , in practice we only have a finite number of measurements at arbitrary points and only know the dynamics partially or not at all . Contributions . An established numerical method for forecasts in systems with fully specified dynamics is the finite element method ( FEM ) ( Brenner et al. , 2008 ) . In this paper , we introduce the first graph-based model for spatio-temporal forecasting that is derived from FEM in a principled way . Our derivation establishes a direct connection between the form of the unknown dynamics and the structure of the model . Through this connection our model can incorporate prior knowledge on the governing physical processes via assumptions on the form of the underlying dynamics . We employ this mechanism to derive a specialized model for transport problems from the convection equation . The way that the model structure arises from the underlying equation makes our models uniquely interpretable . We show that our transport model disentangles convection and the remainder of the learned dynamics such as source/sink behavior , and that the activations of the model correspond to a learned flow field , which can be visualized and analyzed . In experiments on multi-step forecasting of sea surface temperature and gas flow , our model improves upon baselines from recurrent , temporalconvolutional , and continuous-time model classes with further improvement by the transport model . 2 BACKGROUND . 2.1 FINITE ELEMENT METHOD . In the following , we will outline how to approximate a solution u to the dynamics in Eq . ( 1 ) from an initial value u0 by discretizing u in space using finite elements . Let X be a set of points with a triangulation T of d-dimensional , non-overlapping simplices X = { x ( i ) ∈ Rd } Ni=1 T = { ∆ ( j ) | ∆ ( j ) ⊂ X , |∆ ( j ) | = d+ 1 } NT j=1 ( 2 ) such that ∪∆∈T CH ( ∆ ) equals the domain Ω where CH ( ∆ ) is the convex hull of simplex ∆ . So we define a simplex ∆ ( j ) ∈ T representing the j-th mesh cell as the set of vertices of the cell and denote the domain volume covered by the cell by the convex hull CH ( ∆ ( j ) ) of the vertices . We will assume u to be a scalar field , i.e . u : [ 0 , T ] ×Ω→ R. If u is a vector field , we treat it as a system of m scalar fields instead . For a detailed introduction to FEM , we refer the reader to Igel ( 2017 ) . Basis Functions . A priori , we assume that the unknown solution u to our problem lies in an infinitedimensional function space U . The first step in FEM to make the problem numerically feasible is to approximate U with a finite-dimensional linear subspace Ũ . This subspace can then be written in terms of linear combinations of basis functions Ũ = span { ϕ ( 1 ) , ... , ϕ ( N ) } . There are many possible bases and the choice determines various qualities of the resulting procedure such as continuity of the approximation and the sparsity pattern of the mass matrix in Eq . ( 7 ) . In our case , we choose the so-called P1 basis of piecewise linear functions ( hat functions ) , see Fig . 2a ( Igel , 2017 ) . There are as many basis functions as there are points and each is uniquely defined by being linear when restricted to a single cell ∆ ∈ T and the constraint ϕ ( j ) ( x ( i ) ) = { 1 if x ( i ) = x ( j ) 0 otherwise ∀x ( i ) ∈ X . ( 3 ) So the basis function ϕ ( j ) is 1 at x ( j ) , falls linearly to 0 on mesh cells adjacent to x ( j ) and is 0 everywhere else . The resulting finite-dimensional function space Ũ is the space of linear interpolators between values at the vertices , see Fig . 2b . An important property is that if we expand u ∈ Ũ in this basis , the value of u at the i-th node is just its i-th coefficient . u ( x ( i ) ) = ∑N j=1 cjϕ ( j ) ( x ( i ) ) = ci ( 4 ) Galerkin Method . A piecewise linear approximation u ∈ Ũ is not differentiable everywhere and therefore can not fulfill Eq . ( 1 ) exactly . So instead of requiring an exact solution , we ask that the residual R ( u ) = ∂tu−F ( t , x , u , ... ) be orthogonal to the approximation space Ũ with respect to the inner product 〈u , v〉Ω = ∫ Ω u ( x ) · v ( x ) dx at any fixed time t. In effect we are looking for the best possible solution within Ũ . Because Ũ is generated by a finite basis , the orthogonality requirement decomposes into N equations , one for each basis function . 〈R ( u ) , v〉Ω = 0 ∀v ∈ Ũ ⇐⇒ 〈R ( u ) , ϕ ( i ) 〉Ω = 0 ∀i = 1 , ... , N ( 5 ) Plugging the residual back in and using the linearity of the inner product , we can reconstruct a system of equations that resemble the PDE that we started with . 〈∂tu , ϕ ( i ) 〉 = 〈F ( t , x , u , ... ) , ϕ ( i ) 〉Ω ∀i = 1 , ... , N ( 6 ) At this point we can stack the system of N equations into a vector equation . If we plug in the basis expansion ∑N j=1 cjϕ ( j ) for u into the left hand side , we get a linear system A∂tc = m ( 7 ) where Aij = 〈ϕ ( i ) , ϕ ( j ) 〉Ω is the so called mass matrix , c is the vector of basis coefficients of u , and mi = 〈F ( t , x , u , ... ) , ϕ ( i ) 〉Ω captures the effect of the dynamics F . The left hand side evaluates to A∂tc , because the basis functions are constant with respect to time . The right hand side can not be further simplified without additional assumptions on F . Method of Lines . If we can evaluate the right hand sidem , we can solve the linear system in Eq . ( 7 ) for the temporal derivatives of the coefficients of u at each point in time . In fact we have converted the PDE into a system of ordinary differential equations ( ODEs ) which we can solve with an arbitrary ODE solver given an initial value c ( 0 ) as in Fig . 2c . This is known as the method of lines because we solve for u along parallel lines in time . To find a vector field u : [ 0 , T ] × Ω → Rm instead of a scalar field , we treat the m dimensions of u as a system of m scalar fields . This results in m copies of Eq . ( 7 ) , which we need to solve simultaneously . Because the mass matrixA is constant with respect to u , we can combine the system into a matrix equation A∂tC = M ( 8 ) where C , M ∈ RN×m are the stacked c and m vectors , respectively . In summary , the spatial discretization with finite elements allows us to turn the PDE ( 1 ) into the matrix ODE ( 8 ) . 2.2 MESSAGE PASSING NEURAL NETWORKS . Message-Passing Neural Networks ( MPNNs ) are a general framework for learning on graphs that encompass many variants of graph neural networks ( Gilmer et al. , 2017 ) . It prescribes that nodes in a graph iteratively exchange messages and update their state based on the received messages for P steps . For a graph G = ( V , E ) with nodes V and edges E , and initial node states h ( 0 ) v ∀v ∈ V , the p-th propagation step is h ( p ) v = fupd ( h ( p−1 ) v , ∑ { u , v } ∈E fmsg ( h ( p−1 ) u , h ( p−1 ) v ) ) , ( 9 ) where fmsg maps node states and edge attributes to messages and fupd updates a node ’ s state with the aggregated incoming messages . The final node states h ( P ) v can then be interpreted as per-node predictions directly or passed as node embeddings to downstream systems . In this work , we employ a slight generalization of the above to undirected hypergraphs , i.e . graphs where the edges are sets of an arbitrary number of nodes instead of having a cardinality of exactly 2 . For such a hypergraph G = ( V , E ) with nodes V and hyperedges ε = { u , v , w , ... } ∈ E , and initial node states h ( 0 ) v ∀v ∈ V , the p-th propagation step is h ( p ) v = fupd ( h ( p−1 ) v , ∑ ε∈E s.t.v∈ε fmsg ( { h ( p−1 ) u | u ∈ ε } ) v ) . ( 10 ) Note that fmsg jointly computes a separate message for each node v participating in a hyperedge ε .

Summary: This paper proposes a new model for learning partial differential equations from data. The PDE is first discretized then solved as an ODE. The dynamics function is learned with Message-Passing Neural Networks, where the function is split into a sum of physically informed terms. This splitting both improves model performance and makes the model more interpretable by disentangling the dynamics. The model is tested rigorously against multiple baseline models, and the results show the new model performs well.