# A Neural PDE Solver with Temporal Stencil Modeling

Zhiqing Sun  Yiming Yang  Shinjae Yoo

###### Abstract

Numerical simulation of non-linear partial differential equations plays a
crucial role in modeling physical science and engineering phenomena, such as
weather, climate, and aerodynamics. Recent Machine Learning (ML) models
trained on low-resolution spatio-temporal signals have shown new promises in
capturing important dynamics in high-resolution signals, under the condition
that the models can effectively recover the missing details. However, this
study shows that significant information is often lost in the low-resolution
down-sampled features. To address such issues, we propose a new approach,
namely Temporal Stencil Modeling (TSM), which combines the strengths of
advanced time-series sequence modeling (with the HiPPO features) and state-of-
the-art neural PDE solvers (with learnable stencil modeling). TSM aims to
recover the lost information from the PDE trajectories and can be regarded as
a temporal generalization of classic finite volume methods such as WENO. Our
experimental results show that TSM achieves the new state-of-the-art
simulation accuracy for 2-D incompressible Navier-Stokes turbulent flows: it
significantly outperforms the previously reported best results by 19.9% in
terms of the highly-correlated duration time and reduces the inference latency
into 80%. We also show a strong generalization ability of the proposed method
to various out-of-distribution turbulent flow settings. Our code is available
at https://github.com/Edward-Sun/TSM-PDE.

Machine Learning, ICML

Figure 1: Illustration of classic finite volume solvers (in red color),
learnable solvers with vanilla stencil modeling (in blue color) and our
temporal stencil modeling (in green color). While the convective flux
approximation methods are different in each method, the divergence operator,
the explicit time-step operator, and the pressure projection (in yellow color)
are shared between classic solvers and learnable methods. Notice that the
stencil interpolation coefficients in classic solvers such as WENO can also be
data-adaptive (see Sec. 3.1 for more details).

##  1 Introduction

Complex physical systems described by non-linear partial differential
equations (PDEs) are ubiquitous throughout the real world, with applications
ranging from design problems in aeronautics (Rhie & Chow, 1983), medicine
(Sallam & Hwang, 1984), to scientific problems of molecular modeling (Lelievre
& Stoltz, 2016) and astronomical simulations (Courant et al., 1967). Solving
most equations of importance is usually computationally intractable with
direct numerical simulations and the finest features in high resolutions.

Recent advances in machine learning-accelerated PDE solvers (Bar-Sinai et al.
2019; Li et al. 2020c; Kochkov et al. 2021; Brandstetter et al. 2021, inter
alia) have shown that end-to-end neural solvers can efficiently solve
important (mostly temporal) partial differential equations. Unlike classical
finite differences, finite volumes, finite elements, or pseudo-spectral
methods that require a smooth variation on the high-resolution meshes for
guaranteed convergence, neural solvers do not rely on such conditions and are
able to model the underlying physics with under-resolved low resolutions and
produce high-quality simulations with significantly reduced computational
cost.

The power of learnable PDE solvers is usually believed to come from the super-
resolution ability of neural networks, which means that the machine learning
model is capable of recovering the missing details based on the coarse
features (Bar-Sinai et al., 2019; Kochkov et al., 2021). In this paper, we
first empirically verify such capability by explicitly training a super-
resolution model, and then find that since low-resolution down-sampling of the
field can lead to some information loss, a single coarse feature map used by
previous work (Kochkov et al., 2021) is not sufficient enough. We empirically
show that the temporal information in the trajectories and the temporal
feature encoding scheme are crucial for recovering the super-resolution
details faithfully.

Motivated by the above observations, we propose Temporal Stencil Modeling
(TSM), which combines the best of two worlds: stencil learning (i.e., Learned
Interpolation in Kochkov et al. 2021) as that used in a state-of-the-art
neural PDE solver for conservation-form PDEs, and HiPPO (Gu et al., 2020) as a
state-of-the-art time series sequence model. Specifically, in this paper, we
focus on trajectory-enhanced high-quality approximation of the convective flux
within a finite volume method framework. As illustrated in Fig. 1, TSM can be
regarded as a temporal generalization of classic finite volume methods such as
WENO (Liu et al., 1994; Jiang & Shu, 1996) and recently proposed learned
interpolation solvers (Kochkov et al., 2021), both of which adaptively weight
or interpolate the stencils based on the latest states only. On the other
hand, in TSM we use the HiPPO-based temporal features to calculate the
interpolation coefficients for approximating the integrated velocity on each
cell surface. The HiPPO temporal features provide a good representation for
calculating the interpolation coefficients, while the stencil learning
framework ensures that the neural system’s prediction exactly conserves the
Conservation Law and the incompressibility of the fluid. With the abundant
temporal information, we further utilize the temporal bundling technique
(Brandstetter et al., 2021) to avoid over-fitting and improve the prediction
latency for TSM.

Following the precedent work in the field (Li et al., 2020c; Kochkov et al.,
2021; Brandstetter et al., 2021), we evaluate the proposed TSM neural PDE
solver on the 2-D incompressible Navier-Stokes equation, which is the
governing equation for turbulent flows with the conservation of mass and
momentum in a Newtonian fluid. Our empirical evaluation shows that TSM
achieves both state-of-the-art simulation accuracy ($+$19.9%) and inference
speed ($+25\%$). We also show that TSM trained with steady-state flows can
achieve strong generalization performance on out-of-distribution turbulent
flows, including different forcings and different Reynolds numbers.

##  2 Background & Related Work

###  2.1 Navier-Stokes Equation

A time-dependent PDE in the conservation form can be written as

$\partial_{t}\mathbf{u}+\nabla\cdot\mathbf{J}(\mathbf{u})=0$ (1)

where $\mathbf{u}:[0,T]\times\mathbb{X}\rightarrow\mathbb{R}^{n}$ is the
density of the conserved quantity (i.e., the solution), $t\in[0,T]$ is the
temporal dimension, $\mathbb{X}\subset\mathbb{R}^{n}$ is the spatial
dimension, and $\mathbf{J}:\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}$ is the
flux, which represents the quantity that passes or travels (whether it
actually moves or not) through a surface or substance. $\nabla\cdot\mathbf{J}$
is the divergence of $\mathbf{J}$. Specifically, the incompressible, constant
density 2-D Navier Stokes equation for fluids has a conservation form of:

$\partial_{t}\mathbf{u}+\nabla\cdot(\mathbf{u}\otimes\mathbf{u})=\nu\nabla^{2}\mathbf{u}-\frac{1}{\rho}\nabla
p+\mathbf{f}$ (2) $\nabla\cdot\mathbf{u}=0$ (3)

where $\otimes$ denotes the tensor product, $\nu$ is the kinematic viscosity,
$\rho$ is the fluid density, $p$ is the pressure filed, and $\mathbf{f}$ is
the external forcing. In Eq. 2, the left-hand side describes acceleration and
convection, and the right-hand side is in effect a summation of diffusion,
internal forcing source, and external forcing source. Eq. 3 enforces the
incompressibility of the fluid.

