L EARNING THE D YNAMICS OF P HYSICAL S YSTEMS FROM S PARSE O BSERVATIONS WITH F INITE E LEMENT N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
We propose a new method for spatio-temporal forecasting on arbitrarily distributed points
Assuming that the observed system follows an unknown partial differential equation, we derive a continuous-time model for the dynamics of the data via the ﬁnite element method.
The resulting graph neural network estimates the instantaneous effects of the unknown dynamics on each cell in a meshing of the spatial domain
Our model can incorporate prior knowledge via assumptions on the form of the unknown PDE, which induce a structural bias towards learning speciﬁc processes.
Through this mechanism, we derive a transport variant of our model from the convection equation and show that it improves the predictive accuracy on sea surface temperature and gas
ﬂow forecasting against baseline models representing a selection of spatio
-temporal forecasting methods.
A qualitative analysis shows that our model disentangles the data dynamics into their constituent parts, which makes it uniquely interpretable
I NTRODUCTION
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 (Coifﬁer, 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 ﬁll in these gaps with models trained on collected data.
For example, neural networks have been trained for weather forecasts (Shi et al., 2015) and ﬂuid ﬂow simulations (Belbute-Peres et al., 2020), both of which are traditionally outcomes of PDE solvers.
Even the dynamics of discrete dynamical systems such as trafﬁc (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 fulﬁll 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 Ω ⊂ R d that is governed by the PDE
[Equation]
on functions u : [0 , T ] × Ω → R m .
If we have a dense measurement u 0 : Ω → R m of the current state of the system and a solution u that satisﬁes Eq.
(1) for all t ∈ [0 , T ] and also fulﬁlls the initial condition u (0 , x ) =
u 0 ( 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 ﬁnd solutions of Eq.
(1) efﬁciently.
Unfortunately, in practice we only have a ﬁnite 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 speciﬁed dynamics is the ﬁnite element method (FEM) (Brenner et al., 2008).
In this paper, we introduce the ﬁrst 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 ﬂow ﬁeld, which can be visualized and analyzed.
In experiments on multi-step forecasting of sea surface temperature and gas ﬂow, our model improves upon baselines from recurrent, temporalconvolutional, and continuous-time model classes with further improvement by the transport model.
B ACKGROUND 2. F INITE E LEMENT M ETHOD
In the following, we will outline how to approximate a solution u to the dynamics in Eq.
(1) from an initial value u 0 by discretizing u in space using ﬁnite elements.
Let X be a set of points with a triangulation T of d -dimensional, non-overlapping simplices
[Equation]
such that ∪ ∆ ∈T CH(∆) equals the domain Ω where CH(∆) is the convex hull of simplex ∆ .
So we deﬁne 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 ﬁeld, i.e. u : [0 , T ] × Ω → R .
If u is a vector ﬁeld, we treat it as a system of m scalar ﬁelds 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 inﬁnitedimensional function space U .
The ﬁrst step in FEM to make the problem numerically feasible is to approximate U with a ﬁnite-dimensional linear subspace ˜ U .
This subspace can then be written in terms of linear combinations of basis functions ˜
U
= span (cid:8) ϕ
(1) , ..., ϕ ( N ) (cid:9) .
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.
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 deﬁned by being linear when restricted to a single cell ∆ ∈ T and the constraint
[Equation]
So the basis function ϕ ( j ) is 1 at x ( j ) , falls linearly to 0 on mesh cells adjacent to x ( j ) and iseverywhere else.
The resulting ﬁnite-dimensional function space ˜
U is the space of linear interpolators
[Figure]
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 coefﬁcient. [Equation]
[Equation]
Galerkin Method.
A piecewise linear approximation u ∈ ˜ U is not differentiable everywhere and therefore cannot fulﬁll Eq.
(1) exactly.
So instead of requiring an exact solution, we ask that the residual R ( u ) =
∂ t u − F ( t, x , u, ... ) be orthogonal to the approximation space ˜
U with respect to the inner product (cid:104) u, v (cid:105) Ω =
(cid:82) Ω u ( x ) · v ( x ) d x at any ﬁxed time t .
In effect we are looking for the best possible solution within ˜ U .
Because ˜ U is generated by a ﬁnite basis, the orthogonality requirement decomposes into N equations, one for each basis function.
[Equation]
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.
[Equation]
At this point we can stack the system of N equations into a vector equation.
If we plug in the basis expansion (cid:80) Nj =
c j ϕ ( j ) for u into the left hand side, we get a linear system
[Equation]
where A ij = (cid:104) ϕ ( i ) , ϕ ( j ) (cid:105) Ω is the so called mass matrix, c is the vector of basis coefﬁcients of u , and m i = (cid:104) F ( t, x , u, ... ) , ϕ ( i ) (cid:105)
Ω captures the effect of the dynamics F .
The left hand side evaluates to
A ∂ t c , because the basis functions are constant with respect to time.
The right hand side cannot be further simpliﬁed without additional assumptions on F .
Method of Lines.
If we can evaluate the right hand side m , we can solve the linear system in Eq.
(7) for the temporal derivatives of the coefﬁcients 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 ﬁnd a vector ﬁeld u : [0 , T ] × Ω → R m instead of a scalar ﬁeld, we treat the m dimensions of u as a system of m scalar ﬁelds.
This results in m copies of Eq.
(7), which we need to solve simultaneously.
Because the mass matrix A is constant with respect to u , we can combine the system into a matrix equation
[Equation]
[Equation]
where C , M ∈ R N × m are the stacked c and m vectors, respectively.
In summary, the spatial discretization with ﬁnite elements allows us to turn the PDE (1) into the matrix ODE (8).
2. M ESSAGE P ASSING N EURAL N ETWORKS
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
[Equation]
0) v ∀ v ∈ V , the p -th propagation step is
[Equation]
where f msg maps node states and edge attributes to messages and f upd updates a node’s state with the aggregated incoming messages.
The ﬁnal 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
[Equation]
Note that f msg jointly computes a separate message for each node v participating in a hyperedge ε .
F INITE E LEMENT N ETWORKS
Consider a set of nodes X
= { x ( i
) ∈ R d } Ni =1 representing N points in d -dimensional space.
At each node we measure m features y ( t 0 ,i ) ∈
R m .
Our goal is to predict y ( t j ) at future timesteps t j by solving PDE (1).
For it to be well-deﬁned, we need a continuous domain encompassing the nodes X .
By default, we construct such a domain by Delaunay triangulation.
The Delaunay algorithm triangulates the convex hull of the nodes into a set of mesh cells – d -dimensional simplices –
T = { ∆ ( i ) ⊂
X | | ∆ ( i ) | =
d +1 } N T i =1 which we represent as sets of nodes denoting the cell’s vertices.
The PDE will then be deﬁned on the domain Ω = ∪ ∆ ∈T CH(∆) .
One shortcoming of this algorithm is that it can produce highly acute, sliver-like cells on the boundaries of the domain if interior nodes are very close to the boundary of the convex hull.
We remove these slivers in a post-processing step, because they have small area and thus contribute negligible amounts to the domain volume, but can connect nodes which are far apart.
See Appendix F for our algorithm.
As an alternative to Delaunay triangulation, a mesh can be explicitly speciﬁed to convey complex geometries such as disconnected domains, holes in the domain and non-convex shapes in general.
As a ﬁrst step towards solving the PDE, we need to represent y ( t 0 ) as a function on Ω .
We deﬁne the function u 0 ( x ) k
= (cid:80) Ni =1 y ( t 0 )
k ϕ ( i ) ( x ) , such that its coefﬁcients as a vector in the P1 basis encode the data and evaluating it at the nodes u 0 ( x ( i ) )
=
y ( t 0 ,i ) agrees with the data by construction.
In Section 2.1, we have seen that solving PDE (1) approximately under an initial condition u ( t 0 , x ) =
u 0 ( x ) can be cast as an initial value ODE problem over the trajectory of the coefﬁcients of a solution in time.
Above, we encoded the features y ( t 0 ) in the coefﬁcients of u 0 , meaning that Eq.
(8) exactly describes the trajectory of the features through time.
The dynamics of the features are thus given by
[Equation]
(11) where Y ( t ) ik =
y ( t,i ) k is the feature matrix and A ij = (cid:104) ϕ
( i ) , ϕ ( j ) (cid:105) Ω is the mass matrix.
We call the matrix M on the right hand side the message matrix and it is given by
[Equation]
where u encodes the predicted features Y at time t in its coefﬁcients.
Our goal is to solve the feature dynamics (11) with an ODE solver forward in time to compute the trajectory of Y and thereby predict y ( t j ) .
To apply an ODE solver, we have to compute ∂ t Y at each time t which in turn consists of, ﬁrst, evaluating the right hand side and, second, solving the resulting sparse linear system.
Solving Eq.
(11) for ∂ t Y with the exact sparse mass matrix A carries a performance penalty on
GPUs, because sparse matrix solving is difﬁcult to parallelize.
To avoid this operation, we lump the mass matrix.
Lumping of the mass matrix is a standard approximation in FEM for time-dependent
PDEs to reduce computational effort (Lapidus & Pinder, 1999), which leads to good performance in practice.
[Equation]
We employ the direct variant that diagonalizes the mass matrix by row-wise summation, making the inversion of A trivial.
The remaining step in solving Eq.
(11) is evaluating the right hand side.
Let T i = (cid:8) ∆ ∈ T | x ( i ) ∈ ∆ (cid:9) be the set of all mesh cells that include node x ( i ) .
Then we can break the elements of M up into the contributions from individual cells adjacent to each node.
[Equation]
The sum terms capture the effect that the dynamics have on the state of the system such as convection and chemical reaction processes given the geometry of the mesh cell, the values at the nodes, and the basis function to integrate against.
However, the dynamics F are of course unknown and need to be learned from data.
Yet, if we were to introduce a deep model f θ ≈ F at this point, the model would be intolerably slow.
Each evaluation of the model would require numerical integration of a neural network over each mesh cell against multiple basis functions and adaptive-step ODE solvers can require hundreds of evaluations.
We can eliminate this computational effort by assuming that F is cell-wise constant
[Equation]
where I ∆ is the indicator function of cell ∆ .
We will omit the parameters of F ∆ for brevity from now on.
Though this seems like a strong assumption at ﬁrst, we still allow the dynamics to depend on the features, location, and time in arbitrary ways and the approximation becomes arbitrarily good as the mesh resolution increases.
Under this constancy assumption, the message matrix simpliﬁes to
[Equation]
Now we can introduce a deep model f θ ≈ F ∆ and avoid the costly integration in training and application because the coefﬁcients (cid:104) 1 , ϕ ( i ) (cid:105) CH(∆) can be precomputed once per domain.
3. M ODEL
With these theoretical foundations in place, we can formally deﬁne Finite Element Networks (FENs) .
Let G =
( X , T ) be the undirected hypergraph with nodes X and hyperedges T corresponding to the meshing of the domain, so each mesh cell ∆ ∈ T forms a hyperedge between its vertices.
Let f θ be a neural network that estimates the cell-wise dynamics F ∆ from time, cell location, cell shape and the values at the vertices
[Equation]
(17) where µ ∆ is the center of cell ∆ .
We call ˜ F (∆) the reaction term in reference to the convectiondiffusion-reaction model, because it is completely free-form and can in principle model anything such as reactions between features or sources and sinks.
f θ conditions on the cell centers µ ∆ and the difference vectors x
( i ) − µ ∆ separately to help the model distinguish between cell position and shape, and improve spatial generalization.
Deﬁne the following message and update functions in the
MPNN framework
[Equation]
where m ( i ) = (cid:80)
∆ ∈T f msg (∆) x ( i ) are the aggregated messages for node x ( i ) .
Performing one message passing step ( P = 1 ) in this model is exactly equivalent to solving for the feature derivatives
∂ t Y ( t ) in Eq.
(11) with the lumped mass matrix from Eq.
(13) and the simpliﬁed message matrix in Eq.
(16) with f θ as a model for F ∆ .
So making a forecast with FEN for T points in time t
= ( t 0 , t 1 , ..., t T ) based on a sparse measurement y ( t 0 ) means solving an ODE where at each evaluation step t e we run one message passing step in the MPNN described above to compute the current time derivatives of the features ∂ t y | t = t e .
[Equation]
Because FENs model the continuous-time dynamics of the data, they can also be trained on and make predictions at irregular timesteps.
By controlling the information available to f θ , we can equip the model with inductive biases.
If we omit the current time t , the learned dynamics become necessarily autonomous.
Similarly, if we do not pass the cell position µ ∆ , the model will learn stationary dynamics.
An orthogonal and at least as powerful mechanism to induce inductive biases into the model is through assumptions on the form of the unknown dynamics F .
Let us assume that the dynamics are not completely unknown, but that a domain expert told us that the generating dynamics of the data include a convective component.
So the dynamics for feature k are of the form
[Equation]
) where v ( k ) ( t, x , u, ... )
∈ R
d is the velocity ﬁeld of that feature and F (cid:48) represents the still unknown remainder of the dynamics.
The only way to reconcile the convective assumption with the basic assumption of cell-wise constancy of F , is to conclude that the velocity ﬁeld is likewise constant on each cell.
By following the same steps with the convection term
[Equation]
as for the unknown dynamics, we arrive at a specialized term for modeling convection.
Let g θ be a neural network that estimates the cell-wise per-attribute ﬂow ﬁeld
[Equation]
We derive the corresponding message function in Appendix A.
[Equation] [Equation]
With this we can extend FEN with specialized transport-modeling capabilities by adding the two message functions and deﬁne Transport-FEN (T-FEN) as a FEN with the message function
[Equation]
[Equation]
Consequently, T-FEN learns both a velocity ﬁeld to capture convection and a reaction term that ﬁts the unknown remainder of the dynamics.
Network Architecture.
We instantiate both f θ and g θ as multilayer perceptrons (MLPs) with tanh non-linearities.
The input arguments are concatenated after sorting the cell vertices by the angle of x ( i ) − µ ∆ in polar coordinates.
This improves generalization because f θ and g θ do not have to learn order invariance of the nodes.
Both MLPs have an input dimension of 1 + d + ( d + 1) · ( m + d ) in the non-autonomous, non-stationary conﬁguration since | ∆ | =
d +1 .
The reaction term f θ computes one coefﬁcient per cell and attribute and therefore has an output dimension of m whereas the transport term estimates a velocity vector for each cell and attribute resulting in an output dimension of md .
The number of hidden layers and other conﬁguration details for each experiment are listed in Appendix C.
We use scikit-fem (Gustafsson & McBain, 2020) to compute the inner products between basis functions over the mesh cells and the dopri5 solver in torchdiffeq (Chen et al., 2018) to solve the resulting ODE.
Training.
We train our models with the Adam optimizer (Kingma & Ba, 2015) by minimizing the multi-step prediction error in terms of the mean L 1 distance between predicted and actual node features
[Equation]
(24) where ˆ y ( t j , · ) is the prediction at time t j .
To accumulate gradients, we backpropagate through the operations of the ODE solver.
We observed that training our models to predict complete sequences from the beginning can get them stuck in local minima that take many epochs to escape.
To stabilize training, we begin by training on subsequences of length s = 3 and iteratively increase s by 1 per epoch until it reaches the maximum training sequence length.
E XPERIMENTS
Baselines.
We evaluate FEN and T-FEN against a selection of models representing a variety of temporal prediction mechanisms: Graph WaveNet (GWN) combines temporal and graph convolutions (Wu et al., 2019); Physics-aware Difference Graph Network (PA-DGN) estimates spatial derivatives as additional features for a recurrent graph network (Seo et al., 2020); the continuous-time MPNNbased model by (Iakovl et al., 2021) which we will call CT-MPNN, uses a general MPNN to learn the continuous-time dynamics of the data.
1 See Appendix C for the conﬁguration details.
GWN and PA-DGN use 12 and 5 timesteps respectively as input while the ODE-based models make their predictions based on a single timestep.
To ensure a fair comparison, we do not evaluate the models on the ﬁrst 12 timesteps of the test set.
Datasets.
For our experiments we have chosen one synthetic and two real-world datasets.
CylinderFlow contains simulated ﬂuid ﬂows around a cylinder along a channel simulated with the inviscid
Navier-Stokes equations (Pfaff et al., 2021).
The Black Sea dataset provides data on daily mean sea surface temperature and water ﬂow velocities on the Black Sea over several years.
ScalarFlow is a collection of 104 reconstructions of smoke plumes from 2 seconds long, multi-view camera recordings (Eckert et al., 2019).
In each recording a fog machine releases fog over a heating element in the center of the domain which then rises up with the hot air through convection and buoyancy.
All our experiments are run on hold-out datasets, i.e. the year 2019 of the Black Sea data and the last 20 recordings of ScalarFlow.
See Appendix B for a detailed description of the datasets, our sparse [Table]
Multi-step Forecasting.
The main objective of our models as well as the baselines are accurate forecasts and we train them to minimize the mean absolute error over 10 prediction steps for the real-world data and 30 steps for CylinderFlow.
We chose these steps because they are long enough that a model needs to learn non-trivial dynamics to perform well and more accurate dynamics lead to better performance.
At the same time they are short enough that not too much uncertainty accumulates.
For example, if we train any of the models on Black Sea on much longer training sequences, they learn only diffusive smoothing because the correlation between the temperature distribution today and in 30 days is basically nil, so that the only winning move is not too play – similar to how forecasting the weather for more than 10 days is basically impossible (Mailier, 2010).
For Table 1 we have evaluated the models on all possible subsequences of length 10 and 30 respectively of the respective test set of each dataset.
The results show that FEN and T-FEN provide the smallest prediction error on all three datasets with a further improvement due to the separate transport term in T-FEN on real-world data.
Additionally, T-FEN learns dynamics that are simpler to solve and require fewer function evaluations than the general FEN on the Black Sea and ScalarFlow datasets that include a temperature and smoke density component respectively that exhibit transport dynamics.
These results highlight that including prior knowledge in the model can enable increased performance in accuracy as well as efﬁciency when the additional assumptions are hold.
[Figure]
and the strongest baseline trained the
[Figure]
Extrapolation.
While we train the models on 10-step forecasting, it is instructive to compare their predictions on longer time horizons.
Fig.
3 shows the predictions of the three strongest models on a smoke plume after 50 steps.
Because these models are all ODE-based, they were conditioned on a single timestep and have been running for hundreds of solver steps at this point.
First, we notice that all three models have the fog rise at about the same, correct speed.
However, with CT-MPNN the fog disappears at the bottom because this model is by construction unable to learn non-stationary dynamics.
FEN avoids this problem but overall the predicted density of the rising fog is more uniform than we see in the data.
T-FEN with its separate transport term correctly predicts the increased density at the bottom over the fog inlet as well as the rise of a high-density region without too much smoothing.
See Appendix E for a comparison with all baselines at multiple time steps.
Super-Resolution.
Graph-based spatio-temporal forecasting models make predictions on all nodes jointly.
Training these models on many points or ﬁne meshes therefore requires a prohibitive amount of memory to store the data points and large backwards graphs for backpropagation.
One way around this problem is to train on coarse meshes and later evaluate the models on the full data.
In Fig.
4 we compare FEN and T-FEN against the strongest baseline on the Black Sea dataset.
While all three models increase in prediction error as the mesh resolution increases, FEN models deteriorate at a markedly slower rate than CT-MPNN.
We ascribe this to the fact that in FEN and T-FEN estimation of the dynamics and its effect on the nodes are to a certain extent separated.
The terms f θ and g θ estimate the dynamics on each cell, but the effect on the node values is then controlled by the inner products of basis functions which incorporate the shape of the mesh cells in a theoretically sound way.
In CT-MPNN on the other hand, the model learns these aspects jointly which makes it more difﬁcult to generalize to different mesh structures.
See Appendix D for the complete results on both datasets and a visual comparison of low and high resolution data.
Interpretability.
In FEN the reaction term f θ models the dynamics F as a black box and as such its activations are equally opaque as the dynamics we derived it from.
However, as we have seen in Section 3.1, by making assumptions on the structure of F we impose structure onto the model and this structure in turn assigns meaning to the model’s components.
For a T-FEN we see that the transport term corresponds to a velocity ﬁeld for each attribute, while the reaction term absorbs the remainder of the dynamics.
Fig.
5 shows that this disentanglement of dynamics works in practice.
In the ScalarFlow data we have a fog inlet at the bottom, i.e. a source of new energy, and from there the fog rises up by convection.
Plotting the contributions of f θ and g θ to ∂ t y separately shows that the model has learned that all fog is transported upwards and that there is a source at the bottom.
Instead of inspecting the contributions to ∂ t y , we can also investigate the learned parameter estimates directly.
Fig.
6 shows that the transport term learns a smooth velocity ﬁeld for the sea surface temperature on the Black Sea.
R ELATED W ORK
Spatio-Temporal Forecasting.
Various DNNs, and in particular GNNs, have been designed to model complex spatial and temporal correlations in data, with applications ranging from forecasting
[Figure] trafﬁc ﬂow (see (Jiang & Luo, 2021) and (Yin et al., 2021) for a comprehensive survey), to forecasting COVID-19 (Kapoor et al., 2020) and sea-surface temperature (de Bézenac et al., 2018). For a detailed review of both classical and deep learning approaches, which can be roughly categorized as based on point-cloud trajectories, regular grids, or irregular grids, see Shi & Yeung (2018).
As one representative of GNN-based methods, which tend to be SOTA, we compare FENs with GWN (Wu et al., 2019).
GWN uses a combination of graph convolution and 1D dilated temporal convolutions to learn spatio-temporal properties of trafﬁc networks.
Neural Networks and Differential Equations.
In 1998 Lagaris et al.
(1998) introduced the ﬁrst neural network approach to solve ODEs and PDEs.
They formulate the optimization goal as fulﬁlling the properties of the differential equation.
Raissi et al.
(2019) extend this idea and propose a framework to respect any given general nonlinear PDE for continuous- and discrete-time models.
Another neural approach to learn PDEs from data is PDE-Net (Long et al., 2018).
Its two main components are: learning the differential operators of PDEs through convolutional kernels and applying a neural network to approximate the nonlinear response functions.
Building on that work,
PDE-Net 2.0 (Long et al., 2019) uses a symbolic network to approximate the nonlinear response function.
de Bézenac et al.
(2018) derive a model from the convection-diffusion equation and estimate the ﬂow ﬁeld on a regular grid with a CNN.
Chen et al.
(2018) leverage ODE solvers in their neural architecture to create a model family of continuous depth called neural ODEs.
Ayed et al.
(2019) build upon neural ODEs to make spatio-temporal forecasts by learning the dynamics of all data points jointly with the structure of the domain encoded via the node positions as additional features.
GNODE (Poli et al., 2019) brings the framework of neural ODEs of Chen et al.
(2018) into the graph domain and provides continuous depth GNNs to make trafﬁc predictions.
In dynamical systems the variable depth component coincides with the time dimension which allows the network to work with observations at varying timesteps.
Iakovl et al.
(2021) extend GNODE with distance vectors between nodes as edge features and use it for PDE prediction.
Mitusch et al.
(2021) derive a model from FEM similar to our work but integrate the learned dynamics against the basis functions numerically and disregard the mesh to avoid designing the discretization into the model.
C ONCLUSION
We have introduced Finite Element Networks as a new spatio-temporal forecasting model based on the established FEM for the solution of PDEs.
In particular, we have shown how FENs can incorporate prior knowledge about the dynamics of the data into their structure through assumptions on the form of the underlying PDE and derived the specialized model T-FEN for data with convection dynamics.
Our experiments have shown that FENs learn the short-term dynamics of sea surface temperature and gas ﬂow more accurately than baseline models from three spatio-temporal forecasting model classes and generalize better to higher resolution meshes at test time.
The T-FEN variant provides even stronger performance and additionally allows meaningful introspection of the model.
Its structure is directly derived from the assumed form of the underlying PDE and as such its transport and reaction terms disentangle the dynamics into a convection component and the remainder such as inlets and sinks.
E THICS S TATEMENT
We propose a general model for spatio-temporal forecasting which can be applied in a variety of settings.
Among these are trafﬁc and crowd ﬂow predictions which have the potential for abuse.
However, due to the close connection to PDEs, we see it primarily as a model for physical systems in earth sciences and engineering, with impacts in the study of climate change and understanding of complex systems.
Finally, we can also imagine its application in modeling the spread of infectious diseases which also employs PDEs today (Viguerie et al., 2021).
R EPRODUCIBILITY
To maximize the reproducibility of our experimental results, we have ﬁxed seeds for every domain subsampling, training run and evaluation and recorded these in conﬁguration ﬁles, so that everything can be re-run.
We obtained all seed values from numpy.random.
SeedSequence to guarantee that the streams of random numbers from the random number generators (RNGs) are independent and have high sample quality.
While we ﬁxed all seeds, model training is still non-deterministic because the message passing step in graph-based models relies on a scatter operation.
Scattering is implemented via atomic operations on GPUs, which can re-order ﬂoating point operations such as additions between runs with the same seed inducing non-determinism (Fey & Lenssen, 2019).
Nonetheless, we have observed that multiple runs from the same seed only diverge slowly as training progresses and the variability in performance between runs from different seeds is small for all models anyway as can be seen in the standard deviation of the mean absolute error in Table 1.
We will provide the source code for our models as well as the baselines, experiment conﬁgurations and a pre-processed version of the raw Black Sea data to the reviewers through OpenReview at the beginning of the discussion period.
We have used an NVIDIA GeForce GTX 1080 Ti for all experiments and evaluations.
A D ERIVATION OF THE T RANSPORT T ERM
Let F be just a convection term.
Because of the linearity of the inner product (cid:104)· , ·(cid:105) Ω , we can ignore other additive components of F in this derivation.
[Equation]
We begin by plugging this into the message matrix.
[Equation] [Equation]
Next, we split the domain into the mesh cells.
[Equation]
Assuming the velocity ﬁeld to be cell-wise constant, we can pull it out of the inner product.
[Equation]
[Equation]
Now we expand u in the P1 basis (remember that we have encoded the data y in the coefﬁcients of u )
[Equation]
(29) and make use of the linearity of the inner product once more.
[Equation]
Next, we use Gauss’ theorem and a 0 -Neumann boundary condition to move the gradient over to ϕ ( i ) and absorb the negative sign.
[Equation]
Because the inner product is restricted to ∆ , we can restrict the inner sum to vertices of ∆ .
[Equation]
If we stack this expression over all features k and translate the result into the MPNN framework, we arrive at the transport term that we present as part of T-FEN in Section 3.1.
B D ATASETS
In this section, we give a brief overview over each of the datasets we used in Section 4 and how we pre-processed them.
Subsampling.
Both datasets that we use provide large amounts of dense data on a 2D or 3D grid, which we subsample to simulate sparse measurements of a physical system and make the amount of data per timestep manageable.
A non-trivial problem is the selection of sampling points and we decided to use the k-Medoids algorithm.
This way we ensure that we use actual data points as sampling points and reach a roughly uniform but still random cover of the domain.
A uniform covering is especially important for the Black Sea dataset because of the non-convex shape of the domain and small bays on the border that we still want to cover. [Figure]
B. C YLINDER F LOW
The CylinderFlow dataset is a collection of simulated ﬂows characterized by velocity and pressure ﬁelds released by Pfaff et al.
The domain is a two-dimensional channel with impenetrable boundaries at the top and bottom and a circular obstacle of varying size and position that blocks the ﬂow.
On the left-hand boundary, the velocity ﬁeld is ﬁxed to an inﬂow proﬁle that is constant over time but varies between samples.
Depending on the obstacle, inﬂow proﬁle, and ﬂow parameters the resulting velocity and pressure ﬁelds exhibit distinct behaviors, varying between stable ﬂows and oscillations.
Pfaff et al.
used the COMSOL solver to generate the data by solving the incompressible
Navier-Stokes equations.
See Fig.
7 for a visualization of the velocity and pressure ﬁelds as well as the simulation mesh.
The training set contains 1000 simulations, while validation and test set have 100 each.
Every sequence consists of 600 steps of the ﬂow sampled at a time resolution of ∆ t = 0 . 01s .
Spatially, the data is discretized on a mesh with 1843 nodes.
A special property of this dataset is that each sequence has a different mesh, because each mesh has a notch representing the cylinder obstacle which varies in size and position.
To enforce the boundary conditions, we ﬁx the velocity ﬁelds for all methods on the boundaries except for the outﬂow, similar to Pfaff et al..
On the boundaries of the channel as well as the cylinder boundary, we ﬁx the velocities to zero.
The velocity inﬂow proﬁle stays constant over time and we hold it ﬁxed during training and inference.
The pressure, however, is unconstrained throughout the domain and can also vary on the boundaries.
Due to the size of the dataset, we train only on subsequences starting at a multiple of 30 steps, in effect eliminating overlapping training sequences and reducing the size of the training dataset by a factor of 30.
On the test set, we still evaluate the models on all subsequences of length 30.
We chose to evaluate the models over 30 steps on this dataset as opposed to 10 steps on the other datasets, because the simulation data is so ﬁnely sampled in time that 10 steps are too few to see signiﬁcant activity.
B. B LACK S EA
The Black Sea data is a reanalysis dataset of daily and monthly mean ﬁelds in the Black Sea of variables such as temperature, ﬂow velocity, and salinity and 31 depth levels from 01/01/1992.
The whole dataset is available as the BLKSEA_MULTIYEAR_PHY_007_004 2 product from the EU
Copernicus Marine Service under a custom license 3 granting permission to freely use the data in research under the condition of crediting the Copernicus Marine Environment Monitoring Service.
The spatial resolution of the raw data is 1/27° × 1/36° on a regular mesh of size 395×215 covering the Black Sea.
Of the resulting 84,925 grid points, about 40,000 contain actual sea data and we subsample these valid points to create our sparse datasets.
We mesh the chosen sample points with the Delaunay algorithm and ﬁlter acute mesh cells on the boundary as described in Appendix F.
Additionally, we remove mesh cells that cover more land than water to avoid connecting nodes that are actually far apart within the domain of the water.
We use daily data from 01/01/2012 to 31/12/2017 for training, 01/01/2018 to 31/12/2018 for validation and 01/01/2019 to 31/12/2019 for testing.
As features we choose the zonal velocities and the water temperature.
Both features are taken at a depth of 12.54m because the dynamics at the surface are strongly inﬂuenced by wind, sun, and cloud cover, which are not part of the data and thus make the dynamics non-deterministic from the perspective of the model.
We normalize the node positions by their mean and average standard deviation in latitudinal and longitudinal direction.
This way we ensure that the aspect ratio stays constant and mesh cells are not deformed.
The velocities are normalized by their global mean and standard deviation over the training date range.
For the temperature, we normalize each node with its daily mean over the training date range and daily standard deviation of all nodes in the training date range.
We choose separate normalization schemes for the following reasons.
Assume that the ﬂow ﬁeld would be perfectly stable.
If we were to normalize node-wise, all normalized velocity features would be zero and no longer informative for the model.
For the temperature, we chose node-wise normalization to emphasize modeling the dynamics of the heat that is in excess of the daily mean temperature ﬁeld on each day.
This puts more emphasis on accurate modeling of individual heat pockets and helps distinguish models that can learn this from models that mostly work through diffusion such as PA-DGN.
B. S CALAR F LOW
ScalarFlow is a dataset of 104 smoke plumes published by Eckert et al.
under the Creative
Commons license CC-BY-NC-SA 4.0 (Eckert et al., 2019).
Eckert et al. created a controlled environment in which fog from a fog machine rises in an air column over a heating element.
They captured the resulting ﬂow on video with a calibrated multi-camera setup and reconstructed the velocity and density ﬁelds of each recording on a dense 100×178×100 grid for 150 timesteps with their proposed method.
At each grid point, the dataset provides 3-dimensional velocity vectors and the current fog density.
The resulting dataset represents high-resolution measurements of a dynamical physical system that is mostly driven by buoyancy and transport processes but also evokes turbulence.
Of the 104 recordings we use the ﬁrst 64 for training, the next 20 for validation and the ﬁnal 20 as a test set.
Before subsampling the grid points, we reduce the data to 2D by averaging over the depth (Kohl et al., 2020) and restrict the data to the central 60 points in x -direction and bottom 140 points in y -direction, because the fog does not reach points outside of that central box in almost any of the recordings.
We normalize the coordinates of the sample points by their mean and average of the variances in x and y direction.
This way we ensure that the aspect ratio of the domain does not change.
We do not normalize the features.
C M ODEL C ONFIGURATIONS C. FEN & T-FEN
Both f θ and g θ are MLPs with 4 hidden layers and tanh non-linearities.
In FEN the layers of f θ have a width of 128.
For T-FEN we chose to reduce the width of the hidden layers of both f θ and g θ to 96, so that both FEN and T-FEN have a comparable number of parameters and we can trace back any difference in performance to difference in model structure and not capacity.
See Table 2 for the parameter counts.
We solve the feature dynamics with a dopri5 solver with an absolute tolerance and relative tolerance of 10 − 6 on CylinderFlow and Black Sea, and 10 − 7 on ScalarFlow.
The models are autonomous and stationary on CylinderFlow; non-autonomous and non-stationary on Black Sea; and non-stationary but autonomous on ScalarFlow.
On Black Sea, the models learn time-dependent dynamics and represent the yearly periodicity of the data, we map the absolute timestamp into yearly time values where 0 represents the 1st of January and 1 the 31st of December.
On initialization, we set the last weights and bias of the last layers of both MLPs to 0 to start with the constant dynamics as a reasonable initialization for training.
For parameter learning, we use the Adam optimizer with a learning rate of 10 − 3 (Kingma & Ba, 2015).
We always use a batch size of 6.
C. B ASELINES
[Table]
CT-MPNN.
Following Iakovl et al.
(2021), our CTMPNN implementation models both the message function and the update function as MLPs with tanh nonlinearities.
Contrary to the default conﬁguration used by Iakovl et al.
(2021), we increased the widths of the
MLPs from 60 to 128 and the message dimension from 40 to 128 to make the model capacity comparable to the other baselines.
Instead of an implicit Adams-Bashforth solver we used the same dopri5 solver as for FEN and
T-FEN with an absolute tolerance of 10 − 6 and a relative tolerance of 10 − 7 .
This both increased the models performance and avoided frequent convergence issues in the step function of the Adams-Bashforth solver.
Finally, we exchanged the RPROP optimizer with a standard Adam optimizer and we compute the gradients the same in FEN and T-FEN by backpropagation through the ODE solver.
The batch size is 6.
PA-DGN.
For PA-DGN we choose two graph network layers of two layers each with a hidden size of 64.
We condition the model on 5 timesteps.
To represent the structure of the domain, we pass the 4-nearest neighbor graph of the nodes to the model following authors of the method (Seo et al., 2020).
The batch size is 3 on both datasets.
GWN.
For GWN we use the improved conﬁguration described by Shleifer et al.
(2019) withconvolution channels, extra skip connections and four of GWN’s WaveNet inspired blocks with two layers each.
The model is conditioned on 12 time steps of data and predicts 10 steps ahead.
For the long-term extrapolation in Appendix E we apply the model auto-regressively to get predictions for more than 10 steps.
We do not learn an adaptive adjacency matrix because that comes with a quadratic memory cost in the number of nodes and Wu et al.
have shown that GWN only performs 1.2 % worse on METR-LA without the adaptive adjacency matrix.
For the computation of the forward and backward transition matrices, we pass in the adjacency matrix derived from the mesh of the domain where the edges are weighted by the Gaussian kernel
[Equation]
where σ is the standard deviation of (cid:13)(cid:13) x
( i ) − x ( j ) (cid:13)(cid:13) 2 .
The batch size is 6.
All baselines are non-autonomous on the Black Sea dataset and receive the same time information as
FEN and T-FEN.
D S UPER -R ESOLUTION
Inspecting the full results in Fig.
8 reveals two things.
While FEN based models perform best in super-resolution on both datasets, GWN and PA-DGN also only deteriorate slightly on ﬁner meshes.
However, this is offset by the fact that these models did not achieve such a good ﬁt to begin with and are still out-performed by FEN-based models on every resolution.
CT-MPNN achieved performance comparable to FEN and T-FEN on the original mesh but generalizes markedly worse to ﬁner meshes.
We suppose that this is due to the fact that CT-MPNN uses a general MPNN internally, while FEN and T-FEN have an inductive bias towards producing physically sensible predictions because of their structural connection to FEM and PDEs.
If we compare coarse and ﬁne subsamplings of the two datasets in Fig.
9 visually, it is apparent that quintupling the number of nodes on ScalarFlow increases the resolution but the data is still qualitatively the same.
On Black Sea on the other hand, the higher resolution plot reveals small-scale vortices that are lost at 1000 nodes.
We believe that this transition between two qualitatively different regimes is the reason why we observe a stronger deterioration on the Black Sea dataset than on
ScalarFlow in Fig.
[Figure]
[Figure]
[Figure]
In Fig.
10 we show 60-step extrapolations of the all models trained on 10-step forecasting.
The ﬁrst thing to note is that PA-DGN struggles with the directionality inherent in the ScalarFlow dataset.
While its prediction evolves over time, it exhibits a mostly diffusive dynamic.
From a visual standpoint, GWN produces the most realistic looking forecasts with ﬁne details and sharp contrasts, especially compared to the ODE-based models that beat GWN in short-term forecasting.
We see two reasons for that.
First, ODE-based models that learn the data dynamics directly, i.e. all three of FEN, T-FEN, and CT-MPNN, have to work with a tight information bottleneck.
Because these models carry information through time only in the features themselves and have no latent state, they can condition their prediction only on a single time step and this bottleneck makes it difﬁcult to conserve ﬁne details over time.
GWN on the other hand predicts 10 timesteps at a time and conditions each forecast on the past 12 timesteps.
In this way, GWN can extract more information from the input data and also carry that information forward 10 steps at a time whereas
ODE-based models have to conserve the information through roughly 100 solver steps to reach t 10 , see the number of function evaluations in Table 1.
Second, FEN, T-FEN, and CT-MPNN are ODE models and as such are biased towards smooth predictions over time and over long time frames this leads to a smoothing in space.
Yet, on short time frames, before smoothing sets in, these models are able to model these physical processes more accurately, because ODEs are natural models for these data.
For a discussion of the differences between the FEN, T-FEN, and CT-MPNN predictions, see Section 4. [Figure]
T
[Figure]
FEN
CT [Figure] MPNN
GWN
PADGN
[Figure]
F F ILTERING H IGHLY A CUTE C ELLS ON THE B OUNDARY
Delaunay triangulation of a set of points can produce very long cells on the boundary as in Fig.when several points close to the boundary are almost co-linear.
This can just occur in the data but it can also happen during generation of synthetic data.
If some boundary nodes in synthetic data are exactly co-linear such as in grid domains, they might lose the exact co-linearity due to the inexactness of ﬂoating point numbers and operations if we, for example, rotate the domain.
To ensure that these artifacts do not impact our models, we pre-process the triangulation with Algorithm 1 and an angle threshold of ε = 10 π 180 .
1: function F ILTER C ELLS (triangulation T , threshold ε > 0 , boundary faces B ) 2: for all ∆ ∈ T do (cid:46)
Count number of faces on the boundary for each cell 3: c ∆ ← |{ ˜∆ | ˜∆ ⊂ ∆ , | ˜∆ | =
d }| 4:
F ← { ∆ | ∆ ∈ T , c ∆ = 1 } 5: while |F| > 0 do 6: ∆ ← pop( F ) 7: if c ∆ (cid:54) = 1 then (cid:46) Check that ∆ is not at an edge or corner of the domain 8: continue 9: ( ˜∆ , ξ ) ← split(∆)
(cid:46) Split boundary face and interior node 10: if M IN B OUNDARY A NGLE ( ˜∆ , ξ ) <
ε then 11:
T ← T \ { ∆ } (cid:46)
Remove cell 12: for all ˆ∆ ⊂
˜∆ | | ˆ∆ | = d − 1 do (cid:46)
Add newly become boundary faces to B 13:
B ← B ∪ { ˆ∆ ∪ { ξ }} 14: for all ∆ (cid:48) ∈ T | ˆ∆ ⊂ ∆ (cid:48) , ξ ∈ ∆ (cid:48) do 15: c ∆ (cid:48) ← c ∆ (cid:48)
+16:
F ← F ∪ { ∆ (cid:48) } (cid:46) Add cells that are now boundary cells to F 17: return T 18: function M IN B OUNDARY
A NGLE ( ˜∆ , ξ ) 19:
B ← project ( ξ, ˜∆) (cid:46) Project ξ onto the hyperplane given by ˜∆ 20: γ ← π 21: for all A ∈ ˜∆ do 22: γ ← min( γ, angle ( AB, Aξ ) (cid:46) Find the angle between the line segments AB and Aξ 23: return γ
[Figure]