MAgNet: Mesh Agnostic Neural PDE Solver
Anonymous Author(s)
AffiliationAddress email
Abstract
The computational complexity of classical numerical methods for solving Partial
Differential Equations (PDE) scales significantly as the resolution increases.
Whenit comes to climate predictions, fine spatio-temporal resolutions are required toresolve all turbulent scales in the fluid simulations.
This makes the task of accurately resolving these scales computationally out of reach even with modernsupercomputers.
As a result, climate modelers solve these PDEs on grids that aretoo coarse (3km to 200km on each side), which hinders the accuracy and usefulnessof the predictions.
In this paper, we leverage the recent advances in Implicit Neural
Representations (INR) to design a novel architecture that predicts the spatiallycontinuous solution of a PDE given a spatial position query.
By augmentingcoordinate-based architectures with Graph Neural Networks (GNN), we enablezero-shot generalization to new non-uniform meshes and long-term predictionsup to 250 frames ahead that are physically consistent.
Our Mesh Agnostic Neural
PDE Solver (MAgNet) is able to make accurate predictions across a variety of PDEsimulation datasets and compares favorably with existing baselines.
Moreover, MAgNet generalizes well to different meshes and resolutions up to four times thosetrained.
Introduction
Partial Differential Equations (PDEs) describe the continuous evolution of multiple variables, e.g.over time and/or space.
They arise everywhere in physics, from quantum mechanics to heat transfer and have several engineering applications in fluid and solid mechanics.
However, most PDEs can’tbe solved analytically so it is necessary to resort to numerical methods.
Since the introduction ofcomputers, many numerical approximations were implemented and new fields emerged such as
Computational Fluid Mechanics (CFD) (Richardson and Lynch, 2007).
The most famous numericalapproximation scheme is the Finite Element Method (FEM) (Courant, 1943; Hrennikoff, 1941).
In the
FEM, the PDE is discretized along with its domain and the problem is transformed into solving a setof matrix equations.
However, the computational complexity scales significantly with the resolution.
For climate predictions, this number can be quite significant if the desired error is to be reached,which renders its use impractical.
In this paper, we propose to learn the continuous solutions for spatio-temporal PDEs.
Previousmethods focused on either generating fixed resolution predictions or generating arbitrary resolutionsolutions on a fixed grid.
PDE models based on Multi-Layer Perceptrons (MLPs) can generatesolutions at any point of the domain (Dissanayake and Phan-Thien, 1994; Lagaris et al., 1998; Raissiet al., 2017a).
However, without imposing a physics-motivated loss that constrains the predictions tofollow the smoothness bias resulting from the PDE, MLPs become less competitive than CNN-basedapproaches especially when the PDE solutions have high-frequency information (Rahaman et al.,2018).
We leverage the recent advances in Implicit Neural Representations ((Tancik et al., 2020), (Chen et al.,2020), (Jiang et al., 2020)) and propose a general purpose model that can not only learn solutions to a
PDE with a resolution it was trained on, but it can also perform zero-shot super-resolution on irregularmeshes.
The added advantage is that we propose a general framework where we can make predictionsgiven any spatial position query for both grid-based architectures like CNNs and graph-based onesable to handle sensors and predictions at arbitrary spatial positions.
Contributions Our main contributions are in the context of machine learning for approximatelybut efficiently solving PDEs and can be summarized as follows:
• We propose a framework that enables grid-based and graph-based architectures to generate
continuous-space PDE solutions given a spatial query at any position.
• We show experimentally that this approach can generalize to resolutions up to four timesthose seen during training in zero-shot super-resolution tasks.
Related Works
Current solvers can require a lot of computations to generate solutions on a fine spatio-temporal grid.
When it comes to climate predictions, General Circulation Models (GCM) are typically used to makeforecasts that span several decades over the whole planet (Phillips, 1956).
These GCMs use PDEs tomodel the climate in the atmosphere-ocean-land system and to solve these PDEs, classical numericalsolvers are used.
However, the quality of predictions is bottlenecked by the grid resolution that is inturn constrained by the available amount of computing power.
Deep learning has recently emergedas an alternative to these classical solvers in hopes of generating data-driven predictions faster andmaking approximations that do not just rely on lower resolution grids but also on the statisticalregularities that underlie the family of PDEs being considered.
Using deep learning also makes itpossible to combine the information in actual sensor data with the physical assumptions embedded inthe classical PDEs.
All of this would enable practitioners to increase the actual resolution further forthe same computational budget, which in turn improves the quality of the predictions.
Machine Learning for PDE solving Dissanayake and Phan-Thien (1994) published one of thefirst papers on PDE solving using neural networks.
They parameterized the solutions to the Poissonand heat transfer equations using an MLP and studied the evolution of the error with the mesh size.
Lagaris et al.
(1998) used MLPs for solving PDEs and ordinary differential equations.
They wrotethe solution as a sum of two components where the first term satisfies boundary conditions and isnot learnable, and the second is parameterized with an MLP and trained to satisfy the equations.
In
Raissi et al.
(2017a)
the authors also parameterized the solution to a PDE using an MLP that takescoordinates as input.
With the help of automatic differentiation, they calculate the PDE residual anduse its MSE loss along with an MSE loss on the boundary conditions.
In follow-up work, Raissi et al.(2017b) also learn the parameters of the PDE (e.g. Reynolds number for Navier-Stokes equations).
The recently introduced Neural Operators framework (Kovachki et al., 2021; Li et al., 2020b,a)attempts to learn operators between spaces of functions.
Li et al.
(2021) use "Fourier Layers" tolearn the solution to a PDE by framing the problem as learning an operator from the space of initialconditions to the space of the PDE solutions.
Their model can learn the solution to PDEs that lie on auniform grid while maintaining their performance in the zero-shot super-resolution setting.
In thesame spirit, Jiang et al.
(2020) developed a model based on Implicit Neural Representations called"MeshFreeFlowNet" where they upsample existing PDE solutions to a higher resolution.
They use3D low-resolution space-time tensors as inputs to a 3DUnet in order to generate a feature map.
Next,some points are sampled uniformly from the corresponding high-resolution tensors and fed to an
MLP called ImNet (Chen and Zhang, 2018).
They train their model using a PDE residual loss andare able to predict the flow field at any spatio-temporal coordinate.
Their approach is closest to theone we propose here.
The main difference is that we perform super-resolution on the spatial queriesand forecast the solution to a PDE instead of only doing super-resolution on the existing sequence.
Brandstetter et al.
(2022) use the message-passing paradigm ((Gilmer et al., 2017), (Watters et al.,2017), (Sanchez-Gonzalez et al., 2020)) to solve 1D PDEs.
They are able to beat state-of-the-art
Fourier Neural Operators (Li et al., 2021) and classical WENO5 solvers while introducing the"pushforward trick" that allows them to generate better long-term rollouts.
Moreover, they present an
[Figure]
added advantage over existing methods since they can learn PDE solutions at any mesh.
However,they are not able to generalize to different resolutions.
Most machine learning approaches require data from a simulator in order to learn the required PDEsolutions and that can be expensive depending on the PDE and the resolution.
Wandel et al.
(2020)alleviate that requirement by using a PDE loss.
Machine Learning for Turbulence
Modeling Recent years have known a surge in machinelearning-based models for modeling turbulence.
Since it is expensive to resolve all relevant scales,some methods were developed that only solve large scales explicitly and separately model sub-gridscales (SGS).
Recently, Novati et al.
(2021) used multi-agent reinforcement learning to learn thedissipation coefficient of the Smagorinsky SGS model (Smagorinsky, 1963) using as reward therecovery of the statistical properties of Direct Numerical Simulations (DNS).
Rasp et al.
(2018) used
MLPs to represent sub-grid processes in clouds and replace previous parametrization models in aglobal general circulation model.
In the same fashion, Park and Choi (2021) used MLPs to learn
DNS sub-grid scale (SGS) stresses using as input filtered flow variables in a turbulent channel flow.
Brenowitz and Bretherton (2018) use MLPs to predicts the apparent sources of heat and moistureusing coarse-grained data and use a multi-step loss to optimize their model.
Wang et al.
usedone-layer CNNs to learn the spatial filter in LES methods and the temporal filter in RANS as wellas the turbulent terms.
A UNet (Ronneberger et al., 2015) is then used as a decoder to get the flowvelocity.
de Bezenac et al.
(2017) predict future frames by deforming the input sequence according tothe advection-diffusion equation.
It yields good results for Sea-Surface Temperature predictions.
Stachenfeld et al.
(2021) use the "encode-process-decode" (Sanchez-Gonzalez et al., 2018, 2020)paradigm along with dilated convolutional networks to capture turbulent dynamics seen in highresolution solutions only by training on low spatial and temporal resolutions.
Their approach beatsexisting neural PDE solvers in addition to the state-of-the-art Athena++ engine (Stone et al., 2020).
Methodology
We present the developed framework that leverages recent advances in Implicit Neural Representations(INR)
(Jiang et al., 2020; Sitzmann et al., 2020; Chen et al., 2020; Tancik et al., 2020) and drawsinspiration from mesh-free methods for PDE solving.
We first start by giving a mathematicaldefinition of a PDE.
Next, we showcase the proposed "MAgNet" and derive two variants: A gridbased architecture and a graph-based one.
3.1 Preliminaries
We define PDE as follows, using D k to denote k -th order derivatives:
Definition 3.1.
Evans (2010) Let U denote an open subset of R n and k ≥ 1 an integer.
An expressionof the form :
[Equation]
is called a k -th order system of PDEs,
where L : R mn k
×
R mn k − 1 × · · · × R mn × R m × U
→ R mis given and u : U → R m , u = ( u 1 , . .
, u m ) is the unknown function to be characterized .
In this paper, we are interested in spatio-temporal PDEs.
In this class of PDEs, the domain is
U =
[0 , + ∞ ] × S (time × space) where S ⊂ R n , n ≥ 1 and, with D k indicating differentiationwrt x , any such PDE can be formulated as:
[Equation]
Numerical PDE simulations have enjoyed a great body of innovations especially where their use isparamount in industrial applications and research.
Mesh-based methods like the FEM numericallycompute the PDE solution on a predefined mesh.
However, when there are regions in the PDE domainthat present large discontinuities, the mesh needs to be modified and provided with many more pointsaround that region in order to obtain acceptable approximations.
Mesh-based methods typically solvethis problem by re-meshing in what is called Adaptive Mesh Refinement (Berger and Oliger, 1984;
Berger and Colella, 1989).
However, this process can be quite expensive, which is why mesh-freemethods have become an attractive option that goes around these limitations.
3.2 MAgNet: Mesh-Agnostic Neural PDE Solver
3.2.1 "Encode-Interpolate-Forecast" framework
Let { x 1 , x 2 , . . .
, x T } ∈ R C × N denote a sequence of T frames that represents the ground-truthdata coming from a PDE simulator or real-world observations.
C denotes the number of physicalchannels, that is the number of physical variables involved in the PDE and N is the number of pointsin the mesh.
These frames are defined on the same mesh, that is the mesh does not change in time.
We call that mesh the parent mesh and denote its normalized coordinates of dimensionality n by{ p i } 1 ≤ i ≤ N ∈
[ − 1 , 1] n .
Let { c i } 1 ≤ i ≤ M ∈ [ − 1 , 1] n denote a set of M coordinates representing thespatial queries.
The task is to predict the solution for subsequent time steps both at: (i) all coordinatesfrom the parent mesh { p i } 1 ≤ i ≤ N , and (ii) coordinates from the spatial queries { c i } 1 ≤ i ≤ M .
At testtime, the model can be queried at any spatially continuous coordinate within the PDE domain toprovide an estimate of the PDE solution at those coordinates.
To perform the prediction, we first estimate the PDE solutions at the spatial queries for the first
T frames and then use that to forecast the PDE solutions at the subsequent timesteps at the querylocations.
We do this through three stages (see Figure 1):1.
Encoding :
The encoder takes as input the given PDE solution { x t } 1 ≤ t ≤ T at each point ofthe parent mesh { p i } 1 ≤ i ≤ N and generates a state-representation of original frames, whichcan be referred to as embeddings, and which we note { z t } 1 ≤ t ≤ T .
This representation willbe used in the interpolation step to find the PDE solution at the spatial queries { c i } 1 ≤ i ≤ M .
Note that in this encoding step, we can generate one embedding for each frame such that wehave T embeddings or summarize all the information in the T frames into one embedding.
We will explain the methodology using T embeddings, as it is easier to grasp the timedimension in this formulation, but the implementation has been done using a summarizedsingle embedding, as mentioned in section 3.2.2.
We also note that the embedded meshremains the same, i.e. we don’t change it by upsampling or downsampling it.2.
Interpolation :
We follow the same approach as Jiang et al.
and
Chen et al.
(2020)by performing an interpolation in the feature space.
Note that in case we generate onerepresentation that summarizes all T frames into one, then z t =
z for t = 1 , . .
, T .
Let{ t k } 1 ≤ k ≤ T denote the timesteps at which the x t are generated.
For each spatial query c i , let N ( c i ) denote the nearest points in the parent mesh p j .
Wegenerate an interpolation of the features z k
[ c i ] at coordinates
c i and at timestep t k as follows:
[Equation]
Where z k [ p j ] and x k [ p j ] denote the embedding and input frame at position p j and timet k respectively.
Moreover, w j are interpolation weights and are positive and sum to one.
Weights are chosen such that points closer to the spatial query have a higher contribution tothe interpolated feature than points farther away from the spatial query.
The g θ is an MLP.
To get the PDE solution x k [ c i ] at coordinate c
i , we use a decoder d θ which is an MLP here:x k
[ c i ] =
d θ ( z k [ c i ]) .
In practice, the number of neighbors that we choose is 2 n where n isthe dimensionality of the coordinates.3.
Forecasting : Now that we generated the PDE solution at the spatial queries c i for all thepast frames, we forecast the PDE solution at future time points at both spatial queries andthe parent mesh coordinates.
Let G denote the Nearest-Neighbors Graph (NNG) that has asnodes all the N locations in the parent mesh (at original coordinates { p i } 1 ≤ i ≤ N ) as well asall the M query points (at locations { c i } 1 ≤ i ≤ M ), with edges that include only the nearestneighbors of each node among the N + M − 1 others.
This corresponds to a new meshrepresented by the graph G .
Let { c ′ i } 1 ≤ i ≤ M + N denote the corresponding new coordinates.
We generate the PDE solution for subsequent time steps on this graph auto-regressivelyusing a decoder ∆ θ as follows:
[Equation]
We train MAgNet by using two losses:
• Interpolation Loss : This loss makes sure that the interpolated points match the ground-truth
and is computed as follows:
[Equation]
Where ˆ x k [ c i ] denotes the interpolated values generated by the model at the spatial queries.
• Forecasting Loss This loss makes sure
that
the model predictions into the future areaccurate.
If H is the horizon of the predictions, then we can express the loss
as follows:
[Equation]
Where ˆ x k + T [ c ′ i ] denotes the forecasted values generated by the model at the graph G whichcombines both spatial queries and the parent mesh.
The final loss is then expressed as:
[Equation]
3.2.2 Implementation Details
In the previous section, we described the general MAgNet framework.
In this section, we present howwe build the inputs to MAgNet as well as the architectural choices for the encoding, interpolation andforecasting modules and suggest two main architectures: MAgNet[CNN] and MAgNet[GNN].
Data pre-processing : We first consider a mesh that contains N ′ points ( N ′ ≥ N ).
We randomlysample N points from the mesh to form the parent mesh.
During training, M spatial queries arerandomly sampled from the N ′ − N remaining points.
We tried multiple values of M (that is thenumber of training spatial queries) to assess its impact on the performance of the method within asensitivity study presented in in Section 4.3.
The data pre-processing is illustrated in Figure [-
 [Figure]-]2 .
