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.
Asan important example, climate predictions require fine spatio-temporal resolutionsto resolve all turbulent scales in the fluid simulations.
This makes the task ofaccurately resolving these scales computationally out of reach even with modernsupercomputers.
As a result, current numerical modelers solve PDEs on gridsthat are too coarse (3km to 200km on each side), which hinders the accuracy andusefulness of the predictions.
In this paper, we leverage the recent advances in
Implicit Neural Representations (INR) to design a novel architecture that predictsthe spatially continuous solution of a PDE given a spatial position query.
Byaugmenting coordinate-based architectures with Graph Neural Networks (GNN),we enable zero-shot generalization to new non-uniform meshes and long-termpredictions up to 250 frames ahead that are physically consistent.
Our Mesh
Agnostic Neural PDE Solver (MAgNet) is able to make accurate predictions acrossa variety of PDE simulation datasets and compares favorably with existing baselines.
Moreover, MAgNet generalizes well to different meshes and resolutions up to fourtimes those trained 1 .
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-temporalgrid.
For example, climate predictions typically use General Circulation Models (GCM) 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, which is a crucial capability of our method.
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).
We take inspiration of this approach but replace the process module by an interpolation module, toallow the model to capture spatial correlations between known points and new query points.
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]
Where ∂ S is the boundary of S , B is a non-linear operator enforcing boundary conditions on u andg : S
→ R m represents the initial condition constraints for the solution u .
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):
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 will
be 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 one representation 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]
[Equation] [Equation]
i
[Equation] [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 are
accurate.
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.2.
The data pre-processing is illustrated in Figure
[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.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.
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.
We evaluate all models on both 1D and 2D simulations (the datasets generated from 1D and 2D PDEsare presented in Appendix A.3).
All training sets in the 1D case contain 2048 simulations and testsets contain 128 simulations.
For the 2D case, training sets contain 1000 simulations and test setscontain 100 simulations.
All models are evaluated using the Mean Absolute Error (MAE) on therolled out predictions averaged across time and space:
[Equation]
Where n t is the total number of frames and T the number of frames input to the models.
[Figure]
We train models for 250 epochs with early stopping with a patience of 40 epochs.
See Appendix A.5and A.6 for more implementation details.
We present a short summary of our results:
• 1D Case : Our model MAgNet is able to robustly generalize to unseen meshes in the regular
case as compared to the SOTA models FNO and MPNN.
The performance is even better inthe irregular case.
The details of the results are given in the following section 4.1.• 2D Case : Our model MAgNet is less competitive than MPNN when it comes to generalizingto unseen meshes in the regular case.
However, in the irregular case, we are more competitiveespecially when trained on sparse meshes.
The results regarding the 2D PDEs are presentedin the Appendix, in section A.4.
4.1 1D Case
For all the subsequent sections, n x and n ′ x denote the training and testing set’s resolutions respectively.
The temporal resolution n t = 250 remains unchanged for all experiments.
4.1.1 General Performance and zero-shot super resolution On Regular Meshes
In this section, we compare MAgNet’s performance on all three datasets.
All models are trained on aresolution of n x = 50 and the PDE solutions lie on a uniform grid.
We test zero-shot super-resolutionon n ′ x ∈ { 40 , 50 , 100 , 200 } .
Results are summarized in Figure 3 and visualizations of the predictionscan be found in Appendix A.1.
MAgNet[CNN] outperforms both baselines on both E1 and E2datasets yet is slightly outperformed by FNO on E3 (Figure 3c).
Nonetheless, MAgNet[CNN]’spredictive performance stays consistent up to n ′ x = 200 while MPNN does not generalize well toresolutions not seen during training.
4.1.2 General Performance and zero-shot super resolution on Irregular Meshes
In this section, we study how MAgNets compare against the other baselines when it comes to makingpredictions on irregular meshes.
In order to do so, we take simulations from the uniform-mesh E1dataset with a resolution of 100 and run the following steps:
Let n x ∈ { 30 , 50 , 70 } .
For each simulation in the E1 dataset, randomly sample the same subset of n x points: themesh remains for each single simulation in the E1 dataset.
The procedure is the same for the test set but we instead take the original E1 test set at a startingresolution of 200 and generate four test sets with irregular meshes for n ′ x ∈ { 40 , 50 , 100 , 200 } .
Thisis different from the test set of the previous section albeit considering the same resolution since thisone has irregular meshes.
We summarize our findings in Table 1.
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 CNN
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
encoder 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.2 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.
[Table]
With Interaction
[Table]
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,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 architectures evenwhen using the CNN variant of the proposed framework, therefore suggesting a novel way of adaptingexisting CNN architectures to make predictions on any mesh.
The main added value of the proposedmethod is its very good performance on irregular meshes, including for super-resolution tasks, asit can be observed for the presented 1D and 2D experiments, when compared to SOTA methods.
Notably, it seems to perform best when handed the smallest amount of data to work with (with a64x64 training resolution), even in the condensed non-uniform case.
This is a very desirable propertyfor real-life data.
A limitation of our work however, is the significance of the learned interpolator.
Indeed, compared with a simple cubic interpolation, the approach introduced here doesn’t seem tooffer a significant advantage and we leave improvement regarding this point for future work.
Anotherimprovement could be seen in the forecasting module.
For now, MAgNet forecasts using a first-orderexplicit time-stepping scheme that is known to suffer from instability problems in numerical PDEand ODE solvers.
Learned solvers seem to somehow circumvent this limitation even when usinglarge time steps (Sanchez-Gonzalez et al., 2020; Brandstetter et al., 2022; Stachenfeld et al., 2021).
In a future work, we wish to explore other time-stepping schemes such as the 4th order Runge-Kuttamethod (Runge, 1895; Kutta, 1901) which is commonly used for solving PDEs.
Finally, we observedthat our MAgNet model appears to perform significantly better in a 1D than in a 2D setting (asshown in Appendix A.4 2D results section).
While it is retaining its superiority in an irregular meshsetting, it appears a little less performant than the MPNN method, even with its GNN variant.
Furtherresearch on why the regular setting poses issues is thus also reserved for future work.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[No] (d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main exper-
imental results (either in the supplemental material or as a URL)?
[Yes] Link to the
code and dataset is provided in the first footnot.
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes] Training and implementation details can be found in the Appendix
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] We ran models for 5 random seeds.
(d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the
supplemental material or as a URL?
[No]
We are cleaning the github repository for public use.
It will include the final code andlink to our datasets.
d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A]
b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
A Appendix
[Figure]
A.1 Visualizations: Models’ predictions in the 1D case
[Figure]
Figures 5 and 6 show the 1D 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 space and y axis time of time is
[Figure]
Figures 7, 8, 9 and 10 show the 2D models’ predictions on each of the methods on the B1 dataset,trained on regular meshes, for four different test resolutions.
The models are all trained on a 64 × 64grid.
Visually MPNN and MAgNet[GNN] predictions appear rather close, and both matchingrelatively well the ground-truth’s.
[Figure]
Figures 11, 12, 13 and 14 show the 2D models’ predictions on the B1 dataset for the Uniform settingon resolutions n test ∈ { 32 2 , 64 2 , 128 2 , 256 2 } respectively (see Section A.4.2):
A.3 PDE Datasets
1D PDE Simulations For the 1D case, We use three of MPNN’s PDE simulations (Brandstetteret al., 2022) as our experimental testbed.
In the same fashion, we are interested in the followingfamily of PDEs:
[Equation]
Where J = 5 , L = 16 and coefficients sampled uniformly in A j ∈ [ − 0 .
5 , 0 . 5] , ω j ∈
[ − . 4 , − 0 . 4] ,l j ∈ { 1 , 2 , 3 } , ϕ j ∈ [0 , 2 π ) and periodic boundary conditions following Brandstetter et al.
Bar-Sinai et al.
For the 1D case, the temporal resolution is set to n t = 250 for the entirestudy.
During testing, all the models are fed a history of T = 25 frames and produce a rollout ofn t − T = 225 frames in the future.
Here, we present the two datasets we work with:
• E1 :
Burgers equation without diffusion α = 1 , β = 0 , γ =
• E2 :
Burgers equation with variable diffusion α = 1 , β
= η, γ = 0 where η ∈ [0 , 0 . 2]
[Figure]
• E3 :
Mixed scenario where α ∈ [0 , 3] , β ∈ [0 , 0 . 4] and γ ∈ [0 , 1] .
2D PDE Simulations For the 2D case, we use the Burgers equation as our experimental testbed.
This dataset was generated using Phiflow 2 :
[Equation]
Where J = 5 , L = 64 and coefficients sampled uniformly in A j ∈ [ − 0 .
5] , l xj , l yj ∈ { 1 , 2 , 3 } ,ϕ xj , ϕ y j ∈ [0 , 2 π ) and periodic boundary conditions.
For the 2D case, the temporal resolution is set ton t = 50 for the entire study.
During testing, all the models are fed a history of T = 10 frames andproduce a rollout of n t − T = 40 frames in the future.
Here, we present the three datasets we workwith:
The models are trained on a 64 × 64 grid, and tested on a 32 × 32 resolution.
The shown time steps are specified on top of the figure.
• B1 :
Burgers equation without diffusion β = 0• B2 :
Burgers equation with variable diffusion β ∈ (0 , 0 . 2]
A.4 Results for the 2D case
In this section, we present results for the 2D PDE simulations.
We use the datasets B1 and B2 as ourexperimental testbeds (see A.3).
We use n train to denote the train resolution and n test for the testresolution.
All models are fed a history of T = 10 frames and are required to generate a rollout ofn t − T = 40 frames in the future.
A.4.1 General performance and zero-shot super resolution on regular meshes
In this section, we compare MAgNet’s performance on both B1 and B2 datasets.
All models weretrained on a resolution of n train = 64 2 and the PDE solutions lie on a uniform grid.
Zero-shot superresolution is tested on n test ∈ { 32 2 , 64 2 , 128 2 , 256 2 } .
We summarize our findings in Figure 15.
Wenotice that MAgNet[CNN] falls behind when it comes to making good predictions and zero-shotsuper-resolution on both datasets while MAgNet[GNN] and MPNN take the lead.
Suprisingly, FNO suffers when it comes to generalizing to unseen resolutions.
Indeed, leveraging interactions betweenpoints after the interpolation module allows MAgNet[GNN] not only to make good predictions butalso to generalize to denser and unseen resolutions.
Visualizations of the prediction for a sample for the B1 dataset for all test resolutions can be found in
Figures 7, 8, 9 and 10 respectively.
A.4.2 General performance and zero-shot super resolution on irregular meshes
In this section, we present results for irregular meshes.
We consider two settings:
• Uniform: N nodes are randomly and uniformly sampled from a regular grid of resolution
• Condensed: N nodes are randomly sampled in a non-uniform way from a regular grid of
resolution 32 2 following the distribution: p ( x, y )
∝ exp (cid:0) − 8 (cid:0) ( x − 0 .
+ ( y − 0 . 25) 2 (cid:1)(cid:1)
We train MPNN and MAgNet[GNN] on four resolutions N ∈ { 64 , 128 , 256 , 512 } where again N isthe number of nodes in the mesh for both Uniform and Condensed settings.
For each training setting,
The models are trained on a 64 × 64 grid, and tested on a 128 × 128 resolution.
The shown time steps are specified on top of the figure.
we evaluate both models on regular grid PDE simulations from the B1 dataset.
Figure 16 shows themesh nodes for the Uniform setting and Figure 17 is for the Condensed setting.
Our findings are summarized in Figure 18.
We notice that our model MAgNet[GNN] is especiallymore appealing in case we have fewer points to work with.
In the Condensed setting which is morerealistic than the uniform one, MAgNet[GNN] is even more competitive than MPNN.
We notehowever, that while MAgNet[GNN]’s performance is better for small test resolutions, it quicklydeteriorates as we test on higher resolutions which is a main limitation of our approach in the 2Dcase.
A.5 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.
The models are trained on a 64 × 64 grid, and tested on a 256 × 256 resolution.
The shown time steps are specified on top of the figure.
A.6 Architectural details
[Figure] A.6.1 MAgNet[CNN]
[Figure] 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 .
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]
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]
6.3 MPNN Brandstetter et al.
We use the same hyperparameters as in Brandstetter et al.
A.6.4 FNO (Li et al., 2021)
We use 5 Fourier Layers and 12 modes in Fourier space.
We use a hidden channel size of 256 .
The shown time steps are specified on top of the figure
[Figure]
(Brandstetter et al., 2022; Gilmer et al., 2017)
[Figure]
MAgNet
] is especially appealing in however
MAgNet[
’s performance deteriorates resolutions which is a main limitation of our approach in