A common technique to solve time-dependent PDEs is the method of lines (MOL)
(Schiesser, 2012), where the basic idea is to replace the spatial (boundary
value) derivatives in the PDE with algebraic approximations. Specifically, we
discretize the spatial domain $\mathbb{X}$ into a grid $X=\mathbb{G}^{n}$,
where $\mathbb{G}$ is a set of grids on $\mathbb{R}$. Each grid cell $g$ in
$\mathbb{G}^{n}$ denote a small non-overlapping volume, whose center is
$\mathbf{x}_{g}$, and the average solution value is calculated as
$\mathbf{u}_{g}^{t}=\int_{g}\mathbf{u}(t,\mathbf{x})d\mathbf{x}$. We then
solve $\partial_{t}\mathbf{u}_{g}^{t}$ for $g\in\mathbb{G}^{n}$ and
$t\in[0,T]$. Since $g\in\mathbb{G}^{n}$ is a set of pre-defined grid points,
the only derivative operator is now in time, making it an ordinary
differential equation (ODEs)-based system that approximates the original PDE.

###  2.2 Classical Solvers for Computational Fluid Dynamics

In Computational Fluid Dynamics (CFD) (Anderson & Wendt, 1995; Pope & Pope,
2000), the Reynolds number $Re=UL/\nu$ dictates the balance between convection
and diffusion, where $U$ and $L$ are the typical velocity and characteristic
linear dimension. When the Reynolds number $Re\gg 1$, the fluids exhibit time-
dependent chaotic behavior, known as turbulence, where the small-scale changes
in the initial conditions can lead to a large difference in the outcome. The
Direct Numerical Simulation (DNS) method solve Eq. 2 directly and is a
general-purpose solver with high stability. However, as $Re$ determines the
smallest spatio-temporal feature scale that need to be captured by DNS, DNS
faces a computational complexity as high as $O(Re^{3})$ (Choi & Moin, 2012)
and cannot scale to large-Reynolds number flows or large-size computation
domains.

###  2.3 Neural PDE Solvers

A wide range of neural network-based solvers have recently been proposed to at
the intersection of PDE solving and machine learning. We roughly classify them
into four categories:

#### Physics-Informed Neural Networks (PINNs)

PINN directly parameterizes the solution $\mathbf{u}$ as a neural network
$F:[0,T]\times\mathbb{X}\rightarrow\mathbb{R}^{n}$ (Weinan & Yu, 2018; Raissi
et al., 2019; Bar & Sochen, 2019; Smith et al., 2020; Wang et al., 2022). They
are closely related to the classic Galerkin methods (Matthies & Keese, 2005),
where the boundary condition date-fitting losses and physics-informed losses
are introduced to train the neural network. These methods suffer from the
parametric dependence issue, that is, for any new initial and boundary
conditions, the optimization problem needs to be solved from scratch, thus
limiting their applications, especially for time-dependent PDEs.

#### Neural Operator Learning

Neural Operator methods learn the mapping from any functional parametric
dependence to the solution as
$F:\left([0,T]\times\mathbb{X}\rightarrow\mathbb{R}^{n}\right)\rightarrow\left([T,T+\Delta
T]\times\mathbb{X}\rightarrow\mathbb{R}^{n}\right)$ (Lu et al., 2019;
Bhattacharya et al., 2020; Patel et al., 2021). These methods are usually not
bounded by fixed resolutions, and learn to directly predict any solution at
time step $t\in[T,T+\Delta T]$. Fourier transform (Li et al., 2020c; Tran et
al., 2021), wavelet transform (Gupta et al., 2021), random features (Nelsen &
Stuart, 2021), attention mechanism (Cao, 2021), or graph neural networks(Li et
al., 2020a, b) are often used in the neural network building blocks. Compared
to neural methods that mimic the method of lines, the operator learning
methods are not designed to generalize to dynamics for out-of-distribution
$t\in[T+\Delta T,+\infty]$, and only exhibit limited accuracy for long
trajectories.

#### Neural Method-of-Lines Solver

Neural Method-of-Lines Solvers are autoregressive models that solve the PDE
iteratively, where the difference from the latest state at time $T$ to the
state at time $T+\Delta t$ is predicted by a neural network
$\mathbf{F}:\left([0,T]\times
X\rightarrow\mathbb{R}^{n}\right)\rightarrow\left(X\rightarrow\mathbb{R}^{n}\right)_{t={T+\Delta
t}}$. The typical choices for $\mathbf{F}$ include modeling the absolute
difference: $\forall g\in X=\mathbb{G}^{n},\ \mathbf{u}_{g}^{T+\Delta
t}=\mathbf{u}_{g}^{T}+\mathbf{F}_{g}(\mathbf{u}_{[0,T]})$ (Wang et al., 2020;
Sanchez-Gonzalez et al., 2020; Stachenfeld et al., 2021) and modeling the
relative difference: $\mathbf{u}_{g}^{T+\Delta t}=\mathbf{u}_{g}^{T}+\Delta t\
\cdot\ \mathbf{F}_{g}(\mathbf{u}_{[0,T]})$ (Brandstetter et al., 2021), where
the latter is believed to have the better consistency property, i.e.,
$\lim_{\Delta t\rightarrow 0}\|\mathbf{u}_{g}^{T+\Delta
t}-\mathbf{u}_{g}^{T}\|=0$.

#### Hybrid Physics-ML

Physics-ML hybrid models is a recent line of work that uses a neural network
to correct the errors in the classic (typically low-resolution) numerical
simulators. Most of these approaches seek to learn the corrections of the
numerical simulators’ outputs (Mishra, 2019; Um et al., 2020; List et al.,
2022; Dresdner et al., 2022; Frezat et al., 2022; Bruno et al., 2022), while
Bar-Sinai et al. (2019); Kochkov et al. (2021) learn to infer the stencils of
advection-diffusion problems in a Finite Volume Method (FVM) framework. The
proposed Temporal Stencil Modeling (TSM) method belongs to the latter
category.

Figure 2: Illustration of super-resolution process from a trajectory of frames
(i.e., temporal super-resolution) or a single-frame (vanilla super-
resolution).

##  3 Temporal Stencil Modeling for PDEs

###  3.1 Neural stencil modeling in finite volume scheme

Finite Volume Method (FVM) is a special MOL technique for conservation form
PDEs, and can be derived from Eq. 1 via Gauss’ theorem, where the integral of
$\mathbf{u}$ (i.e., the averaged vector field of volume) over unit cell
increases only by the net flux into the cell. Recall that the incompressible,
constant density Navier Stokes equation for fluids has a conservation form of:

$\partial_{t}\mathbf{u}+\nabla\cdot(\mathbf{u}\otimes\mathbf{u})=\nu\nabla^{2}\mathbf{u}-\frac{1}{\rho}\nabla
p+\mathbf{f}$ (4)

We can see that in FVM, the cell-average divergence can be calculated by
summing the surface flux, so the problem boils down to estimating the
convective flux $\mathbf{u}\otimes\mathbf{u}$ on each face. This only requires
estimating $\mathbf{u}$ by interpolating the neighboring discretized
velocities, called stencils. The beauty of FVM is that the integral of
$\mathbf{u}$ is exactly conserved, and it can preserve accurate simulation as
long as the flux $\mathbf{u}\otimes\mathbf{u}$ is estimated accurately. Fig. 1
illustrates an implementation (Kochkov et al., 2021) of classic FVM for the
Navier-Stokes equation, where the convection and diffusion operators are based
on finite-difference approximations and modeled by explicit time integration,
and the pressure is implicitly modeled by the projection method (Chorin,
1967). The divergence operator enforces local conservation of momentum
according to a finite volume method, and the pressure projection enforces
incompressibility. The explicit time-step operator allows for the
incorporation of additional time-varying forces. We refer the readers to
(Kochkov et al., 2021) for more details.