MAgNet[CNN] : In this architecture, we follow Chen et al.
and adopt the EDSR architecture(Lim et al., 2017) as our CNN encoder.
We concatenate all frames { x t } 1 ≤ t ≤ T in the channeldimension and feed that to our encoder in order to generate a single representation z .
For theforecasting module, we use the same GNN as in (Sanchez-Gonzalez et al., 2020).
A key advantage ofthis architecture is that it effectively turns existing CNN architectures into mesh-agnostic ones byquerying them at any spatially continuous point of the PDE domain at test time. [Figure]
The distinction is that the number of " training spatial queries " can be less than the total number of spatial queries and we investigate the impact of this number in Section 4.3.
MAgNet[GNN] : This model is similar to MAgNet[CNN] except that instead of using a CNN asan encoder, we use a GNN: the same architecture as in the forecasting module but each architecturehaving its separate set of parameters.
This is better suited for encoding frames with irregular meshes.
Similarly to MAgNet[CNN], we generate a single representation z that summarizes all the informationfrom the frames { x t } 1 ≤ t ≤ T .
Results
In this section, we evaluate MAgNet’s performance against the following baselines:
Fourier Neural Operators (FNO) (Li et al., 2021) : Considered the state-of-the-art model inneural PDE solving, FNO casts the problem of PDE solving as learning an operator from the space ofinitial conditions to the space of the solutions.
It is able to learn PDE solutions that lie on a uniformgrid and can do zero-shot super resolution.
Message-Passing Neural PDE Solvers (MPNN) (Brandstetter et al., 2022) :
Graph Neural
Networks have been used to learn physical simulations with great success (Sanchez-Gonzalez et al.,2020).
Recently, they have been used to learn solutions to PDEs (Brandstetter et al., 2022; SanchezGonzalez et al., 2020).
MPNN-based GNNs coupled with an autoregressive strategy demonstrate asuperior performance to FNO and are able to make long rollouts with the help of the "pushforwardtrick" that only propagates gradient of the last computed frame.
PDE Datasets We use three of MPNN’s PDE simulations (Brandstetter et al., 2022) as our experimental testbed. In the same fashion, we are interested in the following family of PDEs: [Equation] Where J = 5 , L = 16 and coefficients sampled uniformly in A j ∈ [ − 0 .
5 , 0 . 5] , ω j ∈
[ − . 4 , − 0 .
4] ,‘ j ∈ 1 , 2 , 3 , ϕ j ∈ [0 , 2 π ) and periodic boundary conditions following Brandstetter et al. Bar-Sinai et al. We denote the temporal resolution as n t and and set it to n t = 250 for theentire study. During testing, all the models are fed a history of T = 25 frames and produce a rolloutof n t − T = 225 frames in the future.
Here, we present the three datasets we work with: • E1 : Burgers equation without diffusion α = 1 , β = 0 , γ = • E2 :
Burgers equation with variable diffusion α = 1 , β
= η, γ = 0 where η ∈ [0 , 0 .
2]• E3 :
Mixed scenario where α ∈ [0 , 3] , β ∈ [0 , 0 . 4] and γ ∈ [0 , 1] . All training sets contain 2048 simulations and test sets contain 128 simulations.
All models areevaluated using the Mean Absolute Error (MAE) on the rolled out predictions averaged across timeand space:
[Equation]
[Figure] Appendix
For all the subsequent sections, n x and n ′ x denote the training and testing set’s resolutionsrespectively.
The temporal resolution n t = 250remains unchanged for all experiments.
4.1 GeneralPerformance On Regular Meshes
In this section, we compare MAgNet’s performance on all three datasets.
All models aretrained on a resolution of n x = 50 and the PDEsolutions lie on a uniform grid.
We test zero-shotsuper-resolution on n ′ x ∈ { 40 , 50 , 100 , 200 } .
Results are summarized in Figure 3 and visualizations of the predictions can be found in Appendix A.1.
MAgNet[CNN] outperforms bothbaselines on both E1 and E2 datasets yet isslightly outperformed by FNO on E3 (Figure3c).
Nonetheless, MAgNet[CNN]’s predictiveperformance stays consistent up to n ′ x = 200while MPNN does not generalize well to resolutions not seen during training.
4.2 Zero-ShotSuper-Resolution on Irregular Meshes
In this section, we study how MAgNets compareagainst the other baselines when it comes tomaking predictions on irregular meshes.
In orderto do so, we take simulations from the uniformmesh E1 dataset with a resolution of 100 andrun the following steps:
Let n x ∈ { 30 , 50 , 70 }
For each simulation in the E1 dataset, randomly sample the same subset ofn x points: the mesh remains for eachsingle simulation in the E1 dataset.
The procedure is the same for the test set but weinstead take the original E1 test set at a startingresolution of 200 and generate four test sets withirregular meshes for n ′ x ∈ { 40 , 50 , 100 , 200 } .
This is different from the test set of the previoussection albeit considering the same resolutionsince this one has irregular meshes.
We summarize our findings in Table 1.
[Figure] [Figure] [Figure] [Figure]
[Figure] [Figure]
[Figure]
[Figure] [Figure]
[Figure] [Figure] [Figure] [Figure]
[Figure] [Figure] [Figure] [Figure]
[Figure] [Figure] [Figure] [Figure]
MAgNet[GNN] performs better than MAgNet[CNN] on irregular meshes which is expected since
GNN encoders are better suited for this task.
However, surprisingly, even though we use a CNNencoder for MAgNet[CNN], the performance seems to be better in most cases not only compared to
FNO but also MPNN which is a graph-based architecture.
This effectively shows that MAgNet can beused to turn existing CNN architectures into mesh-agnostic solvers.
This is particularly interesting formeteorological applications where one needs to make predictions at the sub-grid level (at a specificcoordinate) while only having access to measurements on a grid.
4.3 Ablation and Sensitivity studies
In this section, we study different architectural choices and the sensitivity of key parameters in
MAgNet.
Basic Interpolators vs Learned Interpolators We investigate the contribution of the interpolation module to the general predictive performance of MAgNet.
We compare the MAgNet[CNN]architecture against three ablated variants:
• KNN :
We use K-Nearest-Neighbors interpolation (Qi et al., 2017) on the original frames
directly to obtain the interpolated values at the spatial queries.
• Linear : We use Linear interpolation on the original frames directly to obtain the interpolatedvalues at the spatial queries.
• Cubic : We use Cubic interpolation on the original frames directly to obtain the interpolatedvalues at the spatial queries.
Everything else is kept the same.
The evaluation is done on the E1 dataset with regular meshesand a resolution of n x = 50 .
Performance is tested on E1 with regular meshes for test resolutionsn ′ x ∈ { 40 , 50 , 100 , 200 } .
Results are summarized in Figure 4a.
Effect of modeling interactions between the parent mesh and the spatial query In section3.2, we presented the developed method which can be used to generate solutions at any spatialquery through the "Encode-Interpolate-Forecast" framework.
This means that we are free to chooseany architecture for the three processes.
In this section we investigate the choice of the "Forecast"architecture on the predictive performance as well as zero-shot super-resolution capabilities.
Wecompare MAgNet[CNN] and a variant that uses LSTM with attention ((Hochreiter and Schmidhuber,1997; Bahdanau et al., 2015) on the spatial queries only.
Results are shown in Table 2.
We see thatleveraging the interaction between the coordinates from the spatial query and those in the parent meshenables the model to give predictions consistent to different resolutions as opposed to generating thesolutions at these queries with no interaction.
Impact of the number of point samples during training We study the impact of the number ofspatial queries used during training ( M ).
We train MAgNet[GNN] on the E1 dataset with a resolutionof n x = 50 on a uniform grid and test on the same resolution.
Our findings are summarized in
Figure4b.
Increasing the number of spatial queries increases the predictive performance as expected.
Moreover, having many queries also decreases the variance of the results.
When we have fewer points,
[Figure]
[Table] Interaction [Table] the random sampling can cause some of these points to be in regions that decrease the loss fasterthan other regions, hence, the model’s performance becomes sensitive to randomization.
However,this effect grows weaker as the number of queries increases since they would uniformly cover moreregions in the mesh.
Limitations and Future Work
In this paper we introduced a novel framework that we call MAgNet for solving PDEs on any mesh,possibly irregular.
We proposed two variants of the architecture which gave promising results onbenchmark datasets.
We were effectively able to beat graph-based and grid-based architectureseven when using the CNN variant of the proposed framework, therefore suggesting a novel wayof adapting existing CNN architectures to make predictions on any mesh.
A key limitation of ourwork however, is the significance of the learned interpolator.
Indeed, compared with a simple cubicinterpolation, the approach introduced here doesn’t seem to offer a significant advantage and weleave improvement regarding this point for future work.
Another improvement could be seen in theforecasting module.
For now, MAgNet forecasts using a first-order explicit time-stepping scheme thatis known to suffer from instability problems in numerical PDE and ODE solvers.
Learned solversseem to somehow circumvent this limitation even when using large time steps (Sanchez-Gonzalezet al., 2020; Brandstetter et al., 2022; Stachenfeld et al., 2021).
In a future work, we wish to exploreother time-stepping schemes such as the 4th order Runge-Kutta method (Runge, 1895; Kutta, 1901)which is commonly used for solving PDEs.
A Appendix
A.1 Visualizations
Figures 5 and 6 show the models’ predictions on each of the test set resolutions n ′ x ∈{ 40 , 50 , 100 , 200 } on the E1 dataset.
MAgNet[CNN] predictions visually match the ground-truth’swhile MPNN’s prediction degrade as the predictions are advanced in time.
The models shown hereare the ones trained on uniform meshes.
We present visualizations for each of the test resolutions n ′ x ∈ { 40 , 50 , 100 , 200 } .
The temporal resolution is fixed at n t = 250 .
The x axis represents space and y axis represents time.
The arrow of time is from top to bottom.
A.2 Training details
We train all models for 250 epochs and early stopping with a patience of 40 epochs.
All models aretrained using Adam Optimizer (Kingma and Ba, 2014) and the StepLR learning scheduler (Paszkeet al., 2019) which decays the learning rate by a factor k every N steps epochs.
All models weretrained on 5 random seeds ∈ { 5 , 10 , 21 , 42 , 2022 } .
We use Pytorch (Paszke et al., 2019) and PytorchLightning (Falcon et al., 2020) for our implementations and we summarize all model’s traininghyper-parameters in Table 3.
A.
Architectural details
A CNN]
Encoder Architecture We adapt the original EDSR (Lim et al., 2017) architecture to work on 1Dsignals instead of 2D and use 4 residual blocks with a hidden dimension of 128 . [Table]
Interpolation Module We use a 4 layers MLP with a hidden size of 64 followed by Layernorm forg θ .
For d θ , we use a 4 layer MLP with a hidden size of
Forecasting Module We use the same architecture as in Sanchez-Gonzalez et al.
Theencoder module uses a 4 layer MLP with a hidden size of 64 and the latent dimension to 32 .
We use5 message-passing steps (with the same parameters for the MLP), the decoder also has a 4 layer MLPwith hidden size of 64 . [Table] [GNN
[Table]
Encoder Architecture We use the same architecture as in Sanchez-Gonzalez et al.
but onlykeep the encoder and processor.
We also use 5 message-passing steps for the processor and use 4layer MLP with hidden size of 128 and a latent dimension of 128 .
Interpolation Module We use a 4 layers MLP with a hidden size of 128 followed by Layernormfor g θ .
For d θ , we use a 4 layer MLP with a hidden size of
Forecasting Module We use the same architecture as in Sanchez-Gonzalez et al.
Theencoder module uses a 4 layer MLP with a hidden size of 64 and the latent dimension set to 128 .
Weuse 5 message-passing steps for the processor (with the same parameters for the MLP), the decoderalso has a 4 layer MLP with hidden size of 128 .
[Table]
Brandstetter et al.
We use the same hyperparameters as in Brandstetter et al.
[Table] (Li et al., 2021)
We use 5 Fourier Layers and 12 modes in Fourier space.
We use a hidden channel size of 256 .