SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Baselines We compare our model to  $\bullet$  NodeMLP, a FeedForward Neural Network that ignores the node neighbors and only learns a local mapping  $\bullet$  GraphSAGE (Hamilton et al., 2017), a popular GNN architecture that uses SAGE convolutions  $\bullet$  MP-PDE (Brandstetter et al., 2022b), a message passing GNN that builds on (Pfaff et al., 2021) for solving PDEs.
Results. We show in Table 1 the performance on the test sets for the two datasets and for both evaluation settings. Overall, CORAL is on par with the best models for this task. For the Full setting, it is best on Cylinder and second on Airfoil behind MP-PDE. However, for the sparse protocol, it can infer the values on the full mesh with the lowest error compared to all other models. Note that this
second setting is more challenging for Cylinder than for Airfoil given their respective average mesh size. This suggests that the interpolation of the model outputs is more robust on the Airfoil dataset, and explains why the performance of NodeMLP remains stable between the two settings. While MP-PDE is close to CORAL in the sparse setting, GraphSAGE fails to generalize, obtaining worse predictions than the local model. This is because the model aggregates neighborhood information regardless of the distance between nodes, while MP-PDE does consider node distance and difference between features.
Table 1: Initial Value Problem - Test results. MSE on normalized data.
Model	Cylinder	Airfoil
Full	Sparse	Full	Sparse
NodeMLP	1.48e-1 ± 2.00e-3	2.29e-1 ± 3.06e-3	2.88e-1 ± 1.08e-2	2.83e-1 ± 2.12e-3
GraphSAGE	7.40e-2 ± 2.22e-3	2.66e-1 ± 5.03e-3	2.47e-1 ± 7.23e-3	5.55e-1 ± 5.54e-2
MP-PDE	8,72e-2 ± 4.65e-3	1.84e-1 ± 4.58e-3	1.97e-1 ± 1.34e-2	3.07e-1 ± 2.56e-2
CORAL	7.03e-2 ± 5.96e-3	1.70e-1 ± 2.53e-2	2.40e-1 ± 4.36e-3	2.43e-1 ± 4.14e-3
4.2 Dynamics Modeling
For the IVP problem, in section 4.1, the objective was to infer directly the state of the system at a given time  $T$  given an initial condition (IC). We can extend this idea to model the dynamics of a physical system over time, so as to forecast state values over a given horizon. We have developed an autoregressive approach operating on the latent code space for this problem. Let us denote  $(u_0, u_{\delta t}, \dots, u_{L\delta t})$  a target sequence of observed functions of size  $L + 1$ . Our objective will be to predict the functions  $u_{k\delta t}, k = 1, \dots, L$ , starting from an initial condition  $u_0$ . For that we will encode  $z_0 = e(u_0)$ , then predict sequentially the latent codes  $z_{k\delta t}, k = 1, \dots, L$  using the processor in an auto regressive manner, and decode the successive values to get the predicted  $\hat{u}_{k\delta t}, k = 1, \dots, L$  at the successive time steps.
4.2.1 Implementation with Neural ODE
The autoregressive processor is implemented by a Neural ODE solver operating in the latent  $z$ -code space. Compared to the plain MLP implementation used for the IVP task, this provides both a natural autoregressive formulation, and overall, an increased flexibility by allowing to forecast at any time in a sequence, including different time steps or irregular time steps. Starting from any latent state  $z_{t}$ , a neural solver predicts state  $z_{t + \tau}$  as  $z_{t + \tau} = z_{t} + \int_{t}^{t + \tau}\zeta_{\psi}(z_{s})ds$  with  $\zeta_{\psi}$  a neural network with parameters to be learned, for any time step  $\tau$ . The autoregressive setting directly follows from this formulation. Starting from  $z_{0}$ , and specifying a series of forecast time steps  $k\delta t$  for  $k = 1,\dots,L$ , the solver call  $NODESolve(\zeta_{\psi},z_0,\{k\delta t\}_{k = 1,\dots,L})$  will compute predictions  $z_{k\delta t},k = 1,\dots,L$  autoregressively, i.e. using  $z_{k\delta t}$  as a starting point for computing  $z_{(k + 1)\delta t}$ . In our experiments we have used a fourth-order Runge-Kutta scheme (RK4) for solving the integral term. Using notations from Section 3.2, the predicted field at time step  $k$  can be obtained as  $\hat{u}_{k\delta t} = \xi \circ g_{\psi}^{k}(e(u_{0}))$  with  $g_{\psi}^{k}$  indicating  $k$  successive applications of processor  $g_{\psi}$ . Note that when solving the IVP problem from Section 4.1, two INRs are used, one for encoding the input function and one for the output function; here a single modulated INR  $f_{\theta ,\phi}$  is used to represent a physical quantity throughout the sequence at any time.  $\theta$  is then shared by all the elements of a sequence and  $\phi$  is computed by a hypernetwork to produce a function specific code.
We use the two-step training procedure from Section 3.4, i.e. first the INR is trained to auto-decode the states of each training trajectory, and then the processor operating over the codes is learned through a Neural ODE solver according to Equation (5). The two training steps are separated and the codes are kept fixed during the second step. This allows for a fast training as the Neural ODE solver operates on the low dimensional code embedding space.
$$
\underset {\psi} {\arg \min } \mathbb {E} _ {u \sim \nu_ {u}, t \sim \mathcal {U} (0, T ]} \| g _ {\psi} \left(\tilde {e} _ {u} \left(u _ {0}\right), t\right) - \tilde {e} _ {u} \left(u _ {t}\right) \| ^ {2} \tag {5}
$$
4.2.2 Experiment details
Datasets We consider two fluid dynamics equations for generating the datasets and refer the reader to Appendix A for additional details.  $\bullet$  2D-Navier-Stokes equation (Navier-Stokes) for a viscous, incompressible fluid in vorticity form on the unit torus:  $\frac{\partial w}{\partial t} + u \cdot \nabla w = \nu \Delta w + f$ ,  $\nabla u = 0$  for  $x \in \Omega, t > 0$ , where  $\nu = 10^{-3}$  is the viscosity coefficient. The train and test sets are composed of 256 and 16 trajectories respectively where we observe the vorticity field for 40 timestamps. The original spatial resolution is  $256 \times 256$  and we sub-sample the data to obtain frames of size  $64 \times 64$ .  $\bullet$  3D-Spherical Shallow-Water equation (Shallow-Water) can be used as an approximation to a flow on the earth's surface. The data consists of the vorticity  $w$ , and height  $h$  of the fluid. The train and test sets are composed respectively of 16 and 2 long trajectories, where we observe the vorticity and height fields for 160 timestamps. The original spatial resolution is  $128~(\mathrm{lat}) \times 256$  (long), which we sub-sample to obtain frames of shape  $64 \times 128$ . We model the dynamics with the complete state  $(h,w)$ . Each trajectory, for both both datasets and for train and test is generated from a different initial condition (IC).
Setting We evaluate the ability of the model to generalize in space and time. • Temporal extrapolation: For both datasets, we consider sub-trajectories of 40 timestamps that we split in two equal parts of size 20, with the first half denoted  $In - t$  and the second one  $Out - t$ . The training-  $In - t$  set is used to train the models at forecasting the horizon  $t = 1$  to  $t = 19$ . At test time, we unroll the dynamics from a new IC until  $t = 39$ . Evaluation in the horizon  $In - t$  assesses CORAL's capacity to forecast within the training horizon.  $Out - t$  allows evaluation beyond  $In - t$ , from  $t = 20$  to  $t = 39$ . • Varying sub-sampling: We randomly sub-sample  $\pi$  percent of a regular mesh to obtain the train grid  $\mathcal{X}_{tr}$ , and a second test grid  $\mathcal{X}_{te}$ , that are shared across trajectories. The train and test grids are different, but have the same level of sparsity. • Up-sampling: We also evaluate the up-sampling capabilities of CORAL in Appendix C. In these experiments, we trained the model on a sparse, low-resolution grid and evaluate its performance on high resolution-grids.
Baselines To assess the performance of CORAL, we implement several baselines: two operator learning models, one mesh-based network and one coordinate-based method.  $\cdot$  DeepONet (Lu et al., 2021): we train DeepONet in an auto-regressive manner with time removed from the trunk net's input.  $\cdot$  FNO (Li et al., 2021): we use an auto-regressive version of the Fourier Neural Operator.  $\cdot$  MP-PDE (Brandstetter et al., 2022b): we use MP-PDE as the irregular mesh-based baseline. We fix MP-PDE's temporal bundling to 1, and train the model with the push-forward trick.  $\cdot$  DINO (Yin et al., 2022): We finally compare CORAL with DINo, an INR-based model designed for dynamics modeling.
4.2.3 Results
Table 2: Temporal Extrapolation - Test results. Metrics in MSE.
Xtr↓ Xte	dataset →	Navier-Stokes	Shallow-Water
In-t	Out-t	In-t	Out-t
π = 100% regular grid	DeepONet	4.72e-2 ± 2.84e-2	9.58e-2 ± 1.83e-2	6.54e-3 ± 4.94e-4	8.93e-3 ± 9.42e-5
FNO	5.68e-4 ± 7.62e-5	8.95e-3 ± 1.50e-3	3.20e-5 ± 2.51e-5	1.17e-4 ± 3.01e-5
MP-PDE	4.39e-4 ± 8.78e-5	4.46e-3 ± 1.28e-3	9.37e-5 ± 5.56e-6	1.53e-3 ± 2.62e-4
DINo	1.27e-3 ± 2.22e-5	1.11e-2 ± 2.28e-3	4.48e-5 ± 2.74e-6	2.63e-3 ± 1.36e-4
CORAL	1.86e-4 ± 1.44e-5	1.02e-3 ± 8.62e-5	3.44e-6 ± 4.01e-7	4.82e-4 ± 5.16e-5
π = 20% irregular grid	DeepONet	8.37e-1 ± 2.07e-2	7.80e-1 ± 2.36e-2	1.05e-2 ± 5.01e-4	1.09e-2 ± 6.16e-4
FNO + lin. int.	3.97e-3 ± 8.03e-4	9.92e-3 ± 2.36e-3	n.a.	n.a.
MP-PDE	3.98e-2 ± 1.69e-2	1.31e-1 ± 5,34e-2	5.28e-3 ± 5.25e-4	2.56e-2 ± 8.23e-3
DINo	9.99e-4 ± 6.71e-3	8.27e-3 ± 5.61e-3	2.20e-3 ± 1.06e-4	4.94e-3 ± 1.92e-4
CORAL	2.18e-3 ± 6.88e-4	6.67e-3 ± 2.01e-3	1.41e-3 ± 1.39e-4	2.11e-3 ± 5.58e-5
π = 5% irregular grid	DeepONet	7.86e-1 ± 5.48e-2	7.48e-1 ± 2.76e-2	1.11e-2 ± 6.94e-4	1.12e-2 ± 7.79e-4
FNO + lin. int.	3.87e-2 ± 1.44e-2	5.19e-2 ± 1.10e-2	n.a.	n.a.
MP-PDE	1.92e-1 ± 9.27e-2	4.73e-1 ± 2.17e-1	1.10e-2 ± 4.23e-3	4.94e-2 ± 2.36e-2
DINo	8.65e-2 ± 1.16e-2	9.36e-2 ± 9.34e-3	1.22e-3 ± 2.05e-4	1.52e-2 ± 3.74e-4
CORAL	2.44e-2 ± 1.96e-2	4.57e-2 ± 1.78e-2	8.77e-3 ± 7.20e-4	1.29e-2 ± 1.92e-3
Table 2 details the performance of the different models in a combined temporal and spatial evaluation setting. • General remarks: CORAL demonstrates strong performance across all scenarios for both datasets. Only DINo exhibits similar properties, i.e., stability across spatial subsamplings and extrapolation horizon. We observe that all models perform degradation with lower sampling ratio. Also, as the models have been trained only on In-t horizon, error accumulates over time and thus leads to lower performance for Out-t evaluation. • Analysis per model: Although achieving strong performance on some specific scenarios, DeepONet, FNO and MP-PDE results are dependent of the training grid, geometries or number of points. FNO, can only be trained and evaluated on regular grids while DeepONet is not designed to be evaluated on a different grid in the branch net. MP-PDE achieves strong performance with enough sample positions, e.g. full grids here, but struggles to compete on irregular grids scenarios in Navier-Stokes. • Inference Time: We report in Appendix C, the inference time of the baselines considered. Despite operator methods have better inference time, CORAL is faster than mesh-free methods like DINo and MP-PDE. • Generalization across samplings: Coordinate-based methods demonstrate robustness when it comes to changes in spatial resolution. In contrast, MP-PDE model exhibits strong overfitting to the training grid, resulting in a decline in performance. Although MP-PDE and DINo may outperform CORAL in some settings, when changing the grid, CORAL remains stable and outperforms the other models. See Appendix C for details.
4.3 Geometry-aware inference
In this section, we wish to infer the steady state of a system from its domain geometry, all other parameters being equal. The domain geometry is partially observed from the data in the form of point clouds or of a structured mesh  $\mathcal{X}_i\subset \Omega_i$ . The position of the nodes depends on the particular object shape. Each mesh  $\mathcal{X}_i$  is obtained by deforming a reference grid  $\mathcal{X}$  to adjust to the shape of the sample object. This grid deformation is the input function of the operator learning setting, while the output function is the physical quantity  $u_{i}$  over the domain  $\Omega_{i}$ . The task objective is to train a model so as to generalize to new geometries, e.g. a new airfoil shape. Once a surrogate model has been trained to learn the influence of the domain geometry on the steady state solution, it can be used to quickly evaluate a new design and to solve inverse design problems (details in Appendix D).
Datasets. We used datasets generated from three different equations by Li et al. (2022a) and provide more details in Appendix A. • Euler equation (NACA-Euler) for a transonic flow over a NACA-airfoil. The measured quantity at each node is the Mach number. • Navier-Stokes Equation (Pipe) for an incompressible flow in a pipe, expressed in velocity form. The measured quantity at each node is the horizontal velocity. • Hyper-elastic material (Elasticity). Each sample represents a solid body with a void in the center of arbitrary shape, on which a tension is applied at the top. The material is the incompressible Rivlin-Saunders material and the measured quantity is the stress value. We use 1000 samples for training and 200 for test with all datasets.
Baselines We use  $\cdot$  Geo-FNO (Li et al., 2022a) and  $\cdot$  Factorized-FNO (Tran et al., 2023) two SOTA models as the main baselines. We also compare our model to regular-grid methods such as  $\cdot$  FNO (Li et al., 2021) and  $\cdot$  UNet (Ronneberger et al., 2015), for which we first interpolate the input.
Table 3: Geometric aware inference - Test results. Relative L2 error.
Model	NACA-Euler	Elasticity	Pipe
FNO	3.85e-2 ± 3.15e-3	4.95e-2 ± 1.21e-3	1.53e-2 ± 8.19e-3
UNet	5.05e-2 ± 1.25e-3	5.34e-2 ± 2.89e-4	2.98e-2 ± 1.08e-2
Geo-FNO	1.58e-2 ± 1.77e-3	3.41e-2 ± 1.93e-2	6.59e-3 ± 4.67e-4
Factorized-FNO	6.20e-3 ± 3.00e-4	1.96e-2 ± 2.00e-2	7.33e-3 ± 4.66e-4
CORAL	5.90e-3 ± 1.00e-4	1.67e-2 ± 4.18e-4	1,20e-2 ± 8.74e-4
Results In Table 3 we can see that CORAL achieves state-of-the-art results on Airfoil and Elasticity, with the lowest relative error among all models. It is slightly below Factorized-FNO and Geo-FNO on Pipe. One possible cause is that this dataset exhibits high frequency only along the vertical dimension, while SIREN might be better suited for isotropic frequencies. Through additional experiments, we demonstrate in Appendix D, how CORAL can also be used for solving an inverse problem
corresponding to a design task: optimize the airfoil geometry to minimize the drag over lift ratio. This additional task further highlights the versatility of this model.
Although a versatile model, CORAL inherits the limitations of INRs concerning the training time and representation power. It is then faster to train than GNNs, but slower than operators such as FNO, DeepONet and of course CNNs which might limit large scale deployments. Also some physical phenomena might not be represented via INRs. Although this is beyond the scope of this paper, it remains to evaluate the methods on large size practical problems. An interesting direction for future work would be to derive an efficient spatial latent representation for INRs, taking inspiration from grid-based representation for INRs (Takikawa et al. (2022), Müller et al. (2022), Saragadam et al. (2022)). Another avenue would be to leverage clifford layers (Brandstetter et al., 2022a) to model interactions between physical fields.
We have presented CORAL, a novel approach for Operator Learning that removes constraints on the input-output mesh. CORAL offers the flexibility to handle spatial sampling or geometry variations, making it applicable to a wide range of scenarios. Through comprehensive evaluations on diverse tasks, we have demonstrated that it consistently achieves state-of-the-art or competitive results compared to baseline methods. By leveraging compact latent codes to represent functions, it enables efficient and fast inference within a condensed representation space.
We acknowledge the financial support provided by DL4CLIM (ANR-19-CHIA-0018-01) and DEEP-NUM (ANR-21-CE23-0017-02) ANR projects, as well as the program France 2030, project 22-PETA-0002. This project was provided with computer and storage resources by GENCI at IDRIS thanks to the grant AD011013522 on the supercomputer Jean Zay's V100 partition.
Kelsey R. Allen, Tatiana Lopez-Guevara, Kimberly Stachenfeld, Alvaro Sanchez-Gonzalez, Peter Battaglia, Jessica Hamrick, and Tobias Pfaff. Physical design using differentiable learned simulators. 2 2022. URL http://arxiv.org/abs/2202.00728.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. CoRR, abs/1506.03099, 2015. URL http://arxiv.org/abs/1506.03099.
Johannes Brandstetter, Rianne van den Berg, Max Welling, and Jayesh K Gupta. Clifford neural layers for pde modeling. In The Eleventh International Conference on Learning Representations, 2022a.
Johannes Brandstetter, Daniel E Worrall, and Max Welling. Message passing neural pde solvers. International Conference on Learning Representations, 2022b.
Keaton J. Burns, Geoffrey M. Vasil, Jeffrey S. Oishi, Daniel Lecoanet, and Benjamin P. Brown. Dedalus: A flexible framework for numerical simulations with spectral methods. Physical Review Research, 2, 2020. ISSN 26431564. doi: 10.1103/PhysRevResearch.2.023068.
Shengze Cai, Zhiping Mao, Zhicheng Wang, Minglang Yin, and George Em Karniadakis. Physics-informed neural networks (pinns) for fluid mechanics: a review. Acta Mechanica Sinica/Lixue Xuebao, 37, 2021. ISSN 16143116. doi: 10.1007/s10409-021-01148-1.
Sangwon Chae, Joonhyeok Shin, Sungjun Kwon, Sangmok Lee, Sungwon Kang, and Donghyun Lee. Pm10 and pm2.5 real-time prediction models using an interpolated convolutional neural network. Scientific Reports, 11, 2021. ISSN 20452322. doi: 10.1038/s41598-021-91253-9.
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with a general conditioning layer. 32nd AAAI Conference on Artificial Intelligence, AAAI 2018, 2018. ISSN 2159-5399. doi: 10.1609/aaai.v32i1.11671.
Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez, and Peter W. Battaglia. Learning mesh-based simulation with graph networks. International Conference on Learning Representations., 10 2021. URL http://arxiv.org/abs/2010.03409.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. volume 9351, 2015. doi: 10.1007/978-3-319-24574-4_28.
Vishwanath Saragadam, Jasper Tan, Guha Balakrishnan, Richard G. Baraniuk, and Ashok Veeraraghavan. Miner: Multiscale implicit neural representations, 2022.
Vincent Sitzmann, Eric R. Chan, Richard Tucker, Noah Snavely, and Gordon Wetzstein. Metasdf: Meta-learning signed distance functions. Advances in Neural Information Processing Systems, 2020-December, 2020a. ISSN 10495258.
Vincent Sitzmann, Julien N.P. Martel, Alexander W. Bergman, David B. Lindell, and Gordon Wetzstein. Implicit neural representations with periodic activation functions. Advances in Neural Information Processing Systems, 2020-December, 2020b. ISSN 10495258.
Towaki Takikawa, Alex Evans, Jonathan Tremblay, Thomas Müller, Morgan McGuire, Alec Jacobson, and Sanja Fidler. Variable bitrate neural fields. ACM Transactions on Graphics, 2022.
Matthew Tancik, Pratul P. Srinivasan, Ben Mildenhall, Sara Fridovich-Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan T. Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. Advances in Neural Information Processing Systems, 2020-December, 2020. ISSN 10495258.
Matthew Tancik, Ben Mildenhall, Terrance Wang, Divi Schmidt, Pratul P. Srinivasan, Jonathan T. Barron, and Ren Ng. Learned initializations for optimizing coordinate-based neural representations. Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition, 2021. ISSN 10636919. doi: 10.1109/CVPR46437.2021.00287.
Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural operators. 11 2023. URL http://arxiv.org/abs/2111.13802.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 2017-December, 2017. ISSN 10495258.
Petar Velickovic, Arantxa Casanova, Pietro Lio, Guillem Cucurull, Adriana Romero, and Yoshua Bengio. Graph attention networks. 2018. doi: 10.1007/978-3-031-01587-8_7.
Yuan Yin, Matthieu Kirchmeyer, Jean-Yves Franceschi, Alain Rakotomamonjy, and Patrick Gallinari. Continuous pde dynamics forecasting with implicit neural representations. International Conference on Learning Representations, 9 2022. URL http://arxiv.org/abs/2209.14855.
Luisa Zintgraf, Kyriacos Shiarlis, Vitaly Kurin, Katja Hofmann, and Shimon Whiteson. Fast context adaptation via meta-learning. 36th International Conference on Machine Learning, ICML 2019, 2019-June, 2019.
A.1 Initial Value Problem
We use the datasets from Pfaff et al. (2021), and take the first and last frames of each trajectory as the input and output data for the initial value problem.
Cylinder The dataset includes computational fluid dynamics (CFD) simulations of the flow around a cylinder, governed by the incompressible Navier-Stokes equation. These simulations were generated using COMSOL software, employing an irregular 2D-triangular mesh. The trajectory consists of 600 timestamps, with a time interval of  $\Delta t = 0.01s$  between each timestamp.
Airfoil The dataset contains CFD simulations of the flow around an airfoil, following the compressible Navier-Stokes equation. These simulations were conducted using SU2 software, using an irregular 2D-triangular mesh. The trajectory encompasses 600 timestamps, with a time interval of  $\Delta t = 0.008s$  between each timestamp.
A.2 Dynamics Modeling
2D-Navier-Stokes (Navier-Stokes) We consider the 2D Navier-Stokes equation as presented in Li et al. (2021); Yin et al. (2022). This equation models the dynamics of an incompressible fluid on a rectangular domain  $\Omega = [-1,1]^2$ . The PDE writes as:
$$
\frac {\partial w (x , t)}{\partial t} = - u (x, t) \nabla w (x, t) + \nu \Delta w (x, t) + f, x \in [ - 1, 1 ] ^ {2}, t \in [ 0, T ] \tag {6}
$$
$$
w (x, t) = \nabla \times u (x, t), x \in [ - 1, 1 ] ^ {2}, t \in [ 0, T ] \tag {7}
$$
$$
\nabla u (x, t) = 0, x \in [ - 1, 1 ] ^ {2}, t \in [ 0, T ] \tag {8}
$$
where  $u$  is the velocity,  $w$  the vorticity.  $\nu$  is the fluid viscosity, and  $f$  is the forcing term, given by:
$$
f \left(x _ {1}, x _ {2}\right) = 0. 1 \left(\sin \left(2 \pi \left(x _ {1} + x _ {2}\right)\right) + \cos \left(2 \pi \left(x _ {1} + x _ {2}\right)\right)\right), \forall x \in \Omega \tag {9}
$$
For this problem, we consider periodic boundary conditions.
By sampling initial conditions as in Li et al. (2021), we generated different trajectories on a  $256 \times 256$  regular spatial grid and with a time resolution  $\delta t = 1$ . We retain the trajectory starting from the 20th timestep so that the dynamics is sufficiently expressed. The final trajectories contains 40 snapshots at time  $t = 20, 21, \dots, 59$ . As explained in section 4, we divide these long trajectories into 2 parts: the 20 first frames are used during the training phase and are denoted as  $In - t$  throughout this paper. The 20 last timesteps are reserved for evaluating the extrapolation capabilities of the models and are the Out-  $t$  part of the trajectories. In total, we collected 256 trajectories for training, and 16 for evaluation.
3D-Spherical Shallow-Water (Shallow-Water). We consider the shallow-water equation on a sphere describing the movements of the Earth's atmosphere:
$$
\frac {\mathrm {d} u}{\mathrm {d} t} = - f \cdot k \times u - g \nabla h + \nu \Delta u \tag {10}
$$
$$
\frac {\mathrm {d} h}{\mathrm {d} t} = - h \nabla \cdot u + \nu \Delta h \tag {11}
$$