Classic FVM solvers use manually designed $n^{th}$-order accurate method to
calculate the interpolation coefficients of stencils to approximate the
convective flux. For example, linear interpolation, upwind interpolation (Lax,
1959), and WENO5 method (Shu, 2003; Gottlieb et al., 2006) can achieve first,
second, and fifth-order accuracy, respectively. Besides, adjusting the
interpolation weights adaptively based on the input data is not new in
numerical simulating turbulent flows. For example, given the interpolation
axis, the upwind interpolation adaptively uses the value from the
previous/next cell along that axis plus a correction for positive/negative
velocity. WENO (weighted essentially non-oscillatory) scheme also computes the
derivative estimates by taking an adaptively-weighted average over multiple
estimates from different neighborhoods. However, the classic FVM solvers are
designed for general cases and can only adaptively adjust the interpolation
coefficients with simple patterns, and thus are sub-optimal when abundant PDE
observation data is available.

In this paper, we follow Kochkov et al. (2021) and aim to learn more accurate
flux approximation in the FVM framework by predicting the learnable
interpolation coefficients for the stencils with neural networks. In
principle, with $3\times 3$ convolutional kernels, a $1,2,3$-layer neural
network is able to perfectly mimic the linear interpolation, Lax-Wendroff
method, and WENO method, respectively (Brandstetter et al., 2021). Such
observation well connects the learnable neural stencil modeling methods to the
classical schemes. However, previous work (Bar-Sinai et al., 2019; Kochkov et
al., 2021) investigates the learned interpolation scheme that only adapts to
the latest state $\mathbf{u}^{T}$, which still uses the same information as
classical solvers. In this paper, we further generalize both classical and
previous learnable stencil interpolation schemes by predicting the
interpolation coefficients with the abundant information from all the previous
trajectories $\\{\mathbf{u}^{t}|t\in[0,T]\\}$.

Figure 3: $64\times 64\rightarrow 2048\times 2048$ super-resolution errors
with $32\Delta t$-step HiPPO features and $32\Delta t$-step raw features.

Table 1: $64\times 64\rightarrow 2048\times 2048$ super-resolution MSE with
different approaches.

Method | MSE
---|---
Bicubic Interpolation | 2.246
CNN w/ 1-step raw features | 0.029
CNN w/ 32-step raw features | 0.015
CNN w/ 32-step HiPPO features | 0.007

Figure 4: Qualitative results of predicted vorticity fields for reference (DNS
$2048\times 2048$), previous learnable sota model (LI $64\times 64$) (Kochkov
et al., 2021), and our method (TSM $64\times 64$), starting from the same
initial condition. The yellow box denotes a vortex that is not captured by LI.

Figure 5: (left) Comparison of the vorticity correlation between prediction
and the ground-truth solution (i.e., DNS $2048\times 2048$). (middle) Energy
spectrum scaled by $k^{5}$ averaged between simulation time 6.0 to 20.0.
(right) Comparison of high vorticity correlation duration v.s. inference
latency.

###  3.2 Temporal Super-resolution with HiPPO features

A fundamental question in neural stencil modeling is why ML models can predict
more accurate flux approximations, and previous work (Kochkov et al., 2021)
attribute their success to the super-resolution power of neural networks, that
is, the machine learning models can recover the missing details from the
coarse features. In this paper, we empirically verify this hypothesis by
explicitly training a super-resolution model.

Specifically, we treat the 2-D fluid velocity map as a $H\times W\times 2$
image, and train a CNN-based U-Net decoder (Ronneberger et al., 2015) to
generate the super-resolution vorticity results. Our results are reported in
Tab. 1 (right), and we find that the super-resolution results generated by
neural networks is nearly $100\times$ better than bicubic interpolation, which
verify the super-resolution power of ML models to recover the details.

Next, since the temporal information is always available for PDE
simulation111Even if we only have access to one step of the ground-truth
trajectory, we can still use the high-resolution DNS to generate a long enough
trajectory to initialize the spatial-temporal model., we investigate whether
the temporal information can further reduce the super-resolution errors, i.e.,
recovering more details from the coarse features. After preliminary
experiments, we decide to keep the convolutional neural networks as the
spatial encoder module for their efficient implementation on GPUs and
translation invariant property, and only change the temporal input features.
Inspired by the recent progress in time series modeling, we consider the
following two types of features as the CNN model’s inputs:

#### Raw features

Following the previous work on video super-resolution (Liao et al., 2015), we
treat the $H\times W\times T\times 2$ -shaped velocity trajectory as an
$H\times W$ image with feature channels $C=T\times 2$.

#### HiPPO features

HiPPO (High-order Polynomial Projection Operators) (Gu et al., 2020, 2021) is
a recently proposed framework for the online compression of continuous time
series by projection onto polynomial bases. It computes the optimal polynomial
coefficients for the scaled Legendre polynomial basis. It has been shown as a
state-of-the-art autoregressive sequence model for raw images (Tay et al.,
2020) and audio (Goel et al., 2022). In this paper, we propose to adopt HiPPO
to encode the raw fluid velocity trajectories. Due to the space limit, we
leave the general description of the HiPPO technique to Appendix C.

#### Results

We report the temporal super-resolution results in Tab. 1. From the table, we
can see that the $32\Delta t$-step raw features can reduce the super-
resolution error by half, and using the HiPPO to encode the time series can
further reduce the error scale by half. We also visualize the errors of
$32\Delta t$-step HiPPO features and $32\Delta t$-step raw features in Fig. 3.
Our temporal super-resolution results show that plenty of information in the
PDE low-resolution temporal trajectories is missed with vanilla stencil
modeling, or will be underutilized with raw features, and HiPPO can better
exploit the temporal information in the velocity trajectories.

###  3.3 Temporal Stencil Modeling with HiPPO Features

As better super-resolution performance indicate that more details can be
recovered from the low-resolution features, we propose to compute the
interpolation coefficients in convective flux with the HiPPO-based temporal
information, which should lead to a more accurate flux approximation.
Incorporating HiPPO features in the neural stencil modeling framework is
straight-forward: with ground-truth initial velocity trajectories
$\mathbf{v}_{[0,T]}$, we first recurrently encode the trajectory (step by step
with Eq. 7 in Appendix C), and use the resulted features
$\mathrm{HiPPO}{(\mathbf{v}_{[0,T]})}$ to compute the interpolation
coefficients for the stencils. Given model-generated new velocity
$\mathbf{v}_{T+\Delta t}$, due to the recurrence property of HiPPO, we can
only apply a single update on $\mathrm{HiPPO}{(\mathbf{v}_{[0,T]})}$ and get
the new encoded feature $\mathrm{HiPPO}{(\mathbf{v}_{[0,T+\Delta t]})}$, which
is very efficient. Fig. 9 in the appendix illustrates such a process.

##  4 Experiments

###  4.1 Experimental setup

#### Simulated data

Following previous work (Kochkov et al., 2021), we train our method with 2-D
Kolmogorov flow, a variant of incompressible Navier-Stokes flow with constant
forcing $\mathbf{f}=\sin(4y)\hat{\mathbf{x}}-0.1\mathbf{u}$. All training and
evaluation data are generated with a JAX-
based222https://github.com/google/jax-cfd finite volume-based direct numerical
simulator in a staggered-square mesh (McDonough, 2007) as briefly described in
Sec. 3.1. We refer the readers to the appendix of (Kochkov et al., 2021) for
more data generation details.

