Continuous PDE Dynamics Forecasting with Implicit Neural Representations
Abstract
Effective data-driven PDE forecasting methods often rely on fixed spatial and / or temporal discretizations.
This raises limitations in real-world applications like weather prediction where flexible extrapolation at arbitrary spatiotemporal locations is required.
We address this problem by introducing a new data-driven approach, DIN O , that models a PDE’s flow with continuous-time dynamics of spatially continuous functions.
This is achieved by embedding spatial observations independently of their discretization via Implicit Neural Representations in a small latent space temporally driven by a learned ODE.
This separate and flexible treatment of time and space makes DIN O the first data-driven model to combine the following advantages.
It extrapolates at arbitrary spatial and temporal locations; it can learn from sparse irregular grids or manifolds; at test time, it generalizes to new grids or resolutions.
DIN O outperforms alternative neural PDE forecasters in a variety of challenging generalization scenarios on representative PDE systems.
Introduction
Modeling the dynamics and predicting the temporal evolution of physical phenomena is paramount in many fields, e.g. climate modeling, biology, fluid mechanics and energy [47].
Classical solutions rely on a well-established physical paradigm: the evolution is described by differential equations derived from physical first principles and then solved using numerical analysis tools, e g. finite elements, finite volumes or spectral methods [33].
The availability of large amounts of data from observations or simulations has motivated data-driven approaches to this problem [5], leading to a rapid development of the field with deep learning.
The main motivations for this research track include developing surrogate or reduced order models that can approximate high-fidelity full order models at reduced computational costs [23], complementing classical solvers, e.g. to account for additional components of the dynamics [49] or improving low fidelity models [10].
Most of these attempts rely on workhorses of deep learning like CNNs [1] or GNNs [26, 37, 4].
They all require prior space discretization either on regular or irregular grids, such that they only capture the dynamics on the train grid and cannot generalize outside it.
Neural operators, a recent trend, learn mappings between function spaces [28, 31] and thus alleviate some limitations of prior discretization approaches.
Yet, they still rely on fixed grid discretization for training and inference: e.g., regular grids for [28] or a free-form but predetermined grid for [31].
Hence, the number and / or location of the sensors has to be fixed across train and test which is restrictive in many situations [38].
Mesh-agnostic approaches for solving canonical PDEs (Partial Differential Equations) are another trend [39, 40].
In contrast to physics-agnostic grid-based approaches, they aim at solving a known
PDE as usual solvers do, and cannot cope with unknown dynamics.
This idea was concurrently developed for computer graphics, e.g. for learning 3D shapes [41, 32, 44] and coined as Implicit
[Table]
Neural Representations (INRs).
When used as solvers, these methods can only tackle a single initial value problem and are not designed for long-term forecasting outside the training horizon.
Due to these limitations, these approaches cannot handle many practical applications such as: different geometries, e.g. PDEs lying on a Euclidean plane or an Earth-like sphere; variable sampling, e.g. irregular observation grids that evolve as in adaptive meshing [3]; scarce training data e.g. with available observations at few spatiotemporal locations; multi-scale phenomena, e.g. in climate modeling, where integrating intertwined subgrid scales a.k.a.
the closure problem is ubiquitous [51].
This motivates new machine learning models which improve existing approaches on these aspects.
In our work, we aim at forecasting PDE-based spatiotemporal physical processes with a versatile model tackling the aforementioned limitations.
We adopt an agnostic approach, i.e. not assuming any prior knowledge on the physics.
We introduce DIN O (Dynamics-aware Implicit Neural representations), a model operating continuously in space and time, with the following contributions.
Continuous flow learning.
DIN O aims at learning the PDE’s flow to forecast its solutions, in a continuous manner so that it can be trained on any spatial and temporal discretization and applied to another.
To this end, DIN O embeds spatial observations into a small latent space via INRs; then it models continuous-time evolution by a learned latent Ordinary Differential Equation (ODE).
Space-time separation.
To efficiently encode different sequences, we propose a novel INR parameterization, amplitude modulation, implementing a space-time separation of variables.
This simplifies the learned dynamics, reduces the number of parameters and greatly improves performance.
Spatiotemporal versatility.
DIN O combines the benefits of prior models, cf.
Table 1.
It tackles new sequences via its amplitude modulation.
Sequential modeling with an ODE makes it extrapolate to unseen times within or beyond the train horizon.
Thanks to INRs’ spatial flexibility, it generalizes to new grids or resolutions, predicts at arbitrary positions and handles sparse irregular grids or manifolds.
Empirical validation.
We demonstrate DIN O ’s versatility and state-of-the-art performance vs prior neural PDE forecasters, representative of grid-based, operator and INR-based methods, via thorough experiments on challenging multi-dimensional PDEs in various spatiotemporal generalization settings.
Problem Description
Problem setting.
We aim at modeling, via a data-driven approach, the temporal evolution of a continuous fully-observed spatiotemporal phenomenon.
It is described by trajectories v : R → V in a set Γ ; we use v t ≜ v ( t ) ∈ V .
Trajectories share the same dynamics but differ by their initial condition v 0 ∈ V .
R is the temporal domain and V is the functional space of the form Ω → R n , where Ω ⊂ R p is a compact domain of spatial coordinates and n the number of observed values.
In other words, v t is a spatial function of x ∈ Ω , with vectorial output v t ( x ) ∈ R n ; cf.
examples of Section 4.1.
To this end, we consider the setting illustrated in Figure 1.
We observe a finite training set of trajectories D , with a free-formed spatial observation grid X tr ⊂ Ω and on discrete times t ∈ T ⊂ [0 , T ] .
At test time, we are only given a new initial condition v 0 , with observed values v| X ts restricted to a new observation grid X ts , potentially different from X tr .
Inference is performed on both train and test trajectories given only the initial condition, on a new free-formed grid X ′ ⊂ Ω and times t ∈ T ′ ⊂
[0 , T ′ ] .
Inference grid
X ′ comprises observed positions (respectively X tr and X ts for train and test trajectories) and unobserved positions corresponding to spatial extrapolation.
Note that the inference temporal horizon is larger than the train one: T < T ′ .
For simplicity, In-s refers to data in X ′ on the observation grid ( X tr for train / X ts for test ), Out-s to data in X ′ outside the observation grid; In-t refers to times within the train horizon T ⊂ [0 , T ] , and Out-t to times in
T ′ \T ⊂ ( T, T ′ ] , beyond T , up to inference horizon T ′ .
In-t
Out-t train test
X tr X tr Ω \X tr Ω \X tr
DD ΩΩ
TT
T ′ T ′ v t v t v 0 v
Out-s
In-s Ω \X ts Ω \X ts
Out-s ΩΩ
X ts X ts
In-s
Evaluation scenarios.
The desired properties in Section 1 call for space-time continuous forecasters.
We select six criteria that should be met; cf.
column titles of Table 1.
First, the model should generalize to new initial i.e to observed
Finally, should adapt to free-form spatial domains i. to geometries (e.g. manifolds) or irregular grids
Objective.
To satisfy these requirements, we learn the flow Φ of the physical system: Φ: ( V × R ) → V , ( v
[Equation]
Learning the flow is a common strategy in sequential models to better generalize beyond the train time horizon.
Yet, so far, it has always been learned with discretized models, which poses generalization issues violating our requirements.
We describe these issues in our related work in Section A.
Model
We present DIN O , the first space / time-continuous model that tackles all prediction tasks of Section 2, without the limitations of prior approaches (Section A).
We specify DIN O ’s inference procedure (Section 3.1, Figure 2 left), then introduce its components (Section 3.2) and how they are trained (Section 3.3, Figure 2 right).
Finally, we detail our implementation based on amplitude modulation, a novel INR parameterization for spatiotemporal data performing separation of variables (Section 3.4).
3.1 Inference Model
As explained in Section 2, we aim at estimating the flow Φ in Eq.
(1), so that our model can be trained on an observed grid X tr and perform inference given a new one X ts , both possibly irregular.
To this end, we leverage a space- and time-continuous formulation, independent of a given data discretization.
At inference, DIN O starts from an initial condition v 0 ∈ V and uses a flow to forecast its dynamics.
DIN O first embeds spatial observations from v 0 into a latent vector α 0 of small dimension via an encoder of spatial functions E φ : V → R d α ( ENC ).
Then, it unrolls a latent time-continuous dynamics model f ψ : R d α → R d α given this initial condition ( DYN ).
Finally, it decodes latent vectors via a decoder D ϕ : R d α → V into a function of space ( DEC ).
At any time t , D ϕ takes as input α t and outputs a function ˜ v t : Ω → R n .
This results in the following model, illustrated in Figure 2 (left):
[Equation]
3.2 Components
We now further detail each component involved at inference from Eq.
Encoder: α t = E φ ( v t ) .
The encoder computes a latent vector α t given observation v t at any time t .
It is used in two different contexts, respectively for train and test.
At train time, given an observed trajectory v T = { v t } t ∈T , it will encode any v t into α t (see Section 3.3).
At inference time, only v dynamics.
Given the decoder D , α t is a solution to the inverse problem D ϕ ( α t ) = v t .
We solve this inverse problem with auto-decoding [34].
Denoting ℓ dec ( ϕ, α t ; v t ) =
∥ D ϕ ( α t ) − v t ∥ 22 the decoding loss where ∥·∥ 2 is the euclidean norm of a function and K the number of update steps, auto-decoding defines E φ as:
[Equation]
where ∀ k > 1 , α kt = α k
[Equation]
and φ = ϕ
[Equation]
In practice, we observe a discretization ( X tr , X ts ) and accordingly approximate the norm in ℓ dec as in Eq.
Compared to auto-encoding, auto-decoding underfits less [20] and is more flexible: without requiring specialized encoders, it handles free-formed observation grids (irregular or on a manifold) with an appropriate decoder.
Decoder: ˜ v t
= D ϕ ( α t ) .
We define a flexible decoder using an INR.
An
INR I θ is a space-continuous function parameterized by θ ∈ R d θ defined on domain Ω as illustrated in Figure 3.
INRs approximate functions independently of the observation grid, e.g. it handles irregular grids and changing observation positions unlike FNO and DeepONet.
Thus, it constitutes a flexible alternative to operators suitable to auto-decoding.
We condition the INR’s parameters θ on α t via a hypernetwork [17] h ϕ : R d α → R d θ .
It generates high-dimensional parameters θ t ∈ R
d θ of the INR given the low-dimensional latent vector
α t ∈ R d α , i.e. θ t = h ϕ ( α t ) .
In summary, the decoder D ϕ , parameterized as h by ϕ , is defined as follows (cf.
details in Section 3.4):
[Equation]
Dynamics
[Equation]
The dynamics model [Equation] defines a flow via a latent ODE.
An initial condition can be defined for all t by encoding with E φ the input function v t .
In practice, to unroll the dynamics in time, we apply integration via an ODE numerical solver [18].
Overall flow.
Combined altogether, our components approximate the data flow Φ in Eq.
[Equation]
where α t =
E φ ( v t ) .
To summarize, DIN O defines a time-continuous latent temporal model with a space-continuous emission function D ϕ , combining the flexibility of space and time continuity.
This is fully novel to our knowledge, as prior latent approaches are discretized (cf.
[14] for state-space models).
3.3 Training
Given these three components ( ENC ), ( DEC ), ( DYN ), we now present their training procedure, illustrated in Figure 2 (right).
We use a fast and simple two-stage optimization, close to recent works in video prediction [48], and provide implementation details in Section F. Given the train sequences D , we first apply auto-decoding across times to obtain the corresponding latent vectors α T = { α vt } t ∈T ,v ∈D , as well as the decoder parameters ϕ .
We then learn the parameters of the dynamics model ψ by modeling the latent flow over α vt for each v ∈ D .
We detail this procedure in Section F, which can be formalized as a bi-level optimization problem solved in parallel:
[Equation]
3.4 Decoder Implementation via Amplitude-Modulated INRs
We now specify our implementation of decoder D ϕ in Eq.
This includes the definition of the
INR architecture I θ and of the hypernetwork h ϕ .
We introduce for the latter a new method called amplitude modulation, which implements a space-time separation of variables.
I θ as FourierNet.
We implement I θ as a FourierNet, a state-of-the-art INR architecture, which instantiates a Multiplicative Filter Network (MFN, [13]).
A FourierNet relies on the recursion in
Eq.
(7), where x ∈ Ω is an input spatial location, z ( l ) ( x ) is the hidden feature vector at layer l for x and s ω ( l ) ( x ) =
[cos( ω ( l ) x ) , sin( ω ( l ) x )] is a Fourier basis, with a frequency matrix ω ( l )
[Equation]
for l ∈ (cid:74) 1 , L −
[Equation]
[Equation]
where we fix W (0) = 0 , b (0)
= 1 , s ω (0) ( x ) = x .
Denoting W =
[ W ( l ) ]
L − 1 l =1 , b =
[ b ( l ) ]
L − 1 l =1 , ω = [ ω ( l ) ]
L − 1 l =1 , we fit a FourierNet to an input function v observed on a grid X by learning { W, b, ω } s.t.
∀ x ∈ X , z ( L ) ( x ) = v ( x ) .
In practice, we observe that fixing ω uniformly sampled in a given interval (cf.
Section F) performs similarly to learning them, so we exclude them from training parameters.
[Equation]
xx
[Equation]
We leverage the interpretability of FourierNets to separate time and space via amplitude modulation.
[13] shows that ∃ M ≫ L ∈ N , ∃{ c ( m ) j } Mm =1 coefficients that depend individually on { W, b } and ∃{ γ ( m ) } Mm =1 parameters that depend individually on those of the filters ω s.t.
the j th dimension of z ( L ) ( x ) can be expressed as:
[Equation]
Eq.
(8) involves a basis of spatial functions { s γ ( m ) } Mm = evaluated at x and their amplitudes c ( m ) j } Mm [Equation]
Note Eq.
(8) can be extended to other choices of s ( l ) [13]. h as amplitude modulation.
h generates the INR’s parameters θ t given α t to model a target function v t .
We implement h as elementwise shift and scale transformations [FiLM, 36] of the linear layers parameters W, b (excluding those of the filters ω ).
Then in Eq.
(8) amplitudes c ( m ) j only depend on time while the basis functions s γ ( m ) only depend on space: this corresponds to separation of variable [25].
We call our technique amplitude modulation.
In practice, as [12], we consider latent shift mappings: Figure 4 and Eq.
(9) extend Eq.
(7) with a shift term µ ( l − 1) t at each layer l , µ ( l − 1) t ≜ W ′
l − 1) α t , where
W ′ =
[ W ′ ( l − 1) ]
L − 1 l =1 is a weight matrix:
[Equation]
[Table]
The INR’s parameters are defined as h ϕ ( α t ) =
{ W ; b + W ′ α t ; ω } where ϕ
= { W, b, W ′ , ω } are h ’s parameters.
Thus, amplitude modulation separates time and space.
We show in Table 5 that it significantly improves performance, particularly time extrapolation.
Experiments
We assess the spatiotemporal versatility of DIN O , following Section 2.
We introduce our experimental setting (Section 4.1), which includes several challenging PDEs, state-of-the-art baselines and forecasting tasks.
Then, we analyze the experimental results (Section 4.2).
To reproduce our results, our code is publicly available at https://github.com/mkirchmeyer/DINo .
4.1 Experimental Setting
Datasets.
We consider the following PDEs defined over a spatial domain Ω , with further details in
Section E. • 2D Wave equation ( Wave ) is a second-order PDE ∂ 2 u ∂t 2 =
c 2 ∆ u .
u is the displacement w.r.t.
the rest position and c is the wave traveling speed.
We consider its first-order form, so that v t = ( u t , ∂u t ∂t ) has a two-dimensional output ( n = 2 ).
• 2D Navier Stokes
([43], Navier-Stokes ) corresponds to an incompressible fluid dynamics d v d t = [Equation] [Equation]
[Equation]
= 0 , where u is the velocity field and v the vorticity.
ν is the viscosity and f is a constant forcing term; n = 1 .
• 3D Spherical shallow water ([16], Shallow-Water ): it involves the vorticity w , tangent to the sphere’s surface, and the thickness of the fluid h .
The input is v t = ( w t , h t ) ( n = 2 ).
Baselines.
We reimplement representative models from Section A and Table 1 and adapt them to our multi-dimensional datasets.
• CNODE [1] combines a CNN and ODE solver to handle regular grids.
• MP-PDE [4] uses a GNN to handle free-formed grids, yet cannot predict outside the observed grid.
We developed an interpolative extension, I-MP-PDE , to handle this limitation; training is done on the bicubic interpolation of the observed grid.
• MNO [27]: an autoregressive version of FNO [28] for regular grids; MNO can be evaluated on new uniform grids.
• DeepONet [31], considered autoregressively [45] where we remove time from the trunk net’s input.
DeepONet can be evaluated on new spatial locations without interpolation.
• SIREN [41] and MFN [13]: two INR methods extended to our setting.
We consider an agnostic setting, i.e. without the knowledge of the PDE and perform sequence conditioning to generalize to more than a trajectory.
This is achieved by learning a latent vector with auto-decoding; it is then concatenated to the spatial coordinates.
Tasks.
We evaluate models on various forecasting tasks which combine the evaluation scenarios of Section 2.
Performance is measured by the prediction Mean Squared Error (MSE) given only an initial condition.
• Space and time extrapolation.
We consider a uniform grid X ′ for inference.
Training is performed on different observations grids X tr subsampled from X ′ with different ratios,
[Table]
uniform
[Table] with resolution [Table] s ∈ { 5% , 25% , 50% , 100% } where s = 100% corresponds to the full inference grid, i.e. X tr = X ′ .
In this setting, we consider that all trajectories ( train and test ) share the same observation grid
X tr = X ts .
We evaluate MSE error on X ′ over the train time interval ( In-t ) and beyond ( Out-t ) at each subsampling ratio.
• Flexibility w.r.t.
input grid.
We vary the test observation grid, i.e.
X ts ̸
= X tr and perform inference on X ′ =
X ts , i.e. on the test observation grid ( In-s ) under two settings: ▷ Generalizing across grids: [Equation] , X ts are subsampled differently from the same uniform grid; s tr (resp.
s ts ) is the train (resp.
test) subsampling ratio.
▷ Generalizing across resolutions:
X tr , X ts are subsampled with the same ratio s from two uniform grids with different resolutions; the train resolution is fixed to r tr = 64 while we vary the test resolution r ts ∈ { 32 , 64 , 256 } .
• Data on manifold.
combines several evaluation scenarios for a PDE on a sphere.
• Finer time resolution. considers a finer inference time grid than the train one; we present this experiment in Section B.
4.2 Results
Space and time extrapolation.
We report prediction MSE in Table 2 for varying subsampling ratios s ∈ { 5 , 100 } % on Navier-Stokes and Wave .
Section C provides a fine-grained evaluation inside the train observation grid ( In-s ) or outside ( Out-s ) and reports additionally the results for s ∈ { 25 , 50 } % .
We visualize some predictions in Section D. DIN O is compared to all baselines when s = 100% , i.e. X ′ =
X tr = X ts , and otherwise it is compared only to models which handle irregular grids and prediction at arbitrary spatial locations (DeepONet, SIREN, MFN, I-MP-PDE). • General analysis. We
observe that all models degrade when the subsampling ratio s decreases.
DIN O performs competitively overall: it achieves the best Out-t performance on all subsampling settings, it outperforms all the baselines on low subsampling ratios and performs comparably to the competitive discretized alternatives (MP-PDE, CNODE) and operator (MNO) when s = 100% , i.e. when observation and inference grids are equal.
Note that this fully observed setting is favorable for
CNODE, MP-PDE and MNO, designed to perform inference on the observation grid.
This can be seen in Table 2, where DIN O is slightly outperformed only for few settings.
MP-PDE is significantly better only on Wave for In-t .
Overall, CNNs and GNNs exhibit good performance for spatially local dynamics like Wave , while INRs (like DIN O ) and MNO are more adapted to global dynamics like
Navier-Stokes .
• Analysis per model.
MP-PDE is the most competitive baseline across datasets as it combines a strong and flexible encoder (GNNs) to a good dynamics model; however, it cannot predict outside the observation grid ( Out-s ).
To keep a strong competitor, we extend this baseline into its interpolative version I-MP-PDE on subsampled settings.
I-MP-PDE is competitive for high subsampling ratios, e.g. s ∈ { 50% , 100% } , but underperforms w.r.t.
DIN O at lower subsampling ratios due to the accumulated interpolation error.
MNO is a competitive baseline on Navier-Stokes , performing on par with MP-PDE and DIN O inside the training horizon ( In-t ); its performance on Out-t degrades more significantly compared to other models, especially DIN O .
DeepONet is more flexible than MP-PDE as it can predict at arbitrary locations.
As no interpolation error is introduced, it outperforms I-MP-PDE for s = 5% on train data.
Yet, we observe that it underperforms especially on Out-t w.r.t.
its alternatives.
Finally, we observe that SIREN and MFN fit correctly the train horizon In-t on train , yet generalize poorly outside this horizon Out-t or on new initial conditions ( test ).
This is in accordance with our analysis of Section A; we highlight that this is not the case for DIN O which extrapolates temporally and generalizes to new initial conditions thanks to its sequential modeling of the flow.
Thus, DIN O is currently the state-of-the-art INR model for spatiotemporal data .
• Modulation.
We observe that modulating both amplitudes and frequencies (row “DIN O (no sep.)” in Table 2)
degrades performance w.r.t.
DIN O (row DIN O in Table 2) that only modulates amplitudes.
Amplitude modulation enables long temporal extrapolation and reduces the number of parameters.
Hence, as opposed to DIN O (no sep.) which is outperformed by some baselines, time-space variable separation in DIN O is an essential ingredient of the model to reach state-of-the-art levels.
Flexibility w.r.t.
input grid.
We consider in Table 3 Navier-Stokes and compare DIN O to the most competitive baselines, MP-PDE and MNO (with s = 100% subsampling ratio).
• Generalizing across grids.
In Table 3a, we consider that the test observation grid X ts is different from the train one
X tr .
This occurs when sensors differ between two observed trajectories.
We vary the subsampling ratio for the train observation grid s tr and the test one s ts .
We report test MSE on new grids X ′ =
X ts .
DIN O is very robust to changing grids between train and test , while MP-PDE’s performance degrades, especially for low subsampling ratios, e.g. 5%.
For reference, we report in Section C, Table 6 (col. 3), the performance when X ′ =
X tr , where MP-PDE is substantially better.
Generalizing across spatial resolutions.
In Table 3b we vary the test resolution r ts .
We train at a resolution r tr = 64 and perform inference at resolutions r ts ∈ { 32 , 64 , 256 } .
For that, we build a high-fidelity 256 × 256 simulation dataset and downscale it to obtain the other resolutions.
We observe that DIN O ’s performance is the stablest across resolutions in the uniform or irregular setting.
MNO is also relatively stable but is only applicable to uniform grids while MP-PDE is particularly brittle, especially for a 5% ratio.[-R e f e r e n c e D I N o t = 
 t =-]