We train the neural models on $Re=1000$ flow data with density $\rho=1$ and
viscosity $\nu=0.001$ on a $2\pi\times 2\pi$ domain, which results in a time-
step of $\Delta t=7.0125\times 10^{-3}$ according to the
Courant–Friedrichs–Lewy (CFD) condition on the $64\times 64$ simulation grid.
For training, we generate $128$ trajectories of fluid dynamics, each starting
with different random initial conditions and simulating with $2048\times 2048$
resolution for $40.0$ time units. We use $16$ trajectories for evaluation.

Table 2: Quantitative comparisons with the metric of high-correlation
($\rho>0.8$) duration (w.r.t the reference DNS-$2048\times 2048$
trajectories). All learnable solvers use the $64\times 64$ grids.

Method | Type | High-corr. duration
---|---|---
DNS-$64\times 64$ | Physics | 2.805
DNS-$128\times 128$ | Physics | 3.983
DNS-$256\times 256$ | Physics | 5.386
DNS-$512\times 512$ | Physics | 6.788
DNS-$1024\times 1024$ | Physics | 8.752
1-step-raw-CNN | ML | 4.824
4-step-raw-CNN | ML | 7.517
32-step-FNO | ML | 6.283
1-step-raw-CNN | LC | 6.900
32-step-FNO | LC | 7.630
1-step-raw-CNN | LI | 7.910
32-step-FNO | TSM | 7.798
4-step-raw-CNN | TSM | 8.359
\+ 4-step temporal-bundle | TSM | 8.303
32-step-HiPPO-CNN | TSM | 9.256
\+ 4-step temporal-bundle | TSM | 9.481

#### Unrolled training

All the learnable solvers are trained with the Mean Squared Error (MSE) loss
on the velocities. Following previous work (Li et al., 2020c; Brandstetter et
al., 2021; Kochkov et al., 2021; Dresdner et al., 2022), we adopt the unrolled
training technique, which requires the learnable solvers to mimic the ground-
truth solution for more than one unrolled decoding step:

$\mathcal{L}(\mathbf{u}^{gt}_{[0,T]})=\frac{1}{N}\sum_{i=1}^{N}\mathrm{MSE}(\mathbf{u}^{gt}(t_{i}),\mathbf{u}^{pred}(t_{i}))$
(5)

where $t_{i}\in\\{T+\Delta t,\dots,T+N\Delta t\\}$ is the unrolled time steps,
$\mathbf{u}_{gt}$ and $\mathbf{u}_{pred}$ are the ground-truth solution and
learnable solver’s prediction, respectively. Unrolled training can improve the
inference performance but makes the training less stable (due to bad initial
predictions), so we use $N=32$ unrolling steps for training.

#### Temporal bundling

In our preliminary experiments, we found that due to the abundant information
in the trajectories, the TSM solvers are more prone to over-fit. Therefore, we
adopt the temporal bundling technique (Brandstetter et al., 2021) to the
learned interpolation coefficients in TSM solvers. Assume that in a step-by-
step prediction scheme, we predict
$\mathbf{u}_{[0,T]}\rightarrow\mathbf{c}(T+\Delta t)$, where $\mathbf{c}$ is
the stencil interpolation coefficients in the convective flux approximation.
In temporal bundling, we predict $K$ steps of the interpolation coefficients
$\mathbf{u}_{[0,T]}\rightarrow\\{\mathbf{c}(T+\Delta t),\mathbf{c}(T+2\Delta
t),\dots,\mathbf{c}(T+K\Delta t)\\}$ in advance, and then time-step forward
the FVM physics model for $K$ steps with pre-computed stencil interpolation
coefficients.

#### Neural network architectures & hyper-parameters

In TSM-$64\times 64$ with a $T$-length trajectory, the input and output shapes
of TSM are $64\times 64\times T\times 2$ and $64\times 64\times C\times 2$,
while the input and output shapes of CNN are $64\times 64\times(C\times 2)$
and $64\times 64\times(8\times(4^{2}-1))$, which represents that for each
resolution, we predict 8 interpolations that need 15 inputs each. For
HiPPO333https://github.com/HazyResearch/state-spaces, we set the hyper-
parameters $a=-0.5,b=1.0,dt=1.0$. For CNN, we use a $6$-layer network with
$3\times 3$ kernels and $256$ channels with periodic
padding444https://github.com/google/jax-cfd/blob/main/jax_cfd/ml/towers.py.

###  4.2 Main results

The classic and neural methods we evaluated can be roughly classified into
four categories: 1) pure Physics models, i.e., the FVM-based Direct Numerical
Simulation (DNS), 2) pure Machine Learning (ML)-based neural method-of-lines
models, 3) Learned Correction (LC) models, which correct the final outputs
(i.e., velocities) of physics models with neural networks (Um et al., 2020),
and 4) Neural Stencil Modeling models, including single-time-step Learned
Interpolation (LI) (Kochkov et al., 2021), and our Temporal Stencil Modeling
(TSM).

For neural network baselines, except the
periodic555https://github.com/google/jax-cfd/blob/main/jax_cfd/ml/layers.py
Convolutional Neural Networks (CNN) (LeCun et al., 1999; Kochkov et al., 2021)
with raw and HiPPO features we already described, we also compare with Fourier
Neural Operators (FNO) (Li et al., 2020c) and Multiwavelet-based model (MWT),
two recent state-of-the-art pure-ML PDE solvers based on spectral and wavelet
features.

We evaluate all the solvers based on their Pearson correlation $\rho$ with the
ground-truth (i.e., DNS with the highest $2048\times 2048$ resolution) flows
in terms of the scalar vorticity field
$\omega=\partial_{x}u_{y}-\partial_{y}u_{x}$. Furthermore, to ease comparing
all the different solvers quantitatively, we focus on their high-correlation
duration time, i.e., the duration of time until correlation $\rho$ drops below
0.8.

The comparison between HiPPO-based TSM and 1-time-step raw-feature LI (Kochkov
et al., 2021) is shown in Fig. 5 and Fig. 5 (left). We can see that our HiPPO
feature-based TSM significantly outperforms the previous state-of-the-art ML-
physics model, especially when trained with a 4-step temporal bundling. From
Fig. 5 (middle), we can see that all the learnable solvers can better capture
the high-frequency features with a similar energy spectrum
$E(\mathbf{k})=\frac{1}{2}|\mathbf{u}(\mathbf{k})|^{2}$ pattern as the high-
resolution ground-truth trajectories. From Fig. 5 (right), we can see that
with the help of temporal bundling, HiPPO-TSM can achieve high simulation
accuracy while reducing the inference latency to 80% when compared to the
original LI.

Figure 6: Temporal stencil modeling performance (high-correlation duration)
with different feature types and different initial trajectory steps. Figure 7:
Generalization test results of neural methods trained on Kolmogorov flows
($Re=1000$) and evaluated on (A) decaying flows (starting from $Re=1000$), (B)
more turbulent Kolmogorov flows ($Re=4000$), and (C) $2\times$ larger domain
Kolmogorov flow ($Re=1000$).

###  4.3 Ablation Study

We present a quantitative comparison for more methods in Tab. 1. We can see
that the learned correction models are always better than pure-ML models, and
neural stencil modeling models are always better than learned correction
models. In terms of neural network architectures, we find that under the same
ML-physics framework, raw-feature CNN is always better than FNO, and HiPPO-
feature CNNs are always better than raw-feature CNNs. Finally, when adopting
temporal bundling, we can see that only HiPPO-TSM can benefit from alleviating
the over-fitting problem, while the performance of raw-feature TSM can be hurt
by temporal bundling.

We also study the impact of the trajectory length on raw-feature and HiPPO-
feature TSM models in Fig. 6. Notice that in raw features, the temporal window
size is fixed during unrolling, while in HiPPO features, the temporal window
size is expanded during unrolling decoding. From the figure, we can see that
the raw-feature CNN achieves the best performance with a window size of 4,
while HiPPO features keep increasing the performance, and reach the peak with
32 initial trajectory length.

###  4.4 Generalization tests

We evaluate the generalization ability of our HiPPO-TSM (4-step bundle) model
and LI trained on Kolmogorov flows ($Re=1000$). Specifically, we consider the
following test cases: (A) decaying flows (starting $Re=1000$), (B) more
turbulent Kolmogorov flows ($Re=4000$), and (C) $2\times$ larger domain
Kolmogorov flows ($Re=1000$). Our results are shown in Fig. 7, from which we
can see that HiPPO-TSM achieves consistent improvement over LI. HiPPO-TSM also
achieves competitive performance to DNS-$1024\times 1024$ or DNS-$2048\times
2048$, depending on the ground truth (i.e., highest resolution) being
DNS-$2048\times 2048$ or DNS-$4096\times 4096$).

###  4.5 Results on Navier-Stokes with $32\times 32$ and $16\times 16$ grids

Please refer to Sec. D, Fig. 8, and Tab. 3 in the appendix.

###  4.6 Results on 1D Kuramoto–Sivashinsky (KS) equation

Please refer to Sec. E and Fig. 10-12 in the appendix.

###  4.7 Results on 3D Navier-Stokes equation

Please refer to Sec. F and Fig. 13-16 in the appendix.

##  5 Conclusion & Future Work

In this paper, we propose a novel Temporal Stencil Modeling (TSM) method for
solving time-dependent PDEs in conservation form. TSM can be regarded as the
temporal generalization of classic finite volume solvers such as WENO (Shu,
2003; Gottlieb et al., 2006) and vanilla neural stencil modeling methods
(Kochkov et al., 2021), in that TSM leverages the temporal information from
trajectories, instead of only using the latest states, to approximate the
(convective) flux more accurately. Our empirical evaluation on 2-D
incompressible Navier-Stokes turbulent flow data shows that both the temporal
information and its temporal feature encoding scheme are crucial to achieving
state-of-the-art simulation accuracy. We also show that TSM has a strong
generalization ability on various out-of-distribution turbulent flows.
Finally, we show that the proposed method works well on 1-D
Kuramoto–Sivashinsky (KS) equation and 3-D Navier-Stokes.

For future work, we plan to evaluate our TSM method with non-periodic boundary
conditions. We are also interested in leveraging the Neural Architecture
Search (NAS) technique to automatically find better features and neural
architectures for solving the Navier-Stokes equation in the TSM framework.

## Acknowledgement

We thank the constructive suggestions from Xihaier Luo (BNL) and the anonymous
reviewers. This research used the Perlmutter supercomputer of the National
Energy Research Scientific Computing Center, a DOE Office of Science User
Facility supported by the Office of Science of the U.S. Department of Energy
under Contract No. DE-AC02-05CH11231 using NERSC award NERSC DDR-ERCAP0022110.