[Figure]
Data on manifold.
We consider in Figure 5 ShallowWater in a super-resolution setting: test resolution is twice the train one, close to weather prediction applications.
We observe an irregular 3D Euclidean coordinate grid
X tr = X ts ⊂ R 3 shared for train and test .
It samples uni- formly Euclidean positions on the sphere, via the quasiuniform skipped latitude-longitude grid [46].
We predict the PDE on test trajectories with a conventional latitudelongitude inference grid X ′ .
At Earth scale, X tr corre- sponds to a resolution of about 300 km, and X ′ to 150 km.
DIN O significantly outperforms I-MP-PDE, making it a viable candidate for this complex setting.
Conclusion
We propose DIN O , a novel space-time continuous datadriven PDE forecaster.
DIN O handles settings encountered in many applications, where existing methods fail.
We assess its extensive spatiotemporal extrapolation on several PDEs and its generalization to unseen sparse irregular meshes and resolutions.
Its competitive results over recent neural PDE forecasters make it a strong alternative for real-world free-formed spatiotemporal conditions.
There are many promising extensions of DIN O , for example improving its generalization to new system parameters [22] or to new boundary conditions.
Acknowledgement
We thank Emmanuel de Bézenac and Jérémie Donà for helpful insights and discussions on this project.
We also acknowledge financial support from DL4CLIM (ANR-19-CHIA-0018-01) and DEEPNUM (ANR-21-CE23-0017-02) ANR projects.
A Related Work
We review current data-driven approaches for PDE modeling and the representative methods listed in
Table 1.
We express the forecasting rule using the notations in Eq.
(1): t is an arbitrary time; τ is an arbitrary time interval; δt is a fixed, predetermined time interval (as a model hyperparameter).
Sequential discretized models.
Most sequential dynamics models are learned on a fixed observed grid X tr and use discretized models, e.g. CNNs or GNNs, to process the observations.
CNNs require observations on a regular grid but can be extended to irregular grids through interpolation [7].
GNNs are more flexible as they handle irregular grids, at an additional memory and computational cost.
Yet, prediction on new grids X ′ ̸
= X tr fails experimentally for both CNNs and GNNs, as these discretized models are biased towards the training grid X tr , as we show in Section 4.
We distinguish two types of temporal models which both extrapolate beyond the train horizon
due
to
their sequential nature. • Autoregressive models v t | X (cid:55)→ v t + δt | X [30, 11, 37, 4].
These models predict the sequence from t only at fixed time increments δt and not in between.
• Time-continuous extensions using numerical solvers ( v t
| X , τ ) (cid:55)→ v t + τ | X [49, 19] solve this limitation as they provide a prediction at arbitrary times, thus remove dependency on the time discretization.
Operator learning.
Recently, operator-based models aim at finding a parameterized mapping between functions.
They define in theory space-continuous models.
First, neural operators [24] attempt to replace standard convolution with continuous alternatives.
Fourier Neural Operator (FNO, [28]) applies convolution in the spectral domain via Fast Fourier Transformation (FFT).
Graph Neural
Operator (GNO, [26]) performs convolution on a local interaction grid described by a graph.
Second,
DeepONet [31] uses a coordinate-based neural network to output a prediction at arbitrary time and space locations given a function observed on a fixed grid.
Three types of temporal models were used for operators with some limitations.
• The standard approach, v 0 (cid:55)→ v t , models the output at a given time t ∈ [0 , T ] within the train horizon [26]. • A sequential extension, v t (cid:55)→ v t + δt , was proposed in [27]. • Finally, a time-continuous version v 0 (cid:55)→ ( t ∈ [0 , T ] (cid:55)→ v t ) in DeepONet propose a solution at arbitrary time and space locations.
The first and third approaches are not designed to generalize beyond the train horizon, i.e. when t > T as they are not sequential.
The second solves this limitation but is only able to predict solutions from t at fixed time increments of δt and not in-between.
Furthermore, all existing approaches make restrictive assumptions on the space discretization.
They lack flexibility when encoding spatial observations: FNO is limited to uniform Cartesian observation grids due to FFT; GNO does not adapt well to changing observation grids as for the GNN-based models in the previous paragraph;
DeepONet is limited to input observations on fixed observation locations – the latter are chosen at random spatial positions but should remain fixed throughout training and testing.
Concurrent to our work, [29] extends Markov FNO [28, 27] and circumvents these issues by embedding spatial observations in a Euclidean space; however, its spatial flexibility remains limited: it does not handle arbitrary topologies, it has a quadratic complexity in the number of points, and, like its predecessors, cannot be evaluated at unobserved locations in the conditioning frame.
Spatiotemporal INRs.
Another class of models is based on coordinate-based NNs, called Implicit
Neural Representations [INRs, 41, 13, 44].
These space-continuous models share a similar objective as operators, despite constituting a separate research field.
INRs for spatiotemporal data take time as an input along spatial coordinates.
Physics-informed neural networks (PINNs, [39]) use this formulation to solve PDEs, yet are limited to a single known differential equation and a set of initial and boundary conditions.
[15] proposes an agnostic INR approach to build reduced order models for electrophysiology.
Extensions for multi-sequence learning, e.g. for video generation [50, 42] or compression [8], learn a latent conditioning variable from an initial condition v 0 , i.e. take the form
[Table]
v 0 (cid:55)→ ( t ∈ [0 , T ] (cid:55)→ v t ) .
Interestingly, these models can predict at an arbitrary time t in the train horizon without unrolling a sequential model up to t .
Yet, as they only learn mappings from an initial condition v 0 to a function of time v t in the train domain, they fail to predict beyond train conditions, as we show in Section 4.
DIN O is a new instance of spatiotemporal INR which solves this limitation via a time-continuous dynamics model of the underlying flow, ( v t , τ ) (cid:55)→ v t + τ .
B Additional Results
We present another experiments, finer time resolution , where we choose the inference time grid T ′ with a finer resolution than the train one T .
We consider in Table 4 a longer and 10 times finer test time grid T ′ than the train grid T on Navier-Stokes .
We observe the same spatial uniform grid across train and test and perform inference on this grid.
We compare DIN O that performs prediction with an
ODE solver, to interpolating coarser predictions obtained at the train resolution (I-DIN O ).
We report the corresponding test MSE.
We observe that the ODE solver accurately extrapolates outside the train temporal grid, outperforming interpolation.
This confirms that DIN O benefits from its continuoustime modeling of the flow, providing consistency and stability across temporal resolutions.
C Full Results
We provide in Table 5 a more detailed version of Table 2 for the space-time extrapolation problem where we report the performance In-s (on the observation grid) and Out-s (outside).
We add s = 50 %.
Then, we report in Table 6, a more detailed version of Table 3a, which includes the results of X ts = X tr .
This corresponds to our generalization across grids problem.
D Prediction Examples
We display the test prediction of DIN O (Figure 6) and I-MP-PDE (Figure 7) on Navier-Stokes for various subsampling levels when X
= X tr = X ts .
We plot the prediction of DIN O on Wave in the same setting in Figure 8.
Predictions are performed on a× 64 uniform grid which defines the observation grid X via different subsampling rates.
Yellow points correspond to the observation grid
X ( In-s ) while purple points indicate off-grid points ( Out-s ).
The prediction for I-MP-PDE at t =is the interpolated initial condition.
E Detailed Description of Datasets
We choose T (resp.
T ′ )
on a regular grid in [0 , T ] (resp.
[0 , T ′ ] ) with a given temporal resolution and fix T ′ = 2 T .
We provide further details on the choice of these parameters and other experimental parameters, such as the number of observed trajectories. 2D Wave equation ( Wave ) .
It is a second-order PDE:
[Equation]
where u is a function of the displacement at each point in space w.r.t.
the rest position, c ∈ R ∗ + is the speed of wave traveling.
We transform the equation to a first-order form, considering the input v t = (cid:16) u t , ∂u t ∂t (cid:17) , so that the dimension of v t ( x ) at each point x ∈ Ω is n = 2 .
[Table]
u t
[Table]
[Table] [Table]
[Figure] [Figure]
We generate our dataset for speed c = 2 with periodic boundary condition.
The domain is Ω = [ − 1 , 1] 2 .
For initial conditions v 0 = (cid:16) u 0 , ∂u t ∂t (cid:12)(cid:12)(cid:12) t =0 (cid:17) , the initial displacement u 0 is Gaussian:
[Equation]
11) where the height of the peak displacement is a ∼ U (2 , 4) , the location of the peak displacement is ( b 1 , b 2 ) ∼ U ( − 1 , 1) , and the standard deviation is r ∼ U (0 . 25 , 0 . 3) .
The initial time derivative is ∂u t ∂t (cid:12)(cid:12)(cid:12) t {+=
+} = 0 .
Each snapshot is generated on a uniform grid of 64 × 64.
Each sequence is generated with fixed interval δt = 0 .
We set the training horizon T = 2 .
25 and the inference horizon
T = 4 .
We generated 512 training trajectories and 32 test trajectories. 2D Navier Stokes ( Navier-Stokes , [43]) .
This dataset describes an incompressible fluid dynamics by:
[Equation]
12) where u is the velocity field and w the vorticity.
u, w lie on a spatial domain with periodic boundary conditions, ν is the viscosity and f is a constant forcing term.
The input v t is w t ( n = 1 ).
ν is the viscosity and f is the constant forcing term in the domain Ω .
[Figure]
The spatial domain is Ω =
[ − 1 , 1] 2 , the viscosity is ν = 1 × 10 − 3 , the forcing term is set as:
[Equation]
The full spatial grid is of dimension 64 × 64 or 256 × 256 according to experiments in Section 4.
We sample initial conditions as in Li et al.
[28] to create different trajectories.
The first 20 steps of the trajectories are cut off as they are too noisy and not informative in terms of dynamics.
Trajectories are collected with δt = 1 .
We set the training horizon T = 19 and the inference horizon T ′ = 39 .
We generated 512 training trajectories and 32 test trajectories. 3D spherical shallow water ( Shallow-Water , [16]) .
The following problem is originally presented for testing numerical models of global shallow-water equations.
The shallow water equations is written as:
[Equation]
(14) where dd t is the material derivative, k is the unit vector orthogonal to the spherical surface, u is the velocity field tangent to the surface of the sphere, which can be transformed into the vorticity w [Equation] while not converged do for [Equation]
while not converged do for
[Equation]
The initial conditions are slightly modified from [16], detailed below, to create symmetric phenomena on the northern and southern hemisphere.
The initial zonal velocity u 0 contains two non-null symmetric bands in the both hemispheres, which are parallel to the circles of latitude.
At each latitude and longitude ϕ, θ ∈ [ − π / 2 , π / 2 ] × [ − π, π ] :
[Equation]
where u max is the maximum velocity, ϕ
= π / 7 ,
ϕ 1 = π / 2 − ϕ 0 , and e n = exp (cid:0) − 4 / ( ϕ 1 − ϕ 0 ) 2 (cid:1) .
The water height h 0 is initialized by solving a boundary value condition problem as in Galewsky et al.
It is then perturbed by adding the following h ′ 0 to h 0 :
[Equation]
where ϕ 2 = π/ 4 , ˆ h = 120 m , α = 1 / 3 , β = 1 / 15 are constants defined in [16].
We simulate this phenomenon with Dedalus [6] on a latitude-longitude (lat-lon) grid.
The size of the grid is 128 (lat) × 256 (lon).
We take different initial conditions by sampling u max ∼ U (60 , 80) to generate long trajectories.
These long trajectories are then sliced into shorter ones.
For simulation, we take one snapshot per hour (of internal simulation time), i.e. δt = 1 h .
We stop the simulation at the 320 th hour.
To construct a dataset rich of dynamical phenomena, we take the snapshots within the last 160 h in a long trajectory and slice them into 8 shorter trajectoires.
Also note that the data is scaled into a reasonable range: the height h is scaled by a factor of 3 × 10 3 , and the vorticity w by a factor 2.
In each short trajectory, T = 9 h and T ′ = 19 h .
In total, we generated 16 long trajectories (i.e. 128 short trajectories) for train, 2 for test (i.e. 16 short trajectories).
F Implementation
Algorithm.
We detail the algorithm of DIN O for training and test via pseudo-code in Algorithm 1.
[-Test: Input : D ′ 
-]
Implementation details.
We use PyTorch [35] to implement DIN O and the baselines.
The dynamics model f ψ is a multilayer perceptron.
Its input and output size are same as the size of latent space d α .
All hidden layers share the same size.
DIN O ’s parameters are initialized with the default initialization in PyTorch.
The frequency parameters in FourierNet are scaled by a factor, considered as a hyperparameter.
For dynamics learning, we use an RK4 integrator and apply exponential Scheduled
Sampling [2] to stabilize training.
ω is fixed at initialization to reduce the number of optimized parameters without loss of performance.
In practice, modulations of α t are learned channel-wise such [Table] that I θ : Ω → R d c has separate parameters per output dimension to make predictions less correlated across channels.
We optimize all parameters using Adam [21] with ( β 1 , β 2 ) =
Hyperparameters.
We list the hyperparameters of DIN O for each dataset in Table 7.
In practice, we observe it is beneficial to decay the learning rates η ϕ , η α when the loss reaches a plateau. Baselines implementation.
We detail in the following the hyperparameters and architectures used in our experiments for the considered baselines
[Table]
which we reimplemented for our paper.
• CNODE is implemented with four 2D convolutional layers with 64 hidden features, ReLU activations, 3 × 3 kernel and zero padding.
Learning rate is fixed to 10 − 3 .
We use an adjoint method for integration like [9].
• MNO: we use the FNO architecture in [28] with three FNO blocks, GeLU activations,modes and a width of 32 .
Learning rate is fixed to 10 − 3 .
• DeepONet: we consider an autoregressive formulation of DeepONet.
We choose a width of 1000 for hidden features with a depth of 4 for both trunk and branch nets with ReLU activations.
Learning rate is fixed to 10 − 5 .
• MP-PDE: we adapt the implementation in [4] to handle 2D and 3D PDEs.
We use a time window of 1 with pushforward trick.
Batch size and number of neighbors are fixed to 8 .
Learning rate is fixed to 10 − 3 .
We use ReLU activations.
• SIREN: to represent data in space and time, SIREN takes space and time coordinates ( x, t ) as input.
To handle multiple trajectories, we concatenate an optimizable per-trajectory context code α to the coordinates like in DIN O .
We fix the hidden layer size of SIREN to 256 .
We initialize the parameters and use the default input scale as in [41].
The size of the context code is d α = 800 .
The learning rate is 10 − 3 .
• MFN: similarly to the previous SIREN baseline, we concatenate the per-trajectory context code to space and time coordinates at the first layer.
The hidden layer size is fixed toand we use the default parameter initialization with a frequency scale of × 64 as DIN O .
The size of the context code is d α = 800 .
The learning rate is 10 − 3 .