## References

  * Alfonsi (2009) Alfonsi, G.  Reynolds-averaged navier–stokes equations for turbulence modeling.  _Applied Mechanics Reviews_ , 62(4), 2009.
  * Anderson & Wendt (1995) Anderson, J. D. and Wendt, J.  _Computational fluid dynamics_ , volume 206.  Springer, 1995.
  * Bar & Sochen (2019) Bar, L. and Sochen, N.  Unsupervised deep learning algorithm for pde-based forward and inverse problems.  _arXiv preprint arXiv:1904.05417_ , 2019.
  * Bar-Sinai et al. (2019) Bar-Sinai, Y., Hoyer, S., Hickey, J., and Brenner, M. P.  Learning data-driven discretizations for partial differential equations.  _Proceedings of the National Academy of Sciences_ , 116(31):15344–15349, 2019.
  * Bhattacharya et al. (2020) Bhattacharya, K., Hosseini, B., Kovachki, N. B., and Stuart, A. M.  Model reduction and neural networks for parametric pdes.  _arXiv preprint arXiv:2005.03180_ , 2020.
  * Boussinesq (1877) Boussinesq, J.  Theorie de l’ecoulement tourbillant.  _Mem. Acad. Sci._ , 23:46, 1877.
  * Brandstetter et al. (2021) Brandstetter, J., Worrall, D. E., and Welling, M.  Message passing neural pde solvers.  In _International Conference on Learning Representations_ , 2021.
  * Bruno et al. (2022) Bruno, O. P., Hesthaven, J. S., and Leibovici, D. V.  Fc-based shock-dynamics solver with neural-network localized artificial-viscosity assignment.  _Journal of Computational Physics: X_ , pp. 100110, 2022.
  * Cao (2021) Cao, S.  Choose a transformer: Fourier or galerkin.  _Advances in Neural Information Processing Systems_ , 34:24924–24940, 2021.
  * Choi & Moin (2012) Choi, H. and Moin, P.  Grid-point requirements for large eddy simulation: Chapman’s estimates revisited.  _Physics of fluids_ , 24(1):011702, 2012.
  * Chorin (1967) Chorin, A. J.  The numerical solution of the navier-stokes equations for an incompressible fluid.  _Bulletin of the American Mathematical Society_ , 73(6):928–931, 1967.
  * Courant et al. (1967) Courant, R., Friedrichs, K., and Lewy, H.  On the partial difference equations of mathematical physics.  _IBM journal of Research and Development_ , 11(2):215–234, 1967.
  * Dresdner et al. (2022) Dresdner, G., Kochkov, D., Norgaard, P., Zepeda-Núñez, L., Smith, J. A., Brenner, M. P., and Hoyer, S.  Learning to correct spectral methods for simulating turbulent flows.  _arXiv preprint arXiv:2207.00556_ , 2022.
  * Frezat et al. (2022) Frezat, H., Sommer, J. L., Fablet, R., Balarac, G., and Lguensat, R.  A posteriori learning for quasi-geostrophic turbulence parametrization.  _arXiv preprint arXiv:2204.03911_ , 2022.
  * Goel et al. (2022) Goel, K., Gu, A., Donahue, C., and Ré, C.  It’s raw! audio generation with state-space models.  _arXiv preprint arXiv:2202.09729_ , 2022.
  * Gottlieb et al. (2006) Gottlieb, S., Mullen, J. S., and Ruuth, S. J.  A fifth order flux implicit weno method.  _Journal of Scientific Computing_ , 27(1):271–287, 2006.
  * Gu et al. (2020) Gu, A., Dao, T., Ermon, S., Rudra, A., and Ré, C.  Hippo: Recurrent memory with optimal polynomial projections.  _Advances in Neural Information Processing Systems_ , 33:1474–1487, 2020.
  * Gu et al. (2021) Gu, A., Goel, K., and Re, C.  Efficiently modeling long sequences with structured state spaces.  In _International Conference on Learning Representations_ , 2021.
  * Gupta et al. (2021) Gupta, G., Xiao, X., and Bogdan, P.  Multiwavelet-based operator learning for differential equations.  _Advances in Neural Information Processing Systems_ , 34:24048–24062, 2021.
  * Jiang & Shu (1996) Jiang, G.-S. and Shu, C.-W.  Efficient implementation of weighted eno schemes.  _Journal of computational physics_ , 126(1):202–228, 1996.
  * Kochkov et al. (2021) Kochkov, D., Smith, J. A., Alieva, A., Wang, Q., Brenner, M. P., and Hoyer, S.  Machine learning–accelerated computational fluid dynamics.  _Proceedings of the National Academy of Sciences_ , 118(21):e2101784118, 2021.
  * Lax (1959) Lax, P.  Systems of conservation laws.  Technical report, LOS ALAMOS NATIONAL LAB NM, 1959.
  * LeCun et al. (1999) LeCun, Y., Haffner, P., Bottou, L., and Bengio, Y.  Object recognition with gradient-based learning.  In _Shape, contour and grouping in computer vision_ , pp. 319–345. Springer, 1999.
  * Lelievre & Stoltz (2016) Lelievre, T. and Stoltz, G.  Partial differential equations and stochastic methods in molecular dynamics.  _Acta Numerica_ , 25:681–880, 2016.
  * Lesieur & Metais (1996) Lesieur, M. and Metais, O.  New trends in large-eddy simulations of turbulence.  _Annual review of fluid mechanics_ , 28(1):45–82, 1996.
  * Li et al. (2020a) Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., and Anandkumar, A.  Neural operator: Graph kernel network for partial differential equations.  _arXiv preprint arXiv:2003.03485_ , 2020a.
  * Li et al. (2020b) Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Stuart, A., Bhattacharya, K., and Anandkumar, A.  Multipole graph neural operator for parametric partial differential equations.  _Advances in Neural Information Processing Systems_ , 33:6755–6766, 2020b.
  * Li et al. (2020c) Li, Z., Kovachki, N. B., Azizzadenesheli, K., Bhattacharya, K., Stuart, A., Anandkumar, A., et al.  Fourier neural operator for parametric partial differential equations.  In _International Conference on Learning Representations_ , 2020c.
  * Liao et al. (2015) Liao, R., Tao, X., Li, R., Ma, Z., and Jia, J.  Video super-resolution via deep draft-ensemble learning.  In _Proceedings of the IEEE international conference on computer vision_ , pp. 531–539, 2015.
  * List et al. (2022) List, B., Chen, L.-W., and Thuerey, N.  Learned turbulence modelling with differentiable fluid solvers.  _arXiv preprint arXiv:2202.06988_ , 2022.
  * Liu et al. (1994) Liu, X.-D., Osher, S., and Chan, T.  Weighted essentially non-oscillatory schemes.  _Journal of computational physics_ , 115(1):200–212, 1994.
  * Lu et al. (2019) Lu, L., Jin, P., and Karniadakis, G. E.  Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators.  _arXiv preprint arXiv:1910.03193_ , 2019.
  * Matthies & Keese (2005) Matthies, H. G. and Keese, A.  Galerkin methods for linear and nonlinear elliptic stochastic partial differential equations.  _Computer methods in applied mechanics and engineering_ , 194(12-16):1295–1331, 2005.
  * McDonough (2007) McDonough, J. M.  Lectures in computational fluid dynamics of incompressible flow: Mathematics, algorithms and implementations.  2007\.
  * Mishra (2019) Mishra, S.  A machine learning framework for data driven acceleration of computations of differential equations.  _Mathematics in Engineering_ , 1(1):118–146, 2019\.
  * Nelsen & Stuart (2021) Nelsen, N. H. and Stuart, A. M.  The random feature model for input-output maps between banach spaces.  _SIAM Journal on Scientific Computing_ , 43(5):A3212–A3243, 2021.
  * Patel et al. (2021) Patel, R. G., Trask, N. A., Wood, M. A., and Cyr, E. C.  A physics-informed operator regression framework for extracting data-driven continuum models.  _Computer Methods in Applied Mechanics and Engineering_ , 373:113500, 2021.
  * Pope & Pope (2000) Pope, S. B. and Pope, S. B.  _Turbulent flows_.  Cambridge university press, 2000.
  * Raissi et al. (2019) Raissi, M., Perdikaris, P., and Karniadakis, G. E.  Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations.  _Journal of Computational physics_ , 378:686–707, 2019\.
  * Rhie & Chow (1983) Rhie, C. M. and Chow, W.-L.  Numerical study of the turbulent flow past an airfoil with trailing edge separation.  _AIAA journal_ , 21(11):1525–1532, 1983.
  * Ronneberger et al. (2015) Ronneberger, O., Fischer, P., and Brox, T.  U-net: Convolutional networks for biomedical image segmentation.  In _International Conference on Medical image computing and computer-assisted intervention_ , pp. 234–241. Springer, 2015.
  * Sallam & Hwang (1984) Sallam, A. M. and Hwang, N. H.  Human red blood cell hemolysis in a turbulent shear flow: contribution of reynolds shear stresses.  _Biorheology_ , 21(6):783–797, 1984.
  * Sanchez-Gonzalez et al. (2020) Sanchez-Gonzalez, A., Godwin, J., Pfaff, T., Ying, R., Leskovec, J., and Battaglia, P.  Learning to simulate complex physics with graph networks.  In _International Conference on Machine Learning_ , pp. 8459–8468. PMLR, 2020.
  * Schiesser (2012) Schiesser, W. E.  _The numerical method of lines: integration of partial differential equations_.  Elsevier, 2012.
  * Shu (2003) Shu, C.-W.  High-order finite difference and finite volume weno schemes and discontinuous galerkin methods for cfd.  _International Journal of Computational Fluid Dynamics_ , 17(2):107–118, 2003.
  * Slotnick et al. (2014) Slotnick, J. P., Khodadoust, A., Alonso, J., Darmofal, D., Gropp, W., Lurie, E., and Mavriplis, D. J.  Cfd vision 2030 study: a path to revolutionary computational aerosciences.  Technical report, 2014.
  * Smagorinsky (1963) Smagorinsky, J.  General circulation experiments with the primitive equations: I. the basic experiment.  _Monthly weather review_ , 91(3):99–164, 1963\.
  * Smith et al. (2020) Smith, J. D., Azizzadenesheli, K., and Ross, Z. E.  Eikonet: Solving the eikonal equation with deep neural networks.  _IEEE Transactions on Geoscience and Remote Sensing_ , 59(12):10685–10696, 2020.
  * Stachenfeld et al. (2021) Stachenfeld, K., Fielding, D. B., Kochkov, D., Cranmer, M., Pfaff, T., Godwin, J., Cui, C., Ho, S., Battaglia, P., and Sanchez-Gonzalez, A.  Learned coarse models for efficient turbulence simulation.  _arXiv preprint arXiv:2112.15275_ , 2021.
  * (50) Takamoto, M., Praditia, T., Leiteritz, R., MacKinlay, D., Alesiani, F., Pflüger, D., and Niepert, M.  Pdebench: An extensive benchmark for scientific machine learning.  In _Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track_.
  * Tay et al. (2020) Tay, Y., Dehghani, M., Abnar, S., Shen, Y., Bahri, D., Pham, P., Rao, J., Yang, L., Ruder, S., and Metzler, D.  Long range arena: A benchmark for efficient transformers.  In _International Conference on Learning Representations_ , 2020.
  * Tran et al. (2021) Tran, A., Mathews, A., Xie, L., and Ong, C. S.  Factorized fourier neural operators.  _arXiv preprint arXiv:2111.13802_ , 2021.
  * Um et al. (2020) Um, K., Brand, R., Fei, Y. R., Holl, P., and Thuerey, N.  Solver-in-the-loop: Learning from differentiable physics to interact with iterative pde-solvers.  _Advances in Neural Information Processing Systems_ , 33:6111–6122, 2020.
  * Wang et al. (2022) Wang, C., Li, S., He, D., and Wang, L.  Is l2 physics-informed loss always suitable for training physics-informed neural network?  _arXiv preprint arXiv:2206.02016_ , 2022.
  * Wang et al. (2020) Wang, R., Kashinath, K., Mustafa, M., Albert, A., and Yu, R.  Towards physics-informed deep learning for turbulent flow prediction.  In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_, pp. 1457–1466, 2020.
  * Weinan & Yu (2018) Weinan, E. and Yu, B.  The deep ritz method: A deep learning-based numerical algorithm for solving variational problems.  _Communications in Mathematics and Statistics_ , 1(6):1–12, 2018.

##  Appendix A Additional related work

###  A.1 Industrial Solvers for Computational Fluid Dynamics

Industrial CFD typically relies on either Reynolds-averaged Navier-Stokes
(RANS) models (Boussinesq, 1877; Alfonsi, 2009), where the fluctuations are
expressed as a function of the eddy viscosity, or coarsely-resolved Large-Eddy
Simulation (LES) (Smagorinsky, 1963; Lesieur & Metais, 1996), where only the
large scales are numerically simulated, and the small ones are modeled (with
an a-priori physics assumption). However, there are severe limits in both
methods associated with their usage for general purposes. For example, RANS is
too simple to model complex flows (Slotnick et al., 2014) and often exhibits
significant errors when dealing with complex pressure-gradient distributions
and complicated geometries. LES can exhibit limited accuracy in their
predictions of high-$Re$ turbulent flows, due to the first-order dependence of
LES on the subgrid-scale (SGS) model.

##  Appendix B Infrastructures

We train and evaluate all the classic and neural Navier-Stokes solvers in 8
Nvidia Tesla V100-32G GPUs. The inference latency is measured by unrolling 2
trajectories for 25.0 simulation time on a single V100 GPU.

##  Appendix C HiPPO features for Temporal Stencil Modeling

For scalar time series $u_{\leq t}:=u(x)|_{x\leq t}$, the HiPPO (Gu et al.,
2020) projection aims to find a coefficient mapping for orthogonal polynomials
such that

$c(t)\in\mathbb{R}^{N}=\arg\min\|u_{\leq t}-g^{(t)}\|_{\mu(t)}$ (6)

where $g^{(t)}=\sum_{n=1}^{N}c_{n}(t)g_{n}(t)$ and $g_{n}(t)$ is the $n^{th}$
Legendre polynomial scaled to the $[0,t]$ domain.
$\mu(t)=\frac{1}{t}\mathbb{I}_{[0,t]}$ is the uniform weight metric. By
solving the corresponding ODE and its discretization, Gu et al. (2020) showed
that the optimal polynomial coefficients $c_{T}$ can be calculated by the
following recurrence:

$\displaystyle c({T+\Delta t})=\left(1-\frac{A}{T/\Delta
t}\right)c(T)+\frac{1}{T/\Delta t}B\cdot u(T)$ (7)

where

$\displaystyle A_{nk}=\begin{cases}(2n+1)^{1/2}(2k+1)^{1/2}&\mbox{if }n>k\\\
n+1&\mbox{if }n=k\\\ 0&\mbox{if }n<k\end{cases},\qquad
B_{n}=(2n+1)^{\frac{1}{2}}$

We refer the readers to (Gu et al., 2020) for the detailed derivations. When
dealing with the multi-variate time series such as temporal PDE, we follow the
original recipe and treat each scalar component independently, that is, we
treat the $H\times W\times T\times 2$ -shaped velocity trajectory as $H\times
W\times 2$ separate time series of length $T$. Finally, we concatenate the
resulted $H\times W\times 2\times C$ features on the velocity dimension (i.e.,
$H\times W\times 2C$) as the input to the CNN.

##  Appendix D Solving Navier-Stokes with $32\times 32$ and $16\times 16$
grids

Since TSM shows significant improvements over vanilla neural stencil modeling
on the $64\times 64$ grid, we are wondering whether TSM can further push the
Pareto frontier by accurately approximating the convective flux in lower
resolution grid.

Specifically, we still use DNS-$2048\times 2048$ as the ground-truth training
data, but train the TSM solver in the $32\times 32$ and $16\times 16$ down-
sampled grids. In order to make a fair and direction comparison to our
original $64\times 64$ model, we additionally train two $32\times
32\rightarrow 64\times 64$ and $16\times 16\rightarrow 64\times 64$ super-
resolution model, and evaluate (on $64\times 64$ vorticity) their super-
resolved results.

We report the results in Fig. 8 and Tab. 3. We can see that the super-
resolution models actually work quite well for $32\times 32\rightarrow
64\times 64$ and $16\times 16\rightarrow 64\times 64$. Besides, though the
lower-resolution neural solvers’ performance is significantly worse than
$64\times 64$, we can see that the improvement from temporal information in
the trajectories is more significant in lower resolutions.

Figure 8: Qualitative results with TSM evaluating on $32\times 32$ and $16\times 16$, and their $64\times 64$ super-resolution results. From the super-resolution results of the $T=0$, we can see that the super-resolution models actually work quite well. Table 3: Quantitative comparisons with the metric of high-correlation ($\rho>0.8$) duration (w.r.t the reference DNS-$2048\times 2048$ trajectories). The results of TSM-$32\times 32$ and TSM-$16\times 16$ are evaluated on their corresponding $64\times 64$ super-resolution results with additionally trained super-resolution models. Method | High-corr. duration
---|---
DNS-$64\times 64$ |  2.805
DNS-$128\times 128$ |  3.983
DNS-$256\times 256$ |  5.386
DNS-$512\times 512$ |  6.788
DNS-$1024\times 1024$ |  8.752
LI-$64\times 64$ |  7.910
TSM-$64\times 64$ |  9.481 (+19.86%)
LI-$32\times 32$ |  5.400
TSM-$32\times 32$ |  6.802 (+25.96%)
LI-$16\times 16$ |  2.805
TSM-$16\times 16$ |  3.576 (+27.50%)
Figure 9: Illustration of using HiPPO features as the CNN inputs for neural stencil modeling. After encoding the HiPPO features for the initial velocity trajectory $\mathbf{v}_{[0,T]}$, given model-generated new velocity, we only need an additional recurrent update step to get the HiPPO features for $\mathbf{v}_{[0,T+\Delta t]}$ (with Eq. 7). Notice that in TSM, we use the fixed optimal HiPPO recurrence for temporal encoding and only the CNN part is learnable.  Table 4: The $p$-values in one-sample T-test for the differences between TSM and other baseline models. The differences in all 16 test trajectories are used for each significance test. We evaluate the significance for four high-correlation thresholds: 0.95, 0.9, 0.8, and 0.7. |  $p$-value in one-sample T-test
---|---
Baseline | $\rho>0.95$ | $\rho>0.9$ | $\rho>0.8$ | $\rho>0.7$
DNS 64x64 | $1.06\times 10^{-11}$ | $1.14\times 10^{-10}$ | $1.20\times 10^{-10}$ | $1.26\times 10^{-10}$
DNS 128x128 | $8.80\times 10^{-11}$ | $2.23\times 10^{-10}$ | $6.75\times 10^{-10}$ | $6.63\times 10^{-10}$
DNS 256x256 | $1.22\times 10^{-10}$ | $2.39\times 10^{-09}$ | $4.78\times 10^{-09}$ | $1.56\times 10^{-09}$
DNS 512x512 | $2.38\times 10^{-09}$ | $1.49\times 10^{-07}$ | $2.27\times 10^{-06}$ | $1.65\times 10^{-06}$
DNS 1024x1024 | $6.63\times 10^{-03}$ | $6.91\times 10^{-03}$ | $1.53\times 10^{-02}$ | $1.65\times 10^{-02}$
LI 64x64 (Kochkov et al. 2021) | $9.06\times 10^{-04}$ | $1.65\times 10^{-03}$ | $3.63\times 10^{-03}$ | $3.08\times 10^{-03}$

Figure 10: Qualitative comparison of TSM and other baselines on 1D KS
equation. The solutions are down-sampled to $32$ grid in the space dimension
for comparison. The dashed vertical yellow line denotes the time step where
the Pearson correlation between the model prediction and ground truth (i.e.,
DNS-1024) is lower than the threshold ($\rho<0.8$). Figure 11: Qualitative
comparison of TSM and other baselines on 1D KS equation. The solutions are
down-sampled to $64$ grid in the space dimension for comparison. The dashed
vertical yellow line denotes the time step where the Pearson correlation
between the model prediction and ground truth (i.e., DNS-1024) is lower than
the threshold ($\rho<0.8$). Figure 12: Quantitative comparison of TSM and
other baselines on the velocity correlation of 1D KS equation. The solutions
are down-sampled to $64$ grid in the space dimension for comparison. Figure
13: Quantitative comparison of TSM and other baselines on the vorticity
correlation (averaged over three directions) of 3D incompressible Navier-
Stokes equation.

##  Appendix E Generalization on 1D Kuramoto–Sivashinsky (KS) equation

To verify the generalization ability of TSM, we further evaluate it on 1D
equations. Following previous work (Bar-Sinai et al., 2019; Stachenfeld et
al., 2021), we choose Kuramoto–Sivashinsky (KS) equation as a representative
1D PDE that can generate unstable and chaotic dynamics. While KS-1D is not
technically turbulent, it is a well-studied chaotic equation that can be used
to assess the generalization ability of our models in 1D cases.

Specifically, the KS equation can be written in the conservation form of

$\frac{\partial v}{\partial t}+\frac{\partial J}{\partial x}=0,\quad\quad
v(x,t=0)=v_{0}(x)$ (8)

where

$J=\frac{v^{2}}{2}+\frac{\partial v}{\partial x}+\frac{\partial^{3}v}{\partial
x^{3}}$ (9)

Following previous work (Bar-Sinai et al., 2019), we consider the 1D KS
equation with periodic boundaries. The domain size is set to $L=20\pi$, and
the initial condition is set to

$v_{0}(x)=\sum_{i=1}^{N}A_{i}\sin\left({2\pi\ell_{i}x/L+\phi_{i}}\right)$ (10)

where $N=10$, and $A$, $\phi$, $\ell$ are sampled from the uniform
distributions of $[-0.5,0.5]$, $[-\pi,\pi]$, and $\\{1,2,3\\}$, respectively.

Similar to the NS solution, we solve the nonlinear convection term by
advecting all velocity components simultaneously, using a high-order scheme
based on the Van-Leer flux limiter or with a learnable interpolator, while the
second and fourth order diffusion are approximated using a second-order
central difference approximation. The Fourier spectral method is not used
because of the precision issues in JAX FFT and IFFT
666https://github.com/google/jax/issues/2952.

We use DNS-1024 as the ground-truth training data, and train the TSM solver
and LI solver in the $32$ and $64$ down-sampled grids. A time-step of $\Delta
t=1.9635\times 10^{-2}$ and $\Delta t=9.81748\times 10^{-3}$ is used for $32$
and $64$ grids, respectively. Following Bar-Sinai et al. (2019), the
interpolation coefficients of a 6-point stencil are predicted by TSM and LI.
For training, we generate $1024$ trajectories of fluid dynamics, each starting
with different random initial conditions and simulating with $1024$ resolution
for $200.0$ time units (after $80.0$ time unit warmup). We use $16$
trajectories for evaluation.

When comparing with other DNS baselines, all solutions are down-sampled to
$32$ or $64$ for comparison. Fig. 10 and Fig. 11 show the results of TSM and
LI compared with other baselines in the $32$ and $64$ grids, respectively. A
quantitative comparison of various solvers under $64$-resolution is also
presented in Fig. 12. We can see that TSM can outperform LI with the same
resolution, and outperform DNS with $4\times\sim 8\times$ resolutions.

##  Appendix F Generalization on 3D Navier-Stokes (NS) equation

To verify the generalization ability of TSM, we further evaluate it on 3D
Navier Stokes equations. Recall that the incompressible, constant density
Navier Stokes (NS) equation for fluids has a conservation form of:

$\partial_{t}\mathbf{u}+\nabla\cdot(\mathbf{u}\otimes\mathbf{u})=\nu\nabla^{2}\mathbf{u}-\frac{1}{\rho}\nabla
p+\mathbf{f}$ (11)

We train our method with 3D incompressible Navier-Stokes flow with a linear
forcing $\mathbf{f}=0.05\mathbf{u}$ to avoid the decaying of
flows777https://github.com/google/jax-
cfd/blob/main/jax_cfd/ml/physics_configs/linear_forcing.gin. The viscosity of
the fluid is set to $\nu=6.65\times 10^{-4}$ and the density $\rho=1$. For
training, we generate $128$ trajectories of fluid dynamics, each starting with
different random initial conditions and simulating with $128\times 128\times
128$ resolution for $10.0$ time units (after $80.0$ time unit warmup). We use
$16$ trajectories for evaluation.

Following previous work (Stachenfeld et al., 2021; Takamoto et al., ), the
ground-truth solution is obtained by simulation on $128\times 128\times 128$
grid. According to the Courant–Friedrichs–Lewy (CFD) condition, we have time-
step $\Delta=1.402497\times 10^{-2}$ on the $32\times 32\times 32$ simulation
grid. For TSM-$32\times 32\times 32$ and LI-$32\times 32\times 32$, most of
the settings in 3D NS follow our setup for the 2D case, except that the
interpolation coefficients are only calculated for a $2\times 2\times 2$-point
stencil to avoid out-of-memory issues.

When comparing with other DNS baselines, all solutions are down-sampled to
$32\times 32\times 32$ for comparison. A quantitative comparison of various
solvers under $32\times 32\times 32$-resolution is presented in Fig. 13. The
qualitative results on the planes of $x=0$, $y=0$, $z=0$ are presented in Fig.
14, Fig. 15, and Fig. 16, respectively. We can see that TSM can outperform LI
and DNS with the same resolution, but cannot beat DNS with $2\times$ higher
resolution. This is consistent with the results reported in previous work
(Stachenfeld et al., 2021).

Figure 14: Qualitative 3D incompressible Navier-Stokes equation results of
predicted vorticity fields on the plane of $x=0$. Figure 15: Qualitative 3D
incompressible Navier-Stokes equation results of predicted vorticity fields on
the plane of $y=0$. Figure 16: Qualitative 3D incompressible Navier-Stokes
equation results of predicted vorticity fields on the plane of $z=0$.