Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
Authors: Yunfei Huang
David S. Greenberg
Venue: ICML 2025 poster (2025)
Keywords: Geometric and Physical Constraints
Neural PDE Surrogates
Symmetry Equivariance
Conservation Laws
Fluid Dynamics
Primary area: applications->chemistry_physics_and_earth_sciences

Abstract: Neural PDE surrogates can improve the cost-accuracy tradeoff of classical solvers, but often generalize poorly to new initial conditions and accumulate errors over time. Physical and symmetry constraints have shown promise in closing this performance gap, but existing techniques for imposing these inductive biases are incompatible with the staggered grids commonly used in computational fluid dynamics. Here we introduce novel input and output layers that respect physical laws and symmetries on the staggered grids, and for the first time systematically investigate how these constraints, individually and in combination, affect the accuracy of PDE surrogates. We focus on two challenging problems: shallow water equations with closed boundaries and decaying incompressible turbulence. Compared to strong baselines, symmetries and physical constraints consistently improve performance across tasks, architectures, autoregressive prediction steps, accuracy measures, and network sizes. Symmetries are more effective than physical constraints, but surrogates with both performed best, even compared to baselines with data augmentation or pushforward training, while themselves benefiting from the pushforward trick. Doubly-constrained surrogates also generalize better to initial conditions and durations beyond the range of the training data, and more accurately predict real-world ocean currents.

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
Yunfei Huang 12
David S. Greenberg' !2
Abstract
1. Introduction
Recently, neural networks have shown promising results in predicting the time evolution of PDE systems, often achieving cost-accuracy tradeoff's that outperform traditional numerical methods (Li et al., 2021; Gupta & Brandstet- ter, 2023; Stachenfeld et al., 2021; Takamoto et al., 2022; Long et al., 2019; Um et al., 2020; Kochkov et al., 2021). However, obtaining accurate and stable autoregressive 'roll- outs' over long durations remains notoriously difficult. Sev- eral techniques have been proposed to address this prob- lem, including physical constraints, symmetry equivariance, time-unrolled training, specialized architectures, data aug- mentation, addition of input noise and generative modeling (Sanchez-Gonzalez et al., 2020; Lippe et al., 2024; Stachen- feld et al., 2021; Kohl et al., 2024; Brandstetter al., 2022b; Fanaskov al., 2023; Bergamin et al., 2024; Sun al., 2023; Hsieh et al., 2019; Tran et al., 2023; Li al., 2023; Bonev et al., 2023). Nonetheless, the relative effectiveness of these strategies remains largely ambiguous, and transpar- ent, systematic comparisons remain elusive.
Neural PDE surrogates can improve the cost- accuracy tradeoff of classical solvers, but often generalize poorly to new initial conditions and accumulate errors over time. Physical and sym- metry constraints have shown promise in closing this performance gap, but existing techniques for imposing these inductive biases are incompati- ble with the staggered grids commonly used in computational fluid dynamics. Here intro- duce novel input and output layers that respect physical laws and symmetries on the staggered grids, and for the first time systematically inves- tigate how these constraints, individually and in combination, affect the accuracy of PDE surro- gates. We focus on two challenging problems: shallow water equations with closed boundaries and decaying incompressible turbulence. Com- pared to strong baselines, symmetries and physi- cal constraints consistently improve performance across tasks, architectures, autoregressive predic- tion steps, accuracy measures, and network sizes. Symmetries are more effective than physical con- straints, but surrogates with both performed best, even compared to baselines with data augmenta- tion pushforward training, while themselves benefiting from the pushforward trick. Doubly- constrained surrogates also generalize better to initial conditions and durations beyond the range of the training data, and more accurately predict real-world ocean currents.
Here, we systematically investigate the utility of symme- try constraints and physical conservation laws, alone and in combination. While both have proven useful for some tasks and architectures, date there have been practically systematic evaluations of their combination. Given the deep connections between conservation laws and PDE sym- metries in physics (Noether, 1918), it is not clear a priori whether these constraints would prove redundant, or com- bine usefully for training PDE surrogates. Across multiple tasks, accuracy measures, architectures, training techniques, and scenarios, we show clear, reproducible and robust benefit from these constraints for long rollout accuracy and generalization performance, and that they combine syner- gistically. To make them broadly applicable, we introduce novel input and output layers that extend these inductive biases to staggered grids for the first time.!
'Helmholtz Centre Hereon, Geesthacht, Germany 'Helmholtz AI. Correspondence to:
Yunfei Huang Greenberg
Code is available at https: //tpttgtttttt///////////t/.hht con/m /mmmmmau dml/ double cowstraint -de--pdrgoaees.....
<yunfei. huang5 @ gmail.com.>> edavid.grreebberg @ de>.
David
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
Input data staggered C-grid
Input layer
Hidden layers
Output layer equivariance physical constraints
Output data staggered C-grid
equivariant filters
regular representation collocated grid
W
W
W
W
Modern U-Net escnn
Physical constraints
M+ pi
9
1
got
1]
U
Ramo W
a1-
Symmetry constraints: (u', ) [8' ] "ol,,,")TT,,
_
K)
Figure I. Symmetry- and chysics-constraine neural surrogate for incompressible flow on a staggered grid. A rotation- equiivrriai input layer maps velocities onto non -on-staggered regular representation, hidden layers employ steerable convolutions and the equivariant output layer enforces conservation laws on mass and momentum (M + pū) as maps to staggered velocities.
2. Background and Related Work
should exist transformations Tg. operating w, re- spectively, such that
Neural PDE surrogates We aim to train neural networks to predict the time evolution of a system of PDEs. We consider time-dependent variable fields wt,,x) E R"" for E E C Rd, 1. E [0, [) and
[f (x) ITg fwl(2) Vg e G. e Q
(3)
That is, transforming the inputs of f will transform its out- puts correspondingly. When w a scalar field, T and T simply resample at coordinates defined by the action of G on Rd:
Ow
Ôt = 2 Vw, Viw w.....
(1)
Starting from initial conditions (ICs) w (0, x) and boundary conditions (BCs) B(w) (t, x) 0, Vz E in, the solution can be advanced with fixed time step:
(T (xx) = w/g 'x)))
(4)
Other field types transform in complex ways. For example, the action of 90" rotation I on a 2D vector field both resamples the field and rotates each vector:
w(t + -) g[w((t,)) )}
where is an update operator. To provide training data and evaluate performance we use a reference solution gener- ated by numerical solver with space- and time-discretized variable fields.
[TR (wi,
T (Wi, w2) (xx) (-w2(Rt 2)) wi (R (R- 'w) (5)
The range of possible actions is described by G's group representations. Efficient, full-featured software packages exist for equivariant convolutions (Cesa et al., 2022) and self- attention (Romero & Cordonnier, 2021), and have proven useful in image classification (Chidester et al., 2019) and segmentation (Veeling et al., 2018), and to improve neural PDE surrogates (Wang et al., 2021; Helwig et al., 2023; Smets et al., 2023; Huang & Greenberg. 2023; Ruhe et al., 2024). Numerical integration methods can also benefit from maintaining PDE symmetries (Rebelo & Valiquette, 2013). We restrict ourselves to the discrete symmetry groups that hold precisely regular grids, though some approaches have been proposed for the continuous symmetries that hold on the original PDEs (Weiler & Cesa, 2019; Thomas et al., 2018; Cesa et al., 2022; Knigge et al., 2024; Horie & Mitsume, 2022; Brandstetter et al., 2022a; Gasteiger et al., 2020; Lino et al., 2022; Toshev et al., 2023).
Recent studies have trained neural surrogates to approximate G (Greenfeld et al., 2019; Gupta & Brandstetter, 2023; List et al., 2024; Lippe et al., 2024: Li et al., 2021; Tripura & Chakraborty, 2023; Raonic et al., 2024). The neural network can also be combined with numerical solver, in so-called "hybrid methods' (Bar-Sinai et al., 2019; Tompson et al., 2017: Kochkov et al., 2021; Bukka et al., 2021; Long et al., 2019).
Training neural surrogates to remain stable and accurate over long autoregressive rollouts remains challenging. Sev- eral techniques have been proposed, including physical con- straints, symmetry constraints, training with input noise, unrolled training and generative modeling. However, clear consensus on the relative effectiieness of these approaches remains elusive, and their application to new tasks is not always straightforward
Non-equivariant architectures can also be trained using data augmentatiin, with g G randomly sampled to transform input-target pairs during training (Brandstetter et al., 2022b). But while equivariant surrogates maintain symmetries pre- cisely for any training or testing data, data augmentation results in imprecise equivariance that may fail to generalize beyond the training data.
Symmetry equivariance Suppose / w - is an oper- ator mapping between two multidimensional variable fields (xx), x(x) defined on Q C Rdd. Then for a group G of invertible transformations on Rd is equivariant if it com- mutes with the actions of G on and z. Concretely, there
2

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
Original Input
2t
W
W
Mi
yi
y/2
3/
Mn W uu' W o
R
81 RooW" of Roo on
Ro
Roc
92
-Riso
Riso
85
Remo xnd BrioW
Rotated input
R(v")
R((u))
Riso W
Ryso W
Riyah
(R(p))
R(a))
R(ye)
M6 W + R(v") - W" R((e)) R(ys)
Ramo
Rero R W
Vi Roo Roo R((") R(gb)
W2 -Rw."" R((e") + R'sW" R(u') R(n))
Input layer
96 RR R((²)- Rem R(v)) R(ya)
Figure 2. Action of inattin-equivariinn input layer on staggered velocity fields (top left). The filter bank is transformed by each g E G to compute G-indexed regular representation y. Rotation- trantionntansformnng inputs (bottom left) yiclds permuted, rotated output channels.
Staggered grids Fluid dynamical systems are often simu- lated using staggered grids (Fig. 8), in which variables such as pressure, density, divergence or velocity along each axis are represented at different locations. This approach can avoid grid-scale numerical artifacts in numerical integration, and is widely used in fluid dynamics (Holl& & Thuerey, 2024; Kochkov et al., 2021; Jasak, 2009; Stone et al., 2020) as well as atmospheric (Jungclaus et al., 2022; De Pondeca et al., 2011; Korn et al., 2022), and ocean models (Korn et al., 2022; Madec et al., 2023a;b). Staggered grids are common in in finite volume solvers, which generally respect conser- vation laws and offer other numerical advantages (Ferziger et al., 2019). However, current software implementations of equivariant network layers (Cesa et al., 2022; Romero & Cordonnier. 2021) cannot be applied to PDE variable fields on equivariant grids. This is because they assume that variables fields are all located at the same points, allowing the action of symmetries on these fields to be broken into two steps: resampling step z - g 'z carried out on the grid itself, and transformation step - Py(w) carried out on PDE field variables w E I at each point single grid point. This leads to overall transformations w - Ty (w 1))). This is a valid assump- tion for PDEs on continuous spatial domains (eq. 5) or for collocated grids (eq. 1 of Weiler & Cesa 2019). But for stag- vector of values w(x) at each grid point x. Instead, each field is defined at different locations, which may be grid cell cen- ters, interfaces or vertices. Thus, the spatial transformation of the grid and the transformation of local field values can- not be disentangled. Applying existing equivariant network layers to staggered PDE fields therefore breaks symmetry constraints.
and generalization capabilities, these laws can be imposed through additional loss terms (Read et al., 2019; Wang et al., 2021; Stachenfeld et al., 2021; Sorourifar al., 2023). Tak- ing the strategy of physics-derived loss terms to its ulti- mate limit, one arrives unsupervised training on PDE- derived losses for discretized (Wandel et al., 2021; Michelis & Katzschmann, 2022) or continuous solutions (Raissi et al., 2019). Alternatively, one can reparameterize network out- puts to respect hard constraints (Mohan et al., 2020; Beucler et al., 2021; Chalapathi et al., 2024; Cranmer et al., 2020; Greydanus al., 2019). Here we focus on constrrnnd supervised learning with space- and time-discretized grids, which has proven more competitive in larger and more com- plex PDE systems (Takamoto et al., 2022).
3. Symmetry- and Plysics-Constrained Neural Surrogates
In this work, we assess the separate and combined bene- fits of symmetries and conservation laws for neural PDE surrogates. To achieve this, we construct equivariant input layers that support staggered grids (Fig. 8), as well as output layers enforcing both equivariance and conservation laws. When comparing to non-equivariant networks, we replace equivariant convolutions using standard convolutions with the same size and padding options (Appendix B), adjusting channel width to match total parameter counts.
such that (z)
gered grids, the PDE fields are not represented
Fig. demonstrates our overall framework for constructing equivariant, conservative neural surrogates. As an illustra- tive example, we show the incompressible Navier Stokes equations, with equivariance in translation and rotation, momentum conservation and a divergence- freggence-freeeeee condition (equivalent to mass conservation). Input data defined on staggered grids are mapped through novel equivariant input layers to a set of convolutional output channels defined at grid cell centers. Each internal activations consist of regu- lar representations: groups of channels indexed by G,² 2 on which G acts by transforming each spatial field and by per- muting the channels according to the group action (Cohen & Welling, 2016; Cesa al., 2022). Essentially, regular repre- sentations are real-valued functions of the discrete symmetry
Unfortunately, existing equivariant network layers (Cesa et al., 2022; Romero & Cordonnier, 2021) assume T, can be described by a resampling operation followed by an independent transformation at each grid point as in Equation 5, but on staggered grids rotation and reflection do not take this form.
Physical constraints Neural surrogates have frequently been applied to physical systems, many of which contain known conservation laws. To improve accuracy, stability,
Technically, by the non non-translational subgroup of G.
3

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
group G. This formulation allows use the preexist- ing library escnn (Cesa al., 2022) for all internal linear transformations between hidden layers. Finally, we employ novel output layers to map the regular representation back to the staggered grid while enforcing conservation laws as hard constraints.
constraints imposed on the weights to satisfy the equivari- ant transformation of vector fields (Eq. 5, details of output layers and their proofs in D). Vertex-centered scalar outputs are computed using even-sized square filters, followed by pooling layers operating over G-indexed channels.
Conservation laws We impose 3 types of conservation laws as hard constraints. For scalar quantities such as fluid surface height Ç. we subtract the global mean of Ç*+1 ç at each time step. For vector fields, we subtract the mean of each velocity component. As mass conservation in incompressible flows is equivalent to divergence- velocity fields, we impose this by learning vector potential a defined at grid vertices, and compute velocities at grid cell interfaces as the curl V x a to satisfy both mass and momentum conservation (Wandel et al., 2021). Further details and discussion of alternative approaches are found in appendix F.
Input layers We consider input data on staggered Arakawa C-grids (Fig. 8b) with square cells, and variable fields defined at cell centers (typically scalar fields such as pressure, surface height, or divergence), at cell interface midpoints (such as velocity components), and at vertices (such as vector potentials). For a n 2D grid of cells, there is (n + 1) X n grid of interfaces in the iy direction (along rows, including boundaries), and a n x (n + 1) grid of interfaces in the 22 direction (along columns).
We designed convolutional input layers to take scalar inputs at cell centers and/or vector fields with components defined at interfaces. Inputs at interfaces are first processed with bank of convolutional filters, each of even size along the coordinate axis orthogonal to single set of interfaces, and of odd size along all other axes (Fig. 2, left). This filter bank is collectively transformed according to each element of the symmetry group G, while being applied to the input data, Note that, similar to the transformation of vector fields (Eq. 5), these filter banks undergo collective transformation by rotations and reflections, not only through resampling, but also through permutation and sign flips (Fig. 2, right). When we rotation-transform input vector fields (Eq. 5), this has the effect of permuting and rotating the outputs of our input layer, as required for an equivariant mapping onto a regular representation (Fig. 2, magenta arrows), which the proof can be found in the Appendix C.2. Inputs at cell centers are processed with separate, standard equivariant convolution layers. Convolutions for both interface and center-defined input variables produce regular representa- tion outputs, which are then combined to compute the total input to the network's 's first hidden layer. We provide imple- mentations of 2D input layers for translation-rotation (p4) and cetriiio-rrottioo-eeeee
Table I. Geometric and physical constraints for SWES
Symmetrics
Conservation laws
n
!
None Ø Mass M
p1/O
p4/O
p4m/0
p1/M
P4/M
p4m/M
Table 2. Geometric and physical constraints for INS
Symmetries
Conservation laws
ii
None Ø Momentum pi Mass/nomentum M+
p1/O
p4/2
p4m/0
pl/pi
p4/pä
p4m/pü
p1/M+pū
p4/M+pū
p4m/M+pü
3.1. Surrogate Architectures
To measure the efficayy of symmetries and physical con- straints we chose flexible base architecture with efficient training and inference that has produced highly competi- tive results: the "modern U-net" (Gupta & Brandstetter, 2023), which modifies the original U-net (Ronneberger et al., 2015) for improved performance as a PDE surrogate (Appendix H). This architecture has shown strong results in (Kohl al., 2024), and similar version performed well in (Lippe al., 2024). We used it without self-attention layers, which did not significantly affect our results. When constructing symmetry-respecting versions of the U-net, we confirmed equivariance held to numerical precision, but only if input/output layers for staggered grids were used (Appendix B).
(p4m). Further details on input layers can be found in the Appendix C.
Output layers We designed convolutional output layers mapping from regular representations to staggered C-grid variables (Fig. 8). As for the input layers, we use separate convolutional filter banks for cell- and interface-centered variables, but now additionally support vertex-centered scalar outputs for the purpose of enforcing physical con- straints (see below). Scalar face-centered outputs are com- puted using pooling layers over a regular representation (Cohen& & Welling, 2016). Vector field outputs at each cell interface are computed as linear combinations of regular representations at the two surrounding cell centers, with
In some experiments, we also compared additional base- lines. The Dilated ResNet architecture (drnet) has has per- formed well as a PDE surrogate (Stachenfeld et al., 2021). and we considered constrained and unconstrained versions. The rotation- -ouation-euiivariant U-net of Wang et al., 2021 was also enforces hard symmetry constraints, but was not intended for staggered grids. Fourier neural operators (FNOs) com-
4

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
bine local operations with filtering in frequency space (Li al., 2021), and we unconstrained versions.
4.1. Closed Shallow Water System
The shallow water equations (SWEs) are widely used to describe quasi-static motion in homogeneous incom- pressible fluid with free surface. We consider nonlinear SWEs in momentum- and mass conservative form on the domain Ω with 'closed' Dirichlet BCs (Song et al., 2018):
3.2. Training
We trained neural surrogates using MSE loss C where N' is the number of discretized 1 18++1 - [| 11
N
PDE field values. All data fields were normalized by sub- tracting the mean and dividing by the standard deviation, with common values for both components of vector fields. We trained on 2 A100 GPUs with the ADAM optimizer (Kingma, 2014), batch size 32 and initial learning rate le-4. We employed early stopping when validation loss did not reduce for 10 epochs, and accepted network weights with the best validation loss throughout the training process.
1 = CD Fuuu// gVG + an V
(6)
a V (hrs)
(7)
dr
u = 0,
on O
(8)
where ( is fluid surface elevation, u = [u, v] is the velocity field, and h respectively represent the undisturbed- and disturbed fluid depth (so that h = d + () and àl a closed domain boundary. an is the horizontal turbulent momentum exchange coefficient, Cn is the bottom drag coefficient and g is gravitational acceleration. SWEs simulations exhibit travelling waves that reflect from domain boundaries, tem- porarily increasing in height as they self-collide. This sys- tem is more challenging than previous SWE tasks with open (Takamoto et al., 2022) or periodic BCs (Gupta & Brandstet- ter, 2023), due to the combination of self-interfering wave patterns, incompressibility and altered dynamics pixels near the domain boundaries.
0.2h
2h
4h
10h
25h
50h
pt/0 reference
p4m//
1.0
d
NRMSE-
Logiloss)
Numerical reference solution Closed BCs and incom- pressibility lead to stiff dynamics, so explicit solvers are inefficient. Instead, we generate data using a semi-implicit scheme (Backhaus, 1983) that represents ( and [u, e| on a staggered Arakawa C-grid (Arakawa, 1977) and solves a sparse linear system at each time step At 300s.
((b)
50
-0.20
((h)
50
-110
Epochs 600
18
1.00x102
36
High-correla tin
Counts
Mass
right correlation-O.8
Grids are 100 100, 100 x 99, and 99 x 100 respectively for Ç. u, and U. We trained on 50 simulations spanning 50 h (600 time steps) each. ICs were Ç = 0 except for a 0.1 m high square-shaped elevation, and [u, u] = 0. The square had side length uniformly distributed from 2-28 grid cells and random position. The solver was implemented in Fortran and required 67 s/IC a compute node with 48 CPUs. Testing and validation data included 10 simu- lations. Surrogates used the solver's time step. Since the time evolution of this SWE systems depends on the location of boundaries, we provide a binary boundary mask to the network as an additional input field with scalar values de- fined grid cell centers. We note that this binary mask is invariant to rotations and reflections.
pin/o
0.99
E.NRMSE <6U0
$110
ref
p1/M
p4/2
p4mf2 p4m/M
FNON dmmaei
-470
$1/00
PTAH
p4/M
Figure 3. p4m/M (symmetry-physsis constraints) outperforms other networks with similar parameter counts on SWEs. (a) Refer- ence surface disturbance with predictions from pl/Ø and p4m/M. (b-c) Accuracy over 50h rollouts, with standard error of the mean over 20 ICs. (d) Training loss over iterations. (e) Histogram of E, E.NRMSE over 20 ICs. (f) Violation of mass conservation for all methods (black line shows reference simulation). (g) High correlation times for each model.
4. PDE Systems
We considered two challenging 2D fluid dynamic PDEs, with the same staggered grid and symmetries but different variables, BCs/ICs, reference solvers and conservation laws. Full sets of constraints for each system and names for each combination appear in Tables (1-2), while PDE parameters and further numerical details appear in Tables (6-7).
Symmetries and conservation laws The shallow water system in Eqs. 7-8 is equivariant to rotations and reflections. Solver equivariance was empirically verified in Fig. 9. The only conserved quantity for SWE is mass (defined as Az?h times fluid density, so that the mean of Ç is also conserved). Momentum is not conserved in this SWE system, and an
5

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
Table NRMSB-C, S, p(S, (). and average absolute mass and total energy errors for SWEs surrogates 1h and 25h. NaN values indicate
some rollouts diverged to infinity.
NRMSE-C
p(S.5)
Mean( Mass--ef.))
Mean( Total energy-ref.)
Model
1h
25h
1h
25h
1h
25h
1h
25h
FNO/2
0.58++0.03
NaN
0.8391+-0.013
NaN
22.09-+0.19
NaN
86592-4748
NaN
drnet/2
0.11
NaN
0.9977110.0006
NaN
4.84:10.01
NaN
18572:124
NaN
p1/2
NaN
0.9957 =0.0012
NaN
11.7::0.99
NaN
45939:336
NaN
PI/M
0.10::0.11
NaN
0.993410.0016
NaN
0.06 #0.01
NaN
128.0 #23.9
NaN
p4/Ø
0.035=00.003
NaN
0.9992+22-4
NaN
0.17
NaN
768.0=111.77
NaN
p4/M
0.034
3743-1386
0.9992:1224 -0.02
0.04 #0.01
223.4
42.4
153.6=17.2
2.6c9 + 1.9e8
p4m/@
0.045+0.005
NaN
0.9993=2c-4
NaN
2.95 #0.01
NaN
115001148
NaN
p4m/M
0.03210.004
0.02
0.999312--
0.987-144
0.03 10.01
0.29
+0.03
19.1
1094 181
59.5s
76.2s
99.7s
Energy spectrum Velocity specirum
10
reference
SNRMSE-
10%
Momentum
ref.
310%
pdenro pdmmhippi
59.5s
E(8)
100
10
101
100
Wavenumber
p2/1
1.0
-RMMS--U
1010
107
plao=DA PL
59.5s
pénuMs pal
p4m/M
0.4
ref.
pllpi
p1/0+E
p4/pīž
p4m/O
p4m/M+pri
FNO/S Wang
p1j0
p1/Mppp 2
p4/0
p4/M+pü
p4m/pä pi
drnetiO
Figure p4m/M+pü outperforms other networks with similar parameter counts on INS. (a) Reference horizontal velocity with predictions
from pl/Ø and p4m/M+pü. (b-c) Accuracy over 50h rollouts, with standard error of the mean over 30 ICs. (d-e) Log-log plots of the average velocity power spectrum
energy spectrum from 30 ICs at 59.5s. Spectra measure strength the chaotic field's features for each wavenumber & (number of cycles across the domain). Both the velocity and energy spectra p4m/M+ppü align best with the reference. Spectra are scaled by A. (f) The momentum during the entire evolution time for p1/0 and p4m models. (g) Comparison data- aata-nugmentttion (DA) and pushforward trick (PF).
Table 4. The NRMSF-u of decaying turbulence at 36s and 60s for Modern U-net, drnet, Modern U-net with pushforward trick (PF) and Modern U-net with data augmentation (DA).
Modern U-Net
Dilated ResNet
Modern U-net
PF
Modern U-net DA
Model
36s
60s
36s
60s
36s
60s
36s
60s
p1/@
1.56 +0.05
3.32 +0012
1.57 =0.44
2.3e4 #1.9e4
1.88 +0.03
12.03 00.42
1.36 +0.03
1.59 =0.12
pl/pū
1.50+1.004
2.161:.22
1.24+0.02
1.48=0.25
1.38 #0.04
1.47 #0.07
1.36 10.03
1.42 =0.07
pl/M+pü
1 0.005
1.71+0.15
1.33++0004
1.53=0.05
1.43 +0.02
1.52 +0.03
1.31 +0.03
1.38 +0.66
pA/a
1.45:0.04
1.4910.05
1.3410.04
1.42-0.05
1.52 10.03
2.33 40.07
1.3940.04
1.4810.07
p4lpū
1.46:10.03
1.50+0.07
1.33+0.04
1.53=0.06
1.49 10.03
1.54
1.35 +0.03
1.41 #0.06
P4/M+pū
1.38+
1.41+0.05
1.35++0.04
1.47=0.04
1.43 +0.02
1.54 +0.04
1.34
1.31 +0.04
p4m/2
1.37+0004
1.42+0.05
1.36+0.04
1.47=0.07
1.66 +0.05
1.65 #0.05
+0.04
1.39 =0.07
pAm/pūï
1.33:10.03
1.40
1.29
1.40-0.06
1.32 10.03
1.37 40.03
40.03
1.4000.06
p4m/M+pü
1.2910.04
1.20+0.05
1.2940.03
1.30-0.07
1.07 10.01
1.06 +0.02
1.24 10.03
1.25 -0.03
eastward wave will reverse and westwards after reflecting from a boundary. In reality, this would be compensated by slight change the momentum of the Earth, but this is not simulated.
4.2. Decaying Turbulence
The incompressible Navier-Stokes equations (INS) describe
momentum balance for incompressible Newtonian fluids.
6

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
Our 2D version relates velocities [u, v] to pressure y:
accuracy, along with errors in mass and total energy.
du δt + (u V u = Vp + vV P
Fig. 3a compares autoregressive rollouts from unconstrained (p1/2) and maximally constrained networks (p4m/M), p4m/M maintained accurate results for much greater time interval, and in this case was visually indistinguish- able from the reference solution throughout the simulation (for other surrogates, see Fig. 15). Over 20 random held- out ICs, p4m/M exhibited lower normalized RMSE values and higher correlations than other networks (Figs. 3b-c)). p4m/M also outperformed p1/0 trained with input noise (Stachenfeld et al., 2021; Lippe et al., 2024), and uncon- strained FNO and drnet architectures. Compared to other networks, p4m/M trained for more epochs before early stop- ping occurred, reached lower validation loss (Fig. 3d) and produced accurate results for a greater fraction of held-out ICs (Fig. 3e). Mass conservation was respected up to numer- ical precision by the original solver and physies-constrained architectures, but not by other surrogates (Fig. 3f). Over- all, we found that symmetry constraints were more effective than conservation laws, but that equivariant surrogates could be further improved by physical constraints.
(9)
u
(10)
where is fluid density and μ is kinematic viscosity. Here we consider the 'decaying turbulence' scenario introduced by (Kochkov et al., 2021). The velocity field is initialized as filtered Gaussian noise containing high spatial frequencies. Predicting the evolution of the velocity field is challenging, since eddy size and Reynolds number change over time as structures in the flow field coalesce, and the velocity field becomes smoother and more uniform over time.
Numerical reference solution We solve Eqs. 9-10 with C-grid staggering of velocities, using jax-cfd (Kochkov et al., 2021). We follow previous data generation procedures (Kochkov et al., 2021; Stachenfeld et al., 2021), with a 576 x 576 grid and 44 ms time step over 224 seconds. Training data were coarsened a time step of 0.84 s, and resolution was reduced to 48x x 48 (Stachenfeld et al., 2021) using face- averaging to conserve momentum and the divergence-free condition. A burn-in of 148 coarsened steps leaves 120 steps for training. We trained on 100 ICs consisting of filtered Gaussian noise with peak spectral density at wavenumber 10 (that is, 10 cycles across the spatial domain). We used 10 initial conditions for testing and validation.
5.2. Decaying Turbulence
We next trained and evaluated neural surrogates for INS. Here we used the velocity fields [u, u] as both inputs and outputs. As for SWEs, we tested pl, p4 and p4m equivariance, but now considered 3 levels of physical con- straints: unconstrained (Ø), momentum conservation (pū) and mass/momentum conservation (M+) Table (2) lists all constraint combinations used for training, which took 0.4 for nonequivariant networks and 1.4 h for equivariant networks on 2 A100 GPUs. Table 4 lists accuracy after 36s and 60s for all surrogates.
Symmetries and conservation laws We empirically ver- ified the INS solver's equivariance (Fig. 10). Conserved quantities include momentum (equivalent to a constant mean velocities since is constant), and mass (through the divergence- free condition).
Figure (4-a) compares autoregressive rollouts from un- constrained (pl/2) and maximally constrained networks (p4m/M+pü). As for the SWES, both constraint types im- proved accuracy and stability of INS surrogates (Fig. 4b-c), and double constraints were best, also outperforming net- works trained with input noise (Stachenfeld et al., 2021). Unconstrained networks were particularly susceptible to numerical instability in this task (rollouts in Fig. 18-19).
5. Results
5.1. Closed Shallow Water System
We first trained and evaluated neural surrogates for the SWE task. We followed a hybrid learning strategy, based on the observation that the semiimplicit numerical inte- gration scheme calculates çi+1 slowly with an iterative solver, but then calculates [u++1, u given g++1 quickly through a mathematical formula. We therefore trained sur- rogates to predict only Ş++1, and calculated [ ge+1] as in the numerical solver (Appendix I). Keeping param- eter counts constant, we compared networks equivariant to 3 symmetry groups: pl (translation only, as in stan- dard CNNs), p4 (translation rotation-reflection)). We also compared mass conserving net- works (M) to those without physical constraints (Ø). Table 1 lists all constraint combinations used for training, which took 0.5 h for non nouieuuivr networks and 2h for equiv- ariant networks on 2 A100 GPUs. Table 3 shows surrogate
Spectral consistency To evaluate the performance of neural surrogates beyond the time at which their predictions decor- relate from the reference solution, we followed previous studies (Kochkov et al., 2021; Lippe et al., 2024; Stachen- feld et al., 2021) in further comparing the power spectra of predicted velocity fields, and of energy fields 1 / to those of the reference solver. Even after average correlation with the reference solution reached 0, we found that p4m/M+pri networks matched the spectra of the reference solver bet- ter than other methods, consistently across multiple roll- out times and especially at the highest spatial frequencies
and p4m (translation-

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
0.2h
4h
10h
17h
25h
40h
0.8s
12.6s
38.5s
59.5s
82.9s
99.7s
reference
Volocity spectrum p4m/M+pü p1/0 reference
p1/0
p4m/M
NRMSE
b
1.0
p(5.C)
p1/0
p1/M
10
10
Energy spectrum
10th
1010
p1/M-ppi
ref.
P1/0+E
(kk)
pl/ts=
P4/M
p4/0
10
99.7s
10
99.7s
paIMepi
P413
pátpi
p4m/0
3 ((h)
50
0.20
(h)
50
p4m/M
10
10
10
10'
Wavenumber
Wavenumber k:
Figure 5. Generalization beyond training data. (a) SWE rollouts from p1/03 p4m/M L-shaped ICs. (b-c) Accuracy of cach network
six generalization tests (Appendix 1.. (d) INS rollouts from p1/2 and p4m/M+pu on ICs with peak wavenumher (c-f) Velocity- and
energy spectra INS 99.7 s, averaged over ICs.
Table Comparison nRMSE for zonal velocity, energy spectrum enror (ESE), and p(ù, u) real ocean currents predicted 1 120h ahead. Architecture and hyperparameters for Equ Unet are described in (Wang et al., 2021)
12h
120h
Model
nRMSE
ESE
p(ù, u)
nRMSE
ESE
p(ū, u)
p1/03
1.52++0.04
5.33++6.23
0.01++0.2
1.66++0.03
13.9+11.4
-0.02++0.02
Equ Unet (Wang et al., 2021)
0.85+0.16 0.91
0.31+0.02 0.80
1.1420.03
1.07#1.11
0.00+00.02
p4m/M+pü
0.88
0.45:10.03
1.10=00.0
0.94
0.11
rogates (Fig. 4g): data augmentation and the pushforward trick. We applied data augmentation using the p4m symme- try group, such that velocity fields the staggered C-grid were transformed consistently with the numerical solver: S Tw()) T, S. We applied the pushforward trick as in (Brandstetter et al., 2022c), with the MSE loss computed after two autoregessive time steps, but gradients backpropa- gated only one step. Both of these training modes improved performance compared to standard training of p1/2, but could not match the accuracy of p4m/M++rü with standard training (Fig. 4g, Table 4). Data augmentation had, as ex- pected, no effect on p4m surrogates, but pushforward train- ing p4m/M+pr produced the most accurate surrogate overall, showing that doubly constrained networks benefit from autoregressive training. We present further results and details these modes in appendices J.9-J.10.
NAMSE-W
At ()
plù 1)
0.1M
2M
8.5M
100
400
760
Network size
Training data size
=4.22
12.66
p1/0
pám/M+ppi
p1/0 p4m/M
Figure 6. Accuracy of symmetry- and physics-constraine INS models across data and network sizes, at 4.2 and 12.6 s. (a- b) NRMSF-u and p(û. u) vs. network size for pl/6 and p4m/M+pri. (c-d) NRMSE-u and p(ū, u) for p1/2 and p4m/M+pu vs. training datasets size.
(Figs. 4d-e, additional spectra in Fig. 21). Conservation of momentum was respected only by momentum-conssrainned networks (Fig. 4f, 20). Training p4m/M+pū surrogates with input noise resulted in lower accuracy but excellent long-term numerical stability (Fig. 22).
A Alternate base architecture Beyond evaluating the utility of constraining modern U-nets, we also evaluated symme- tries and physical constraints for Dilated ResNet surrogates of INS(Table 4). Performance was strikingly consistent with previous results, with symmetries more effective than phys- ical constraints but further benefits observable when com- bining both. Thus, the separate and combined effects of our constraints were consistent across these two architectures.
Specialized training modes We further compared the per- formance gains from symmetries and physical constraints to those offered by specialized training modes for PDE sur-
8

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
6h
36h
60h
120h
180h
300h
5.5. Predicting Ocean Current Observations
reference
We trained neural surrogates to predict real-world observa- tions of ocean currents (Wang et al., 2021) (data and training details in Appendix J.14). Doubly-constrained p4m/M+pü predicted future observations better than p1/2, and also out- performed the equivariant network proposed in Wang et al., 2021, in terms of NRMSE and correlation (Fig. 7) as well as Energy Spectrum Error (ESE, Table 5).
p4m/M-pph
NAMSE Go
6. Discussion
ref.
pfit. ai..a)
pt/O0
We enforced hard constraints on symmetries and conserva- tion laws for neural PDE surrogates. We extended the appli- cability of previous techniques to staggered grids, and sys- tematically tested performance across tasks and constraints. Symmetries were more effective, but conservation laws were not redundant. Double constraints matched reference sim- ulations, individually and statistically, better than multiple architectures, pushforward training, and data augmentation. In our challenging PDE tasks, insufficiently constrained sur- rogates diverged towards infinity before completing their rollouts (Table 3), as did the PDEs' conserved quantities themselves in some cases (Fig. 3f, 4f). These results un- derline physics- and symmetry-constraie promising strategy when long-term accuracy is required but data expensive or limited, as in weather, climate and industrial fluid mechanics.
Equñot
p4m/M+ piž
0
(b)
300
(M)
300
Figure 7. Prediction of ocean current observations by p1/2,
Wang's Equir (Unet) and p4m/M+ppü. (a) Zonal ocean currents from reference observations and predictions. (b-c) NRMSE-u and correlation for forecasts up 300 h ahead (averaged over 30 ICs).
We present additional results for drnets in Appendix J.8.
5.3. Generalization to Novel ICs
Closed Shallow Water System Fig. 5a shows an initial condition consisting of an 'L'-shaped elevation (details in Fig. 31). Randomly varying the elevation's 's location and shape, we found p4m/M to outperform alternatives with equal parameter counts (Fig. 5b-c). Additional results are shown in Fig. 29-30.
surrogates as a
Limitations & Future work For large enough net- works and datasets, constraints might be learned from the data (Stachenfeld et al., 2021; Watt-Meyer et al., 2023), but the benefit of constraints grows with rollout length even for large networks and datasets. Thus, constraints are likely relevant for longer time scales, e.g. for seasonal forecasting and climate projections (Kochkov et al., 2024; Watt-Meyer et al., 2023; Nguyen et al., 2023). How constraints limit error accumulation remains unclear, but empirical investiga- tions of how error accumulation correlates with constraint violations over time and ICs could provide some clarity.
Decaying Turbulence We tested surrogates on ICs with peak wavenumber changed from 10 to p4m/M+pu matched the reference solver (Fig. 5d) and its spectra (Figs. 5e-f) most closely. Additional results are shown in Figs. 32-34.
5.4. Effects of Network and Dataset Size
To investigate how the benefits of symmetries and physical constraints scale with network and dataset size, trained p1/Ø and p4m/M+pr networks with 0.1M, 2M and and 8.5M parameters on 100 INS simulations. At both 4.2 and 12.6 we observed lower errors and high correlations for p4m/M+pū for all network sizes. The relative advantage of p4m/M+pü over pl/Ø was greatest for smaller networks and longer forecast horizons, and overall performance was best for larger networks. CPU and GPU inference speeds reported in Appendix J.13.
We considered mass and momentum conservation, and sym- metries of square 2D grids. Future work could pursue other PDEs such as hyperbolic equations (Takamoto et al., 2022), energy conservation (Cranmer et al., 2020), local vs. global conservation (McGreivy & Hakim, 2023), continuous sym- metry groups (Cohen et al., 2018; Esteves et al., 2018), alternative grids and meshes (Cohen et al., 2019; De Haan et al., 2021; Horie & Mitsume, 2022), generalization to new geometries (Wandel et al., 2021; Horie & Mitsume, 2022), unrolled training (Brandstetter et al., 2022c; List et al., 2024), invariant measure learning (Schiff et al., 2024) and generative modeling (Lippe et al., 2024; Kohl et al., 2024).
Training 0. 0.1M. Maaiingg..M---aMaMMa M-parameter p1/@ and p4m/M+prū networks on 100, 400 and 760 simulations showed constraints en- hance performance robustly across dataset size (Fig. 6c-d). Improvements were greater on larger datasets and longer rollouts (additional results and spectra in Fig. 35).
9

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
Acknowledgements
Brandstetter, J., Worrall, D. E., and Welling, M. Message passing neural pde solvers. In International Conference on Learning Representations, 2022c.
We thank A.C. Bekar and Kumar for comments on the manuscript.
Bukka, S. R., Gupta, R., Magee, A. R., and Jaiman, R. K. Assessment of unsteady flow predictions using hybrid deep learning based reduced-order models. Physies of Fluids, 33(1), 2021.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must specifically highlighted here.
Cesa, G., Lang, I.., and Weiler, M. A program to build e (n)- equivariant steerable cnns. International conference on learning representations, 2022.
References
Chalapathi, N., Du, Y., and Krishnapriyan, A. S. Scal- ing tmysics-informer hard constraints with mixture-of- experts. In International Conference on Learning Repre- sentations, 2024.
Arakawa, A. Computational design of the basic dynamical processes of the ucla general circulation model. Methods in Computational Physicx/Academic Press, 1977.
Chidester, B.,
Zhou, T., Do, M. N., and Ma, J. Rotation equivariant and invariant neural networks for microscopy image analysis. Bioinformaties, 35(14):150-i5377 2019.
Backhaus, J. O. A semi-implicit scheme for the shallow water equations for application to shelf sea modelling. Continental Shelf Research, 2(4):243- 254, 1983.
Cohen, T. and Welling, M. Group equivariant convolu- tional networks. In International Conference on Machine Learning, pp. 2990- 2990-2999. PMLR, 2016.
Bar-Sinai, Y., Hoyer, S., Hickey, J., and Brenner, M. P. Learning data-driven discretizations for partial differen- tial equations. Proceedings of the National Academy of Sciences, 116(3))115344-15522 2019.
Cohen, T., Weiler, M., Kicanaoglu, B., and Welling, M.
Gauge equivariant convolutional networks and the icosa- hedral enn. In International conference on Machine learn- ing, pp. 1321-1330. PMLR, 2019.
Bergamin, F., Diaconu, C., Shysheya, A., Perdikaris, P., Hernández-Lobato. J. M., Turner, R. E., and Mathieu, E. Guided autoregressive diffusion models with appli- cations to pde simulation. In ICLR 2024 Workshop A14Diferentialliuaaiin In Science, 2024.
Cohen, T. S., Geiger, M., Köhler, J., and Welling, M. Spher- ical cnns. In International Conference on Learning Rep- resentations, 2018.
Beucler, T., Pritchard, M., Rasp. S., Ott, J., Baldi, P., and Gentine, P. Enforcing analytic constraints neural net- works emulating physical systems. Physical Review Let- ters, 126(9))98332,, 2021.
Cranmer, M., Greydanus, S., Hoyer, S., Battaglia, P., Spergel, D., and Ho, S. Lagrangian neural networks, 2020. URL https: arwiv.avv..oo org/ abs 2000.. 04630.
De Haan, P., Weiler, M., Cohen, T., and Welling, M. Gauge equivariant mesh cnns: Anisotropic convolutions geo- metric graphs. In International Conference on Learming Representations, 2021.
Bonev, B., Kurth, T., Hundt, C., Pathak, J., Baust, M., Kashinath, K., and Anandkumar, A. Spherical fourier neu- ral operators: Learning stable dynamics on the sphere. In International conference on machine learning, pp. 2806- 2823. PMLR, 2023.
De Pondeca, M. S., Manikin, G. S., DiMego, G., Benjamin, S. G., Parrish, D. F., Purser, R. J., Wu, W.-S., Horel, J. D., Myrick, D. T., Lin, Y., et al. The real-time mesoscale analysis noaa's national centers for environmental pre- diction: current status and development. Weather and Forecasting, 26(5):593-612, 2011.
Brandstetter, J., Hesselink, R., van der Pol, E., Bekkers, E. J., and Welling, M. Geometric and physical quan- tities improve e(3) equivariant message passing. In International Conference on Learning Representations, 2022a. URL htt.tttpt////ptteteeeeeeet ww.pe.review.net..e.. id= @d=_xwr8gOBeV1.
Esteves, C.,
C., Makadia, A., and Dani- ilidis, K. Learning (3) equivariant representations with spherical enns. In Proceedings of the European Confer- ence Computer Vision (ECCV), pp. 52-68, 2018.
Brandstetter, J., Welling, M., and Worrall, D. E. Lie point symmetry data augmentation for neural pde solvers. In International Conference on Machine Learning, pp. 2241- 2256. PMLR, 2022b.
Fanaskov, V., Yu, T., Rudikov, A., and Oseledets, General covariance data augmentation for neural pde solvers. In
10

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
International Conference Machine Learning, pp. 9665- 9688. PMLR, 2023.
Knigge, D. M., Wessels, D., Valperga, R., Papa, S., Sonke, J.-J., Bekkers, E. J., and Gavves, S. Space-time continu- ous pde forecasting using equivariant neural fields. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Ferziger, H., Perić, M., and Street, R. L.. Computational methods for fluid dynamics. springer, 2019.
Gasteiger, J., GroB, J., and Günnemann, S. Directional message passing for molecular graphs. In International Conference on Learning Representations, 2020.
Kochkov, D., Smith, J. A., Alieva, A., Wang, Q.. Brenner, M. P., and Hoyer, S. Machine aarming-accelerated com- putational fluid dynamics. Proceedings of the National Academy of Sciences, (18011)e21001888111 2021.
Greenfeld, D., Galun, M., Basri, R., Yavneh, I., and Kimmel, R. Learning to optimize multigrid pde solvers. In Interna tional Conference on Machine Learning, pp. 2415-2423. PMLR, 2019.
Kochkov, D., Yuval, J., Langmore, I., Norgaard, P., Smith, J., Mooers, G., Klöwer, M.. Lottes, J., Rasp, S., Düben, P., al. Neural general circulation models for weather and climate. Nature, pp- 1-7, 2024.
Greydanus, S., Dzamba, M., and Yosinski, J. Hamiltonian neural networks. Advances in neural information process- ing systems, 32, 2019.
Kohl, G., Chen, L.,, and Thuerey, N. Benchmarking au- toregressive conditional diffusion models for turbulent flow simulation. In ICML 2024 Al for Science Workshop, 2024.
Gupta, J. K. and Brandstetter, J. Towards multi- spatiotemporal scale generalized pde modeling. Transac- tions on Machine Learning Research, 2023.
Korn, P., Brüggemann, N., Jungelaus, J. H., Lorenz, S., Gutjahr, O., Haak, H., Linardakis, I.., Mehlmann, C., Mikolajewicz, U., Notz, D., et al. Icon-o: The ocean component of the icon earth system model-global sim- ulation characteristies and local telescoping capability, Journal of Advances in Modeling Earth Systems, 14(10): 22021MS002952, 2022.
Helwig, J., Zhang, X., Fu, C., Kurtin, J.. Wojtowytsch, S., and Ji, S. Group equivariant fourier neural operators for partial differential equations. In International Conference on Machine Learning, pp. 12907-12930, 2023.
Holl, P. and Thuerey, N. now (PhiFlow): Differentiable sim-
ulations for pytorch, tensorflow and jax. In International Conference on Machine Learning. PMLR, 2024.
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, 2021.
Horie, M. and Mitsume, N. Physics-embedded neural net- works: Graph neural pde solvers with mixed boundary conditions. Advances in Neural Information Processing Systems, 35:23218 2-22299, 2022.
Hsieh, J.-T., Zhao, S., Eismann, S., Mirabella, L., and Er- mon, S. Learning neural pde solvers with convergence guarantees. International Conference on Learning Representations, 2019.
Li, Z., Peng. W., Yuan, Z., and Wang. J. Long-term pre- dictions turbulence by implicit u-net enhanced fourier neural operator. Physics of Fluids, 35(7), 2023.
Lino, M., Fotiadis, S., Bharath, A. A., and Cantwell, C. D. Multi- rotation- -tation-equivariant graph neural networks for unsteady eulerian fluid dynamics. Physics of Fluids, 34(8), 2022.
Huang, Y. and Greenberg, D. S. Symmetry constraints enhance long-term stability and accuracy in unsupervised learning of geophysical fluid flows. Authorea Preprints, 2023.
Lippe, P., Veeling, B., Perdikaris, P., Turner, R., and Brand- stetter, J. Pde-refiner: Achieving accurate long rollouts with neural pde solvers. Advances in Neural Information Processing Systems, 36, 2024.
Jasak, H. Openfoam: Open source cfd in research and industry. International journal of naval urchitecture and ocean engineering, 1(2):89 -94999994229 2009.
Jungclaus, J. H., Lorenz, S. Schmidt, H., Brovkin, V., Brüggemannn N., Chegini, F., Crüger, T., De- Vrese, P., Gayler, V., Giorgetta, M. A., et al. The icon earth system model version 1.0. Journal of Advances in Modeling Earth Systems, MSM0)22222MM02811111 2022.
List, B., Chen, L.-.., Bali, K., and Thuerey, N. How temporal unrolling supports neural physics simulators. CoRR, 2024.
Long, Z., Lu, Y., and Dong, B. Pde-net 2.0: Learning pdes from data with a numeric-symbolic hybrid deep network. Journal of Computational Physics, 399:108922 2019.
Kingma, D. P. Adam: A method for stochastic optimization. arXiv preprint arXiv: /4// 29800, 2014.

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
Madec, G., Bell, M., Blaker, A., Bricaud, C., Bruciaferri, D., Castrillo, M., Calvert, D., Chanut, J., Clementi, E., Coward, A., Epicoco, I., Éthé, C., Ganderton, J., Harle, J.. Hutchinson. K., Iovino, D., Lea, D., Lovato, T., Mar- tin, M., Martin, N., Mele, F., Martins, D., Masson, S., Mathiot, P., Mele, F., Mocavero, S., Müller, S., Nurser, A. G., Paronuzzi, S., Peltier, M., Person, R., Rousset, C., Rynders, S., Samson, G., Téchené, S., Vancoppenolle, M., and Wilson, C. Nemo ocean engine reference manual, July 2023a.
Read, S., Jia, X., Willard, J., Appling, A. P., Zwart, J. A., Oliver, S. K., Karpatne, A., Hansen, G. J., Hanson, P. C., Watkins, W., et al. Process-guided deep learning pre- dictions of lake water temperature. Water Resources Research, 55(11)99173-9)
2019.
Rebelo, R. and Valiquette, F. Symmetry preserving numer- ical schemes for partial differential equations and their numerical tests. Journal of Difference Equations and Applications, 19(5):738-757, 2013.
Madec, G., Bell, M., Blaker, A., Bricaud, C.,, Bruciaferri, D., Castrillo, M., Calvert, D., Chanut, J., Clementi, E., Coward, A., et al. Nemo ocean engine reference manual. Mrešno]. Available: https://zenodo. org/record/8167700, 23, 2023b.
Romero, D. W. and Cordonnier, J.-B. Group equivariant stand-alone self-attention for vision. In International Conference on Learning Representations, 2021.
Ronneberger, O., Fischer, P., and Brox, T. U-net: Con- volutional networks for biomedical image segmenta- tion. In Medical image computing and computer-assisted intervention
Marullo, S., Santoleri, R., Ciani, D., Le Borgne, P., Péré, S., Pinardi, N., Tonani, M., and Nardone, G. Combin- ing model and geostationary satellite data to reconstruct hourly sst field over the mediterranean sea. Remote sens- ing of environment, 146:1 1-23, 2014.
2015: 18th international confer- ence, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pp. 234 241. Springer, 2015.
Ruhe, D., Brandstetter, J., and Forré, P. Clifford group equiv- ariant neural networks. Advances in Neural Information Processing Systems, 36, 2024.
McGreivy, N. and Hakim, A. Invariant preservation in machine learned pde solvers via error correction. In ICLR Workshop Physies for Machine Learning, 2023.
Sanchez- Zonzalzz, A., Godwin, J., Pfaff, T., Ying, R., Leskovee, J., and Battaglia, P. Learning to simulate com- plex physics with graph networks. International Con- ference on Machine Learning, pp. 8459-8468. PMLR, 2020.
Michelis, M. Y. and Katzschmann, R. K. Physics- constrained unsupervised learning of partial differential equations using meshes, 2022. URL https:: /arxiv.. org/abs/ 2203. 6628.
Mohan, A. T., Lubbers, N., Livesou, D., and Chertkov, M. Embedding hard physical constraints in neural net- work coarse- graining of 3d turbulence. arXiv preprint arXiv:2002.0002/ 2020.
Schiff, Y., Wan, Z. Y., Parker, J. B., Hoyer, S., Kuleshov, V., Sha, F. and Zepeda-Núñez, L. DySLIM: Dynamics Sta- ble Learning by Invariant Measure for Chaotic Systems, June 2024. URL ht:p: /ttpp:://axiv..../// orw.v////rrabb///222 04467. 67 arXiv:2402.04467.
abs/
Nguyen, T., Brandstetter, J., Kapoor, A., Gupta, J. K., and Grover, A. Climax: A foundation model for weather and climate. In International Conference on Machine Learning, pp. 25904 -59938. PMLR, 2023.
Smets, B. M., Portegies, J., Bekkers, E. J., and Duits, R. Pde- based group equivariant convolutional neural networks. Journal of Mathematical Imaging and Vision, 65(1))209- 239, 2023.
Noether, E. Invariante variationsprobllmee Nachrichten von der Gesellschaft der Wissenschuften zu Göttingen, Seachhematsch--Plyiiilllin
Klasse, 1918:235-257, 1918. URL http: http///udmllllle.e
Song, T., Main, A., Scovazzi, G., and Ricchiuto, M. The shifted boundary method for hyperbolic systems: Embed- ded domain computations of linear waves and shallow water flows. Journal of Computational Physics, 369: 45-79, 2018.
59024.
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.
Sorourifar, F., Peng, Y., Castillo, I., Bui, L., Venegas, J. and Paulson, A. Physics-enhanced neural ordinary differential equations: Application to industrial chemical reaction systems. Industrial & Engineering Chemistry Research, (1333::15563-55777 2023.
Raonic, B., Molinaro, R., De Ryck, T., Rohner, T., Bar- tolucci, F., Alaifari, R., Mishra, S., and de Bézenac, E. Convolutional neural operators for robust and accurate learning of pdes. Advances in Neural Information Pro- cessing Systems, 36, 2024.
Stachenfeld, K., Fielding, D. B., Kochkov, D., Cranmer, M., Pfaff, T., Godwin, J., Cui, C., Ho, S., Battaglia,
12

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
P., and Sanchez-Gonzalez, A. Learned coarse mod- els for efficient turbulence simulation. arXiv preprint 1rXiv:2112.15275, 2021.
Wandel, N., Weinmann, M., and Klein, R. Learning in- compressible fluid dynamics from scratch-towards fast, differentiable fluid models that generalize. In Interna- tional Conference on Learning Representations, 2021.
Stone, J. M., Tomida, K., White, J., and Felker, K. G. The athena++ adaptive mesh refinement framework: Design and magnettohydrodyyamiic solvers. The Astrophysical Journal Supplement Series, 249(1):4, June 2020, doi: 10. 3847/1533-4365//abbab URL https: fooi.org 10. 384 3w4733242221 43654433662229333
Wang. R., Walters, R., and Yu, R. Incorporating symmetry into deep dynamics models for improved generalization. In International Conference Learning Representations, 2021.
Watt-Meyer, O., Dresdner, G., McGibbon, J., Clark, S. K., Henn, B., Duncan, J., Brenowitz, N. D., Kashinath, K., Pritchard, M. S., Bonev, B., et al. Ace: A fast, skillful learned global atmospheric model for climate prediction. CoRR, 2023.
Sun, Z., Yang. Y., and Yoo, S. A neural pde solver with temporal stencil modeling. In International Conference on Machine Learning, pp. 33135-33155, PMLR, 2023.
Takamoto, M., Praditia, T., Leiteritz, R., MacKinlay, D., Alesiani, F., Pflüger. D., and Niepert, M. Pdebench: An extensive benchmark for scientific machine learning. Advances in Neural Information Processing Systems, 35: 1596-1611, 2022.
Weiler, M. and Cesa, G. General e (2)-equivariant steerable cmns. Advances in neural information processing systems, 32, 2019.
Thomas, N., Smidt, T., Kearnes, S., Yang. L., Li, L., Kohlhoff, K., and Riley, P. Tensor field networks: Rotation-and vaaslation-equivarti neural networks for 3d point clouds. arXiv preprint orXiv:1802.088 2018.
Tompson, J., Schlachter, K., Sprechmann, P., and Perlin, K. Accelerating eulerian fluid simulation with convolu- tional networks. In International conference on machine learning, pp. 3424-3433. PMLR, 2017.
Toshev, A. P., Galletti, G., Brandstetter, J., Adami, S., and Adams, N. A. Learning lagrangian fluid mechanics with e3)-evivvrr graph neural networks. In International Conference on Geometric Science of Information, pp. 332-341. Springer, 2023.
Tran, A., Mathews, A., Xie, L., and Ong, C. S. Factorized fourier neural operators. In International Conference on Learning Representations, 2023.
Tripura, T. and Chakraborty, S. Wavelet neural operator for solving parametric partial differential equations in com- putational mechanics problems. Computer Methods in Applied Mechanics and Engineering, 404: 404:115758 115783, 2023.
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.
Veeling, B. S., Linmans, J., Winkens, J., Cohen, T., and Welling, M. Rotation equivariant cnns for digital pathol- ogy. In Medical Image Computing and Computer As- sisted MICCAA 2018: 21st International Conference, Granada, Spain, September 16-20, 2018, Proceedings, Part I, pp. 210-218. Springer, 2018
13

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
A. Symmetries of SWEs and INS with C-grid staggering
A.1. Grid Discretizations
The Arakawa C-grid (Arakawa, 1977) is a discretization technique frequently employed in numerical simulation, partic- ularly in fluid dynamics. In this section, the C-grid staggering for shallow water equations (SWEs) and incompressible Navier-Stokes equations (INS) presented (Figure 8).
C-grid staggering of SWEs
b
C-grid staggering of INSs
V1
U12
V11
U12
a13
y
S11
S12
S1
a11
Q1
P11
P12
u11
U12
U13
U11
U12
U13
221
U22
U21
U22
a23
S21
C2
223
u121
A.22
P21
P22
U2
U22
U23
U22
U222
U23
U31
U32
U31
U32
S
S32
C3
a31
0:32
a.33
Figure 8. Arakawa's C-grid staggering approach has been employed for time integrating the of shallow water equations (SWEs) and incompressible Navier-Stokes equations (INS). (a) Red square points denote the vertical displacement of the free surface, denoted by Ç. The circular points correspond to the velocities and v. (b) In the discretization of the INS system, the square points represents the vector potential while the velocity is indicated by the circular points. The black stars indicate pressure, which is calculated by the numerical solver each time step but used prognostic variable solve for velocity in future time steps.
A.2. The symmetries of shallow water equations
The shallow water system with closed boundary is represented by Equations (6-8). For specified set of PDE fields (i.e., and v) at given time I, the solution of the equation for Ç at the subsequent time interval can be expressed as Se v). The symmetry transformations (flip, rotation, and flip-rotation) of SWEs with the numerical solver S can be expressed as follows:
flip: F()) u(S(().u,))
(11)
rotation: )((R(() R(u)) = R(SS((S flip rotation: (C(R(F(()) R(F())) = =((F(S((G,,,)))
(12)
(13)
In this section, the numerical solution of the equations of motion for the variable in the subsequent time step will be denoted by Sç. The actions of the flipping operator, F, and the rotation operator, R, will also be defined for these PDE fields. Figure show the equivariance of the numerical solver to these transformations.
14

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
<
U
V
Sc(S, u)
F( M, v))
Flip
numerical solver
flip
222
F(C)
F(u)
F())
F())
T1
numerical solver
Identical
Rotation
U
v
Sc(S, u, 2))
R(Sc(C,
numerical solver
rotation
R(()
R(()
R(u)
SC(R(C), R(v). R(w))
numerical
solver
identical
Rotation-filp
S
U
V
Sc(S, u,
((,
numerical solver
flip rotation
R(F())
R(F()))
R(F(u)
(((R(F(C) R(F(U)) R(F(U))
numerical
solver
Identical
Figure 9. Empirical validation of the symmetries of the numerical SWE solver is demonstrated through three transformations: flip, rotation, fotp-rotation. These plots correspond Equations (11-13).
A.3. Symmetries of the incompressible Navier-Stokes equations
The incompressible Navier-Stokes equations are presented Equations (10-9). Given a set initial conditions u and a particular time f, the solutions to the velocity equation are expressed as S, (u, v)and S, v) time 1. The mathematical description of the flip, rotation, and flip-rotation symmetries of INS is follows:
S.(-F(u),F()) = - o(),(u,))
flip
(14)
F(()) F(So(u,v)
-.(R()).-((u) R(, v)) -R(() R(,
rotation
(15)
S.(R(((6)),-R(( -R(-) R(F(S. v))) -R(-F)))) = -R( v)))
flip rotation
(16)
As shown in Figure 10, the symmetries of the incompressible Navier-Stolr Stokes equations encompass are generated by flipping and rotation.
15

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
Filp
u
Su(u, v)
(u,u, v)
F(S.(u,o)
F(S,
numerical solver
fllp
F(u)
F(2)
F(u), F(o))
S.( F(e))
numerical solver
Identical
Rotation
Su,u, v)
Sy (u, v)
R(S.(u.))
R( (v,
numerical solver
rotation
R(v)
R(x)
S.(R(v), -R(u))
(R(R(),, - R(u))
numerical solver
Identical
Rotation-fllp
Sy(u, v))
S, v)
R(F(S.(M 1)))
RC ((S.(u. o))
2
numerical solver
filp rotation
R(F(c))
R8( F(a)))
(R(FF(()) R F(N))) (R((F())) RR(-RRF F())))
numerical solver
Identical
Figure 10. Empirical confirmation of the equivariance of the INS solver.
B. Padding options
In certain numerical solvers, although C-grid staggering is employed, the software generates output of an identical size for each component of the vector field. This phenomenon necessitates thorough examination of the selected conventions for padding and boundary representation in the outputs. In such instances, the implementation of padding technique becomes imperative to restore the vector field to its standard C-grid staggering. One potential solution is the implementation of periodic padding for the periodic boundary conditions (BCs). It is imperative to emphasize that the physical properties of the data, such as the divergence in the incompressible Ravier-Stokes equations and the BCs, must remain unaltered during
16

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
the implementation of this correction.
C. Group equivariant input layers for SWE and INS
In this section, we present explicit formulas for equivariant input layers and show that they satisfy equivariance relations. For the sake of brevity, we have elected to include only the proof for the p4 input layer for SWEs, but extending this proof to p4m or INS straightforward.
C.1. Group equivariant input layer for SWEs
Given that the input data for shallow water equations (SWEs) utilizes C-grid staggering, as illustrated in Figure 8, is necessary to construct an input layer that aligns with the C-grid staggering while preserving equivariance. On the C-grid, the velocity components and possess distinct dimensions. To address this, we propose the utilization of two rectangular filters, denoted -. and W , for the variables and i, respectively. The symbol is employed to denote all values along given axis. The filter W is Gi Co K x S array, where cy denotes the batch size, en signifies the number of input channels, and K x S represents the filter size.
For instance, on a periodic 2D grid tiling the torus with 3 x 3 grid cells, the sizes are 4 x 3 for u and 3 x 4 for v. While the height field Ç is defined at grid cell centers, u, v are defined on the interfaces (edges) between the grid cells. In the context of performing group transformations in the input layer, imperative to exchange the filters for u and v to ensure compatibility with the dimensions of the input variables.
p4 We first introduce an input layer of the p4 group transformation for SWEs, which has four channels obtained from four different rotated filters. The output y of this input layer is defined as follows (* denotes convolution):
y
(WS. Sihe * Gir.))
(W Wiir. * ui...)
Jair vi,..) + bj.
(17)
cn
R90 not
(WS, 2.b1.,. * Sis.))
+
-B900" (W5)..) Uiis))
==
(18)
10000 (itt...) vey.))
+ +b3:
2a -1
V/....
1800 (WS 2oh..)) Gi...)
+
Σ
(
(-11100 (W * u....)
1=0
(19)
p180 (Wy...) Vi... + by,
1=0
-1
6 -1
V/,3,r
R2770
(WS
(. Sir..) +
Σ
T2270'
(W)...)) :i...)
1=0
2=0
(20)
G
270 (W). vi.... + by.
1=
where WSS 10yy is square filter for Ç, for example 4 x 4. b is a clement vector.
17

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
p4m Next, we define an input layer for p4m that uses the same logic as the p4 input layer. It has different group transformations, including flips and rotations.
in
Vi.0..
W Moi.
St...)
(W J.b. * ui...) +
(Wy) J.d. * vi...) + bj.
(21)
==0
V)1...
(in jib...)) * Se...)
(Fii( (air.)) * us...))
(22)
(Fmp (W), s.d...) * (i,..) + bj.
==0
cin
R200 (W j,h.,,) * Siss) + E
Rood (Wit...) * uuys))
(23)
Ro00 jit.. vh..)) +bz:
Via..
Fmip ( R900
(WS, j.,,, * Se...
+
Fnip R990 (W )) Us..)
1=0
(24)
Faip R (W,),,))) Beyy))
+ by.
1=0
c -1
Uj.4.,
R2 (WS J.b...) S1...) +
R2100 (W 3ir..) uu...)
==
(25)
R180 (Wy...)) * vbb..) + by.
1=0
Fnipl ( 180 (WS (W.,,))) Sir.) +
(Ra (Wa,..)) wice)
1=0
1=0
(26)
Prip (12
(W (W)...)) * Virr.)
+ رز
==
-1
V
R2270" (WS,..) Girs)
R2270 (WJ...) mi...
1=0
(27)
-1
122700 (WJ), Jit...) * (e...)) ) + oj.
1=0
Fmppl (R2709 (W5 (WS * Siree) +
Frap R2700 (W (W)...)) + ui..)
(28)
Faip R2700 (W).,.)) ver.) bj.
where the filters Wt. J. and Wy.... are rectangles and the filter W s.h. a square.
18

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
C.2. Proof of p4 equivariance for the SWE p4 input layer
Here we prove group equivariance (Cohen & Welling, 2016) for the SWE p4 input layer. To this end, the following equation must be demonstrated: y( (R()). -R(v). R(u)) R(y(S. u, u)) where y is the output of the p4 input layer for SWEs. This equation is simply the rotation symmetry of shallow water system depicted in Equation 12.
The forms of the ordinary input layer, denoted by u(S, u, v) in Equations (17-20), have previously been described. Applying them to a rotated inputs consisting of Roo (Si,,), - 12000 (u,,,)). ). and Ror (us,.)) we get:
W Janr Rooi (Ge...) + Σ
W Juiss - -RR00" (11,..)
(29)
(we. Jih. R (ui...) + bj.
in -1
[R00r (W * Ro0 (Gt...) +
Σ
(-
Reod (Wi.,.)) -R (vt,..)
(30)
Read War.) Ruod (us...) + bj.
(Rib0"" (W n2o0 (Ge..)) +
(- RR (W)....) -R000 (Ui..))
1=0
(31)
(R1,00 12180 (WJ...)) 12000
(ui..)) bb
1) =
12270' (WS Joh..)) * 12900 (Ge...) +
R2270
(Wr...) - R2000 (vi,.))
1=0
1=0
(32)
G
R2700 (W....)) * R7000 (us..)) + bj-
1=0
19

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
Applying the same a 90-degree rotation the regular representation output for unrotated inputs, get:
R
(w)....)) R
(Wi.. * (Gis)) +
Whi.e John U...))
(33)
Wiiss * Ui.,.) +bg
0=
cin
Rio (,),..,) = R
R000 (W (.... * S1...)
+
Σ
(- (W)r..)) )
k=
1=
(34)
Roo (W),,
vs:..) b3
c
cin
1 (w)2...) = RRo0
$000
Σ
12880" (W J,,,,) Si...) + Σ (
(-Reat (W * win.)
i=0
(35)
12180* (.....) * v + bb
= W......
i=0
c -1
c -1
R2o (w),...) = -Ro00"
R2200 (WS, (.... Ser..) +
Σ
(aane wi...)
i=n
1=0
(36)
c -1
(- Ren (WJ)..) * viyy) + bj
1 0y.,.
=
Applying the definition of the rotation operator's is action on the regular representation (rotation + permutation), we see the channels of these two regular representations match, so we arrive at the formula:
Re00 (y' (Si,.v Wiper v)...)) ij (Figo (Gi...), -R (v.,,.), Redd (ua,..)
(37)
Thus have proved the group equivariance of the p4 input layer in shallow water equations.
C.3. Group equivariant input layer for INS
p4 The system state for the incompressible Navier-Stokes equations consists of velocity fields u and v, which possess disparate dimensions in the C-grid staggering. Consequently, two rectangular filters, denoted as W J and W required for the velocity field. According to the symmetries of rotation of INS in Equation (15), we first construct p4 input layer for INS as follows:
V/.
WWas. * ui..))
+
Wia * vi..) + bj:
(38)
-1
[2900 (Wij..)) Us...
1-2000 (Wit...) * wh...) + bj,
(39)
1=0
ca -1
6 -1
v). 11.2.r
(-
180 (Was...) * Ui..)
+
1180 (W (W * vs.:)) + bj:
(40)
1=0
1=0
c -1
c -1
3/
R270
(W
* Uir.))
+
[2270
(WT....) vi...) + by-
(41)
p4m Subsequently, in accordance with the flip-rotation symmetries of INS in Equation (16), the p4m input layer is defined
20

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
by the following eight equations:
-1
yj.,,
(W Jay. u.....) )
+
(W Jáy. * Uk...) + by,
(42)
==
(Fri (W j.b..;) * up
+
(Fri (h * vs..) + by.
(43)
Oin
Ro00 (W),..)) * Ui...)
Agoo (Wit.)) * vh..) + by,
(44)
Fitip Rugo (Wjt,))) Uz...))
+
Fitip Ruae" (Wir.))) vs...) + bj.
(45)
i=0
c
-1
Uj4,1
RRid0" (W Wisse) (iss))
+
1s00 (W * oi.:) + bj.
(46)
1=
c -1
Vj,b,r
Frip 12800 (Wa.. * Uiy.'
+
Frii (R R2180 (W 3) * vi:.) + bj,
(47)
cin
V/,0,,
[2270' (Wit..) * Uigs))
R270"
(WHi...) * oi..) + bj:
(48)
Uj.7,1
(F (R270" (W * Ui
Frip R2270" (W (Wit. vis.) bj.
(49)
C.4. Empirical validation of equivariance for SWE and INS input layers
In the preceding sections, the forms of group equivariant input layers for SWES and INS were introduced. In this section, the equivariance of input layers for SWEs and INS is demonstrated empirically. Figure 11 illustrates the group equivariant input layer of SWEs, specifically the p4 and p4m layers. Given that the output of the input layer possesses four and eight group- group-indexed channels in the regular representations for p4 and p4m, respectively, an average is taken of the representations to produce a single equivariant scalar field output.
21

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
u
V
average(p4)
R(average(p1))
Rotation
ML first layer
rotation
R(()
R(v)
R(u)
average (p4)
ML first layer
Identical
U
v
(verage(p4m)
R( (aaverage(p(mm))
Rotation-filp
ML first layer
fllp rotation
R(F(())
R(F(v))
R(F(M))
averago(p4m)
ML first layer
Identical
Figure 11. Example illustrating equivariance of the SWE input layer for groups p4 and p4m staggered grid input data. The equivariance relation observed hold, numerical precision. this case, a G-averaged regular representation the input layer's output employed to visualize the equivariance.
Figure 12 illustrates the equivariance of p4 and p4m group equivariant input layers for INS. can be seen that the input layer is equivariant in both cases.
22

---

The mathematical formulas of these input layers are given in Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
average(p4)
R( R(average (p4)
u
Rotation
ML solver
rotation
R(0)
R((u)
average (p4)
ML solver
Identical
averago (p4m)
R ( (aa(erraag)())))
U
V
Rotation-fllp
ML solver
flip rotation
R(F(v))
R(-F(ww)
average Average(p4m) (p4m)
ML solver
Identical
Figure 12. Group equivariance of INS input layers for p4 and p4m. equations 42-49.
D. Group equivariant output layer on the C-grid staggering for INS
As shown in Fig. 1, the internal (hidden) layers of our equivariant PDE surrogates comprise modern U-Net. In the case of a scalar field output defined at grid cell centers, direct utilization of the function r2_ trivial_reeppp in esenn is permissible, for example, for Ç in SWEs. However, for vector fields with C-grid staggering, the vector field implementation in esenn, referred to as ac. irrep (1), cannot used because not on the C-grid and does not satisfy symmetries of the discretized PDEs. Therefore, the output layers of p4 and p4m for the vector field are constructed on the C-grid in INS using the following formulations:
p4: U440555 Pi+1..0 Pi...i U4.+0.3 = Pi.++12 PiJ.3
(50)
U140.5. = Pi+1.j.1 Pi,13 + Pi+1, - Pi.j.
p4m:
(51)
U..j+0.5 Pi.j+1,2 Pt.jA + Pi,++166 Ps.j.0
The index i and j denote the spatial position, while l indicates the group element index. i + 0.5 and + 0.5 denote the position of outputs the C-grids for and which are staggered by half grid cell relative p. These layers satisfy the group equivariance property while providing staggered outputs on the C-grid. Fig. (1) show how the output layer make part of the surrogate architecture (blue box).
For regular representations p, the elements Pt.j defined follows:
D.1. Proofs of group equivariance for output layers with C-grid staggering
Equations 50-51 give the formulas defining the output layers for p4 and p4m. In this section, we will demonstrate that the equivariance relations are satisfied as intended.
I. VECTOR OUTPUT LAYER WITH P4 REGULAR REPRESENTATION INPUTS
Given inputs Wn, Vilns the regular representation layer contains four channels indexed by the finite rotational component of p4 symmetry group: P.,.0. P. P..2, and p. When the input undergoes transformation to R(v), - RRuu, the four
23

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
regular representations are converted to p.,.o- P.,.1- P..-2 and P.,.3 respectively. By the assumption of equivariance for all network components from the inputs to p, the following relations hold:
R(p...) P.,,3
(52)
R(p...) = P.,.00
(53)
R(p..2) =P...1
(54)
R(p...)) -P...
(55)
This equivariance relation is of the same form (staggered vector field mapped onto collocated regular representation) as the one investigated empirically in Fig. (2).
To prove the output layers defined by Eq. 50-51 are equivariant, we first consider wider class of linear output layers, identifying constraints on their linear weights equivalent to equivariance, then choose our output layer among these. The wider class of output layers consider (not all of which equivariant), are defined by setting the outputs 0440.5.j + defined at grid cell interfaces to be equal to linear combinations to the regular representation values at the grid cell centers on either side of the interfaces. Thus, we have:
044055.
CoPi+1..h
daPi.jk
(56)
k=O
3
Σ k=0
fuPi.,K
(57)
k=
where Cb. du, ex. and fs are freely chosen coefficients.
Rotate-transforming the inputs, we end up with a transformed regular representation ý that provides input to the final output layer defined in Eq. 56-57, which produces final outputs 1, C. We have equivariance and only if:
(ū, 0) (Tieetooy [(u,
(58)
where 0 < < 4 and TRactor is defined as in Eq. 5. By the definition of the rotation operator's s action on regular representations and vector fields, we therefore arrive at the following constraints:
Ci d2 = 63 f4
(59)
C2 d3 C4 f1
(60)
C3 d4 = f2
(61)
C1 d1 = = f3
(62)
This the set of local linear readout layers (Eq. 56-57) that are p4 equivariant is 4-dimensional, and we could pick any of them or learn 4 coefficients. In this work, opt for simplicity, setting с₁ and c2 03 C4 0. Consequently, the equivariant vector output from the p4 regular representation can be expressed as:
4440.5. = Pi+1..11 - Pi.i3
(63)
U1,+0..5 = Pr.++,2 - Pi.j.4
(64)
It is important note that this is merely particular form of output layer for vector fields the C-grid staggering. Other forms of the output layer are attainable by selecting different values for d,, e, and fi, which can also be learned.
D.1.2. VECTOR OUTPUT LAYER WITH P4M REGULAR REPRESENTATION INPUTS
The proof for this case proceeds similarly to the p4 case. We begin by again considering a class of linear output layers mapping from the collocated regular representation to the staggered velocity fields. But whereas Eq. 56-57 c, d, e f each have four coefficients, for p4m they have eight each according the group dimension of the regular representation:
U:+0.5.j
CaPi+ 1.
da d.Pi.j.k
(655
k=0
k=0
Ui.++0.5
exPi
Σ
fiPij
(66)
k=
k=
24

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
Here have equivariance if and only if:
(ú, 8) = (Tieetorr
( ((,, v)]
(67)
where 0- < < 2,0 < and Tp is the flipping operator for vector fields. Equating elements of the transformed outputs as required for each of the 8 possible transformations, we arrive at the constraints:
C d2 63 /4 = dg e7 fo
(68)
Q d3
61
f1 G6 d7
eg fo
(69)
= dp 6
f2 = = do
E5 = f6
(70)
4 = d1 62 Ja = = d5 66 /T
(71)
As before, any choice of will satisfy equivariance. We again take the simple case of c1 and C2 = the vector output layers on C-grids for and from the p4m regular representation are written as follows:
C1 0. Thus,
4440.5.5 Pi+1..1 PiJ.3 + P.+1j,3 Pi..77
(72)
04,+0.5 Pt,+1,2 Paj.1 Paj+1,6 Pt.0..
(73)
Note that this only one possible form for the output layer.
E. Group equivariance of PDE surrogate architectures
As illustrated in Fig. 13, the group equivariance of a neural network requires equivariance for the input layer, internal layers and the output layer. Fig. 13a shows a network incorporating specialized input/output layers for velocity fields on a staggered C-grid. Here the network exhibits a precise equivariance relationship that holds up to numerical precision. Fig. 13b show the same network, but with a standard equivariant layer designed to operate on vector fields without staggering. Here the equivariance relation no longer holds.
rotated orig higginal P
7++1
@++1
R(a+1) - -R(
input layer
escnn p4
output layer
rotation
Modern U-net
p4 C-grid (Eq. 50)
Rip
Rou
0
identical
R(e²)
I((x))
f
G++1
G-+1
R(
I(
escnn p4
output layer
rotation
vector field escnn
Modern U-net
p4 C-grid (Eq. 50)
iptater
difference
I(('))
-R((e))
Figure 13. Equivariance of neural PDE surrogates with (a) and without (b) specialized input layers for velocity fields on staggered C-grids. Both networks used the same modern u-net architecture, but for (b) standard equivariant vector field input from assumes a collocated grid. Both networks were tested for equivariance in a randomly initialized state, without any training. The same input fields used input both before (upper rows) and after transforming the vector fields using 90 degree rotation (bottom rows). Both networks use an output layer designed for velocity fields on staggered C-grid. For the network in (a), rotation-trant transformation of the output fields produced by original input fields matched the output fields arising from rotation-transfomme transformed inputs up machine precision. For the network in (b), the cquivariance relation did not hold.
used, which
25

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
F. Imposing physical constraints on neural PDE surrogates
In the shallow water system, mass is conserved variable. To ensure the preservation of this variable during training and inference, the mean of the update the mass field is subtracted before the update is applied:
5 S* dç- mean(dC)
(74)
To conserve momentum without conserving mass in the INS system, perform similar global mean correction.
u ++ u' du mean(du); oftl = u² + du mean((e)
(75)
To also impose mass conservation (equivalent to the divergence free property of the velocity field), We learning a scalar potential a and define the velocity field as the curl of this scalar field: [ x(+1] = x [0,0, aj as Wandel et al., 2021.
da Oy
da
u ++ =
w =
(76)
dx
These physical constraint layers are added after the equivariant output layers as needed. An example implementation for INS shown in the blue box in Fig. (1).
As an alternative approach, we might have also learned fluxes at the C-grid interfaces for conserved quantities at the cell centers, or fluxes at the vertices for conserved quantities the interfaces. This approach would be analogous a finite volume solver. key advantage of this approach is that it is computed locally (McGreivy & Hakim, 2023), which avoids unphysical action at large distances and facilitates easier generalization of the domain size after training, We leave this avenue of exploration for future work, with the expectation that further improvements in accuracy could be achieved.
G. Simulation parameters
In this section, we list simulation parameters used to solve the shallow water equations and the incompressible Havier-Stoker equations in the numerical solvers. We also show additional settings used generate training data. Table 6 shows the parameters for solving SWEs. Table 7 shows the parameters used to solve the INS and to generate the training data.
Table 6. Simulation parameters for SWES
Parameters
Explanation
Value
L
simulation domain
1000 1000 (Km)
d
undisturbed water depth bottom drag coefficient acceleration due to gravity space step time step implicit weighting
100 (m)
CD
1.0e - 3
9.81 mm/s²)
9
Az
10 (Km)
At
300 (8)
Wimp
0.5
Table 7. Simulation- and data generation parameters for INS
Parameters
Explanation
Value
simulation domain
2m x 2n
P
density
μ
viscosity
1 3
T
simulation time
224.34
Atsolver
the time step of numerical solver
0.00436
M x M
the grids of numerical solver
576 x 576
Fyalver
the space step of numerical solver
0.0109
At ml
the time step of ML model the grids of ML model the space step of ML solver
0.8375
m m
48 x 48
Azmi
0.1308
26

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
H. Architectural details of (group equivariant) modern U-Net
modified the Modern U-Net (Gupta Brandstetter, 2023) to create a symmetry- and physics-constraines version. Table 8 describes network hyperparameters, and how these were adjusted depending on the chosen symmetry group in order match total parameter counts.
Table 8. Detailed architectural settings for (group equivariant) modern U-Net. In this table, we compare the hyperparameters of modern U-Net with those p4- and p4m--quivvariatt versions.
Hyperparameter
Modern U-net
Modern U-net p4
Modern U-net p4m
Number of time steps the input (time history (int))
Number of time steps in the output (time_future (int))
Base hidden layer channel count (before multiplier)
4
2
2
Activation function
Gelu
Gelu
Gelu
Normalization
True
True
True
List of channel multipliers for each resolution Self-attention used bottleneck
(1,2,4)
(1,3,3)
(1,4,1)
False
False
False
Number of residual blocks each resolution
Kernel size
(3, 3)
(3, 3)
(3, 3)
Padding
(1, I)
(1, 1)
(1, 1)
Padding mode Convolution type for internal layers
Circular
Circular
Circular
Conv2d
ro22d0nR((N=4)
fljpRo22dOmR2(N44)
I. A hybrid method for the inference of shallow water system
Figure shows hybrid method used predict the solution of shallow water system. In our neural PDE surrogate, we have only one output Ç and we have three inputs u, v and Ç. An additional physics-based solution step required compute uf and of from the ç* produced by the surrogate. These additional calculations performed only for autoregressive rollouts with trained networks, not during training. Following Backhaus, 1983, first calculate interim solution u", v² using an explicit time step:
u" u" ACCD u [l"l
A1g
Wimp)
Atan V²
(77)
dr
1 v² Un AtcD
8d |00 Atg(1
Wimp)
Atas 2
(78)
h
Oy
here the superscript n indicates the time step. Next, the interim velocities are combined with the surrogate's estimate of Ç to produce semi-implicit solution for the velocities:
ac++1
n++ u²
AtgWimp
(79)
dr
aea+1
a+l v²
Atgwimp
(80)
Oy
u²
Neural integrator
Numerical solver
Initial condition u
St
2
Figure 14. The figure describes the configuration of the hybrid method for the inference of shallow water systems from a trained model. The fluid surface clevation ç* is predicted by the neural network. The calculation of x' and o² achieved through the application of the following Equations (79-80) with the known value of ç*.
27

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
J. Additional results
J.1. SWE rollouts
Fig. 15 shows predicted surface heights for the shallow water system with additional time steps and models. The primary text presents two models, p1/Ø and p4m/M, which are evaluated against the reference at six time steps, as illustrated in Figure 3-a. Figure 15 extends this by presenting the reference, p1/Ø, p1/M, p1/Ø + e, p4/Ø, p4/M, p4m/2, p4m/M, p FNO/2, and drnet/22 with the reference nine time steps.
These rollouts illustrate only one of thirty distinct initial conditions. In this particular instance, the p4m/M model is the most stable and accurate model in comparison to all other methods. As previously mentioned, the main text also displayed statistical scores calculated over twenty distinct initial conditions, for example in Figure 15c-g.
28

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
0.2h
0.9h
1.8h
2.6h
6.8h
10.0h
16.8h
25.0h
50.0h
reference
p1/0
p1/M
p1/0+E
p4/0
p4/M
p4m/O
p4m/M
FNO/0 p4m/M+E
drnet/0
Figure 15.. Rollout predictions for the fluid surface elevation Ç for all surrogates on the closed boundary shallow water system at various forecast horizons. The predictions of trained models compared the reference. The results obtained from this analysis clearly demonstrate that the p4m/M method exhibits superior performance in terms of stability and accuracy when compared to all other methods. This figure expanded version of Fig. 3a in the main text.
29

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
J.2. Training with input noise for SWE
We compare our unconstrained and physicslyymmnttry--onssaaannddd SWE surrogates, with their noisy variants pl/ø, pl/Ø (Fig. 16). We added Gaussian noise drawn from ([ 0, 0.0001) during the training process. Training with input noise yielded moderately low error for long very rollouts, but yielded lower accuracy compared to the noise-free model for shorter rollouts, especially when constraints were used. Predictions from the noisy model demonstrate reduced accuracy, even at the early stages of the rollout.
10
b
reference
NRMSE-
p1/0
p(S.5)
p1/0+E
p4m/M
p4m/M+6
-0.2
t (h)
50
t th)
50
Figure 16. Comparison of reference simulations and neural surrogates trained with input noise, denoted p1/Ø + and p4m/M++e, as well their their noise-free counterparts p1/2 and p4m/M, based the metrics NRMSE-C and p(ć. <). The NRMSE-C metric reveals that p4m/M++ exhibits a comparatively diminished error over an extended time span in comparison to alternative methodologies. p4m/M delivers precise prediction for the first 25 hours.
J.3. Time evolution of conserved quantities for SWE
Fig. 17 shows the mass, momentum, and total energy for the closed boundary shallow water system for various surrogates and the reference solution over the course of 50 simulated hours 17. The mass and total energy remain constant, while momentum oscillates slightly (black curves). The p4m/M model gives the closets match to the reference solution.
106
6 104
1.94 110
1.000
a
b
C
Momentum
Total energy
ref.
p4/0
Mass
p1/O
p4/M
pl/M
p4m/
p1/0+E
p4m/M
FNO/Ø
drnet/0
0.989
1.942
8 (h)
50
t (h)
50
(h)
50
Figure 17. Mass, momentum and total energy for the shallow water equations for all surrogates tested over the course of 50 simulated hours. The p4m/M model, which has been identified as the most effective model, exhibits superior performance in reproducing the mass, momentum and total energy of the solver.
J.4. Decaying turbulence rollouts
Fig. (18-19), show rollout of velocities and v for additional time points and surrogates, extending Fig. 4 from the main text. We compared the predictions velocity u and v from thirteen surrogates with the reference solution: p1/2, pl/pú, pl/M+pü, pl/Ø + e, p4/2, p4/pū, p4/M+pū, p4m/2, p4m/pü, p4m/M+pü p4m/M+pü + e, FNO/Z and dmet//2. These results show just one of 30 initial conditions.
30

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
0.8s
10.1s
19.3s
28.5s
37.7s
46.9s
56.1s
65.3s
74.5s
83.8s
93.0s
99.7s
reference
p1/o
p1/0+E p1/M+pü p1/p
p4/0
paspi
p4ml0 p4/M+pi
p4m/pii
p4m/M+ p4m/M
pi
Pi
FNO/Ø
drnet/0
Figure Rollout predictions for u from thirteen models with approximately M. parameters each for decaying turbulence at varying time steps. The top row shows the reference simulation. This plot extends Fig. 4 from the main text.
31

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
0.8s
10.1s
19.3s
28.5s
37.7s
46.9s
56.1s
65.3s
74.5s
83.8s
93.0s
99.7s
reference
p1/0
p1/0+E p1/M+pü pi/pi
p4/0
p4/p
p4m/0 p4/M+pü
pü p4m/pt
FNO/O p4m/M+ p4m/M ++
pu+
drnet/o
Figure 19. Rollout predictions for v from thirteen models with approximately 0.1M parameters each for decaying turbulence at varying time steps. The top row shows the reference simulation.
32

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
J.S. Momentum in surrogate rollouts for decaying turbulence
Figure. 20 shows reference and predicted momentum for various surrogates. The non-physiiaallyyconstaaaee models (pl/2, p4/2, and p4m/2) exhibit substantial increase in momentum compared the reference.
0.1 X
x10-2 2
b
Momentum
Momentum
Momentum
reference
reference
reference p1/M+pi
pllpi
-10
(s)
100
(s)
100
f (()
-110
x10
e
Momentum
Momentum
Momentum
reference p4/0
reference p4/pūž
reference
p4/M+pu
(()
100
(s)
100
t(s)
100
10
X10-1
h
Momentum
Momentum
Momentum
reference p4m/O
reference
reference
p4m/M ppi
(s)
100
-50
(8)
100
(8)
100
Figure 20. Momentum over time for all surrogates the decaying incompressible turbulence task, compared to momentum of reference solution (black).
J.6. Velocity and Energy Spectra for Decaying Turbulence
Figure (21) shows the velocity and energy spectra (u(k)k² and E(k)&²) at additional time steps, expanding on the two time steps previously shown in Figure 4-d,e of the main text. Our analysis demonstrates that the model p4m/M+pü (blue curve) consistently provides the closest match to the reference across all time steps for both the velocity and energy spectra.
33

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
0.8s
28.5s
56.1s
99.7s
Velocity spectrum
10
10a
a
10a
109
d
ref.
p1/0+E
10
pl/0
p4/0
(k)
10
10
10
107
10
p1/pū
p1/M+pū
p4lpä
p4m/o
10%
10%
p4/M+pii
p4m/pü
106
p4m/M+pt
105
10
10
109
10
10 &:
10
10a
10"
10
10
10
Wavenumber
0.8s
28.5s
56.1s
83.8s
99.7s
Energy spectrum
109
1094
10
10
10t1
h
1044
10
ks)
B(k)
10
108
10
10
10
107
107
107
10
106
109
10y
10
10
10'
Wavenumber k
Figure 21. Velocity and energy spectra at additional time points. Results are consistent with the those presented in the main text, specifically Figures (4-d.e). p4m/M+pr (blue curves) most closely matches the reference spectra (black).
J.7. Training with input noise for decaying turbulence
We investigated the effect of input noise during the training process for the decaying turbulence task (Fig. 22). We trained with input noise drawn from N(μ 0, 0.001). Training with input noise did not produce clear increase in accuracy for unconstrained surrogates.
NRMSE-U N
1.0
a
reference
p u)
p1/O
p1/0+E
p4m/M+ pii
p4m/M+pū
-0.4
t (s))
100
(()
100
Figure 22. Comparison of surrogates trained with noise clean data: pl/@ e, p4m/M+pu e, p1/0 and p4m/M+pūū. We find that training with noise did not bring about clear improvements to accuracy, though unconstrained models were slightly more accurate for long rollouts.
J.8. Dilated ResNet as surrogate for decaying turbulence
In the main text, we used the robust Modern U-net architecture for the study of INS. In this section, an alternative architecture, DilatedRes Net Stachenfeld et al., 2021; Kohl et al., 2024, is modified to incorporate equivariance with respect to p4 and p4m. The input layer, output layer, and physical constraints of these models are consistent with those outlined in the main text. In this section, focus investigating nine Dilated ResNet models: p1/Ø, pl/pii, pl/M+pü, p4/2, p4/púi p4/M+pri, p4m/Ø, pám/pnū p4m/M+prü.
Fig. (23-24) present a single instance ofvelocities and predicted by nine Dilated ResNet models, as well as the reference simulation. The network size of these models is approximately 0. M. which comparable to the modern U-Net in Figure 4-a. The p4m/M+pū with Dilated ResNet exhibits the best rollouts for both u and when compared to the reference, as illustrated Figures (23-24).
34

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
0,8s
10.1s
19.3s
28.5s
37.7s
46.9s
56.1s
65.3s
74.5s
83.8s
93.0s
99.7s
reference
p1/0
p4/0 p1/M+pi pi/pi
pA/pi
p4/M+pü
p4m/
p4m/M4 p4m/pu pi
Figure 23. Rollouts u by nine Dilated ResNet surrogates for decaying turbulence. The network size these surrogates was approximately O.I. parameters. In this particular instance, the p4m/M+pūū model is the most similar the reference solution.
35

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
0.8s
10.1s
19.3s
28.5s
37.7s
46.9s
56.1s
65.3s
74.5s
83.8s
93.0s
99.7s
reference
p1/0
p1/M+piL p1/pi
p4/0
p4/pi
p4/M+pË +pr
p4m/o
p4m/M4 p4m/pi pii
Figure 24, Rollouts by nine Dilated ResNet surrogates for decaying turbulence, same data Fig.
Over 30 random initial conditions, we computed accuracy metrics for these dmet surrogates (Fig. 25). The worst dilated ResNet model was p1/0. and the best was p4m/M+pü.
NRMSE-- 15
1.0
a
b
ref.
p4/0
p( u)
p1/O
p4/pū
p1/pi
p4m/0
pl/M+pē
p4m/pt
p4/M+pū
p4m/M+pu
-0.4
((s)
100
t (()
100
Figure 25, NRMSE-U u and p(ū, u) for the Dilated ResNet surrogates on the decaying turbulence task. Results are derived from 30 random
initial conditions.
We also computed velocity and energy spectra (u(k)&5 and E(k)k*) for the Dilated ResNet models decaying turbulence. Doubly constrained surrogates reproduced these spectra most faithfully (Fig. 26.
36

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
0.8s
28.5s
56.1s
99.7s
Velocity spectrum
10"
a
109
b
10t5
10ah
10
10
ref.
p4/0
ks
10
10
p1/0
p4/pi
10%
10 19
pl/pt
p4m/0
10
p1/Mapu
p4m/ pti
109
106
10
p4/M-pü
p4m/M+pii
106
10
10a
100
10'
10P
10
100
10
10
Wavenumber
0.8s
28.5s
102s
56.1s
1027
83.8s
99.7s
Energy spectrum
e
10
10
109
1039
10
10a
B(k)k
109
10
10
10
10a2
10
10
10
10
10
10a
107
10
10
10th
10'
10
10
10"
10
10"
10
109
10
Wavenumber h:
Figure 26. Velocity and energy spectra (u(b)k² and ((:))=)) for the Dilated ResNet models of decaying turbulence.
J.9. Pushforward Training
Here we analyze surrogates trained with the pushforward trick (Brandstetter et al., 2022c) in greater detail. We consider the modern U-Net base architecture and its constrained variants on the decaying turbulence task. Pushforward training employs an auto-regressive rollout, using model outputs as inputs for the next time step time step, does not propagate gradients backwards in time. Following Brandstetter al., 2022c, roll for time steps.
With PF training, the maximally constrained model p4m/M+püü+PF exhibited significantly reduced NRMSE compared to less constrained models, while the improvement was less clear for correlation.
1.0
NRMSE 22
b
ref,
p4/0+PF
pl u)
p1/0+PF
p4/pūt+PF
pl/pu+PF
p4m/O+PF
p1/M+pÜ+PE
p4m/putPP
p4/M+pü+PP +
p4m//+pp-+PF
0.4
((8)
100
t(s)
100
Figure 27. NRMSE-u and p(û, u) for pushforward training of PDE surrogates.
J.10. Data augmentation
Data augmentation is a regularization technique that uses transformed copies of data points train a neural network, to reduce overfitting and improve generalization. It has been employed train PDE surrogates (Brandstetter et al., 2022b; Fanaskov et al., 2023). We augment with random rotations, as well as flipping with probability 0.5. The transformation is applied to both inputs and outputs for each training data point, with transformations applied to vector fields as previously described.
As expected, data augmentation had no effect on fully equivariant surrogates, but did bring about some improvement non equivariant models (Fig. 28).
37

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
1.0
NRMSE 21
b
ref.
pI/pũ+DA
p(ú. ④)
p1/0+DA
p1/M+pŪ++A
p1/pũdDA
-p4m/O+DA
p1/M+pÜ+DA
--p4m/pū +DA
p1/0+DA
-4mm/++Ø+DA
p4m/M+pu
0.40
((s)
100
t (8)
100
Figure 28. Performance data dala-augmented models.
J.11. Generalization tasks
We tested surrogates on additional ICs that present in their training datasets, for the SWE and INS cases.
J.11.1 GENERALIZAAIION FOR SWES
We evaluated generalization capability for SWE surrogates in 3 experiments with 20 ICs each. First, we used a single rectangle with a size of 0.1 m as the IC, with its location chosen at random (Fig. 29). Second, we used the sum of two rectangles, each measuring 0.1 m, with their locations selected at random (Fig. 30). Finally, we used two overlapping rectangular elevations (summed over the region of overlap, Fig.31), which proved the most challenging configuration. The p4m/M surrogate proved superior in each case.
38

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
0.2h
4.0h
10.0h
25.0h
33.3h
41.7h
50.0h
reference
p1/0
p1/M
p4/0
P4/M
p4m/@
p4m/M
Figure 29, Rollouts demonstrating generalization for the SWEs, generated by all surrogates from a single rectangular-shan shaped elevation IC, shown various time intervals.
39

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
0.2h
2.0h
4.0h
10.0h
25.0h
33.3h
41.7h
50.0h
reference
p1/0
p1/M
p1/0+E
p4/0
p4/M
p4m/0
p4m/M
Figure Rollouts demonstrating generalization for the SWEs, generated by surrogates from a single consisting of two rectangular elevations, are shown at various time intervals.
40

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
0.2h
2.0h
4.0h
10.0h
25.0h
33.3h
41.7h
50.0h
reference
p1/0
p1/M
p1/0+E
p4/0
p4/M
p4m/O
p4m/M
Figure 31. Rollouts demonstrating generalization for the SWEs, generated all surrogates from single IC consisting of two overlapping (summed) rectangular elevations, are shown at various time intervals.
J. 1.2. GENERALIZATION FOR DECAYING TURBULENCE
We evaluated generalization capability for INS surrogates, using 10 ICs with peak wavenumber 8, Figs. 32-33 show examples of the reference solution and surrogate rollouts for this data.

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
0.8s
10.1s
19.3s
28.5s
37.7s
46.9s
56.1s
65.3s
74.5s
83.8s
99.7s
reference
p1/o
pl/pā
p1/M+pï
p1/0+E
p4/o
p4/pī
p4/M+pū
p4m/o
P4m/M² p4m/pū pü
Figure 32. Rollout performance of networks with various physical and symmetry constraints for the generalization test of decaying turbulence. The figures depict the cvolution of the field variable u.
42

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
0.8s
10.1s
19.3s
28.5s
37.7s
46.9s
56.1s
65.3s
74.5s
83.8s
99.7s
reference
p1/o
p1/pi
p1/M+pū
p1/0+E
p4/0
p4/piž
p4/M+pï
p4m/o
p4m/M+ p4m/piž pu
Figure 33. Rollout performance of networks with various physical and symmetry constraints the generalization test of decaying turbulence. The figures depict the evolution of the field variable
Fig. 34 shows velocity and energy spectra for the reference solution and surrogates in the generalization test for INS. These spectra are expanded versions of those in Fig. 5e-f in the main 4xxtp4m/M++pi (blue curves) exhibits the most precise alignment with the reference spectra (black) for both velocity and energy.
43

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
Velocity spectrum
10
a
0.8s
109
b
28.5s
10
56.1s
10
99.7s
10a
10a
10
10
p1/0
ref.
p4/o
pI/0+ē
(kk)
102
10?
10
pl/piž
p4lpū
106
10%
106
10
p1/M+pä
p4/M+pü
p4mipi
p4m/0
104
p4m/M+pii
10
10
10
10"
10
10"
10
100
10
Wavenumber k 0.8s
10
28.5s
56.1s
83.8s
99.7s
Energy spectrum
10
f
g
1003
h
10
109
10
10t1
10"1
10
10
108
10
10
10
10
107
10
10
109
105
10
10
10'
10
10
10
10
10
10
10
10
Wavenumber
Figure 34. Velocity (u) and energy power spectra over extended range of rollouts for the generalization of decaying turbulence case. This analysis extends the findings presented in Figure 5-(c.f) of the primary text. p4m/M+pri exhibits the best match to the energy and velocity spectra reference solutions.
J.12. Effects of network and dataset size on rollout performance
Fig. 35 show the effect of network and training dataset size rollout performance, using INS surrogates based on modern U-net with 0. MM, 2M and 8.5M parameters. We consider training dataset sizes 100, 400 and 760 unique ICs. The findings indicate that increasing the network size or the training dataset size enhances the rollout performance of the network. Additionally, networks with physical and symmetry constraints exhibits superior performance in each case.
44

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
Effect of networks size
1.0
b
p1/0 90114 parameters
p1/O 2224642 parameters
NRMSE-U
p w
p1/0 8884482 parameters
p4m/M+ 86995 parameters
p4m/M+p 2111619 parameters
p4m/M+pü 8415199 parameters
0
(s)
100
{s)
100
0.8s
28.5s
56.1s
83.8s
99.7s
Energy spectrum
1040_
e
10
10a4
g
109
108
10
10
[k
10%
105
109
10
10a
10
107
106
10%
109
10a
10c
100
10
10
10
10
10
10
10
10
10
Wavenumber h:
Effect of data size
1.0
h
p1/0 100 simulations
NRMSE-U
p4m/M+pū 100 simulations
plu ul)
p1/0 400 simulations
p4m/M+pül 400 simulations
p1/0 760 simulations
p4m/M+pü 760 simulations
-0.4
(s)
100
(8)
100
0.8s
28.5s
56.1s
83.8s
99.7s
Energy spectrum
g
k
10
m
1013
n
10
10
10%
10
10t1
E: A.
108
10
109
107-
10a
107
107
10
106
105
10a
106
10
10
10
10
10
10
10
10
10
10
Figure 35. Top: effect network size on NRMSE-u, p(ú, u) (a, b) and the energy spectrum (e-g) PDE surrogates. Bottom: effect training data size NRMSE-1, p(û, u) (h, i), and the energy spectrum (g.s). All results reported for 99.7s rollouts.
J.13. Inference time per time step pl/0 and p4m/M+pü at various network sizes
Table J.13 reports inference time per time step of pl/Ø and p4m/M+pül with various network sizes on CPU (Intel Xeon Platinum 8160) and GPU (Nvidia A100 40 GB) nodes. Inference speed was computed from an autoregressive rollout over 120 time steps. On the CPU, p4m/M+pü is consistently slower than the speed of p1/Ø for all network sizes, but the GPU both networks are dominated by overhead costs that do not scale with network size to 8.5 M parameters. Inference on the CPU is slower than on the GPU for both pl/Ø and p4m/M+pü. These results indicate that the advantages conferred by physical and symmetry constraints do not adversely impact GPU-based predictions, at least for the domain and network sizes tested. The relative independence of GPU inference speed on network size suggests that overhead costs, kernel launches and memory transfers are likely bottlenecks, and further network scaling and optimization would be required to identify the true
45

---

Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates
costs of constrained unconstrained inference in the limit of large network sizes.
Table Inference time per time step of p1/2 and p4m//M+pü on CPU (Intel Xeon Platinum 8160) and GPU Vidia A100 40 GB) nodes for various network sizes. Inference speed computed for autoregressive rollout 120 time Steps.
pl/Ø
p4m/M+pü
Networks size
CPU
GPU
CPU
GPU
0.1 M
9.440.9ms 26.4:14. Ims
6.2:t0.7ms
20.6::3.2ms
8.2 =m,4ms
2.0 M
6.3:0.33ms
71.2::59.5ms
12.0:t3.3ms
8.5 M
27.5++3.8ms
6.4-H002ms
145.7=-9..7s
8.4-1-0.1m
J.14. Real Ocean Dynamics
We used the ocean current velocity data sourced from the Global Ocean Physics Analysis and Forecast (Marullo et al., 2014), and followed (Wang et al., 2021) for data selection and processing. The 6-hourly data from 2022-06-01 to 2025-04-05 was selected for analysis from three different ocean regions. The corresponding latitude and longitude ranges for these regions are listed below: (30 -42, 168 -100), (-49-0 -37, 80 -92), and (-51 -39, -30 -18). The data set under consideration has 144 x 144 pixel velocity fields. In order to reduce the size of the training set, was downscaled to a coarse grid of 48 x 48 for training, testing and validation purposes.
Once surrogate training was completed, forecasts were made from a different region and time period. The latitude and longitude ranges for prediction were (-44 -32, -130 - 118), and the time range 2023-06-01 to 2025-04-05.
46


Reviews:
[{"review_id": "xohTZz0n80", "reviewer": "Reviewer_ehyY", "summary": "This paper propose to integrate rotation symmetry of staggered grid into PDE surrogate models. Additionally, the models also encode physics constraints in the network readout. The experiments are conducted on closed shallow water equations and decaying turbulence.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the careful reading and constructive feedback.\n\n> technical challenges of extending group CNN to staggered C-grid not sufficiently highlighted\n\nWe agree and have revised and expanded the last sentence of the paragraph labeled \"Staggered Grids\" in sec. 2 to read as follows: \n\n\"However, current software implementations of equivariant network layers cannot be applyed to PDE variable fields on equivariant grids. This is because they assume that variables fields are all located at the same points, allowing the action of symmetries on these fields to be broken into two steps: a resampling step $x \\rightarrow g^{-1}x$ carried out on the grid itself, and a transformation step $w\\rightarrow \\rho_g(w)$  carried out on PDE field variables $w\\in\\mathbb R^m$ at each point single grid point. This leads to overall transformations $w\\rightarrow \\mathcal T_g w$, such that $[\\mathcal T_g w] (x) = \\rho_g(w(g^{-1}x))$. This is a valid assumption for PDEs on continous spatial domains (eq. 5) or for colocated grids (Weiler, 2021, eq. 1). But for staggered grids, the PDE fields are not represented as a vector of values $w(x)$ at each grid point $x$. Instead, each field is defined at different locations, which may be grid cell centers, interfaces or verticies. Thus, the spatial transformation of the grid and the transformation of local field values cannot be disentangled. Applying existing equivariant network layers to staggered PDE fields therefore breaks symmetry constraints.\"\n\nWe also added a [new supplementary figure](https://tinyurl.com/2rn8ydp5), showing how equivariance breaks down when applying previous equivariant input layers to input data on a staggered grid.\n\n> Moreover, how to enforce the conservation laws is only briefly mentioned in the method section.\n\nWhile space constraints prevent us from giving full details in the main text, we have revised the relevant material for clarity and to more explicitly list the types of constraints we can impose while maintianing equivariance on staggered grids. We now include in section 3 a summary of the full stragegy described in appendix E. The paragraph beginning \"conservation laws\" now reads:\n\n\"We impose 3 types of conservation laws as hard constraints. For scalar quantities such as fluid surface height $\\zeta$, we subtract the global mean of $\\zeta^{t+1}-\\zeta^t$ at each time step. For vector fields, we subtract the mean of each velocity component. As mass conservation in incompressible flows is equivalent to divergence-free velocity fields, we impose this by learning a vector potential $a$ defined at grid vertices, and compute velocities at grid cell interfaces as the curl $\\nabla \\times a$ to satisfy both mass and momentum conservation (Wandel et al., 2020). Further details and discussion of alternative approaches are found in appendix E.\"\n\nWe also discuss other possible physical constraints in the discussion.\n\n> Section 4.1: \"Simulations in Fortran required 67 seconds on the CPU\", is the time for simulating one trajectory or the training set?\n\nFor a one trajectory, now clarified.\n\n> Table 3: from the result (especially at 25h) it looks like the baseline models barely work. Although I think symmetry is important but the effect is rather surprising.\n\nThis is a challenging task and no models tested were accurate over the full 50 simualted hours. NaN values in the table indicate that some examples from the test set diverged to infinity, which is now clarified. Symmetry was essential to achieving accurate results at 25h, but physical constraints were also important, as rotation-reflection equivariant without physical constraints also diverged within 25h. As shown in Fig. 3f, this can partly be explained by the fact that in many non-mass-conserving networks total mass diverged to infinity. We now mention this in the discussion.\n\n> The authors could consider moving figure 7b and description of staggered C-grid to the main text.\n\nWe agree and plan to add additional detail on the staggered C-grid to fig. 1. \n\n> Minor: is this the right paper template? The footnote at first page and the running title seem to be missing.\n\nWe appreciate the notice and will correct this if accepted.\n\n> Assuming one has a good understanding of group equivariant CNNs, how would the authors explain the extension to staggered C-grid?\n\nStaggered grids invalidate the assumption that symmetries can be brokedn into a transformation of space and a transformation of PDE fields at each point. Equivariance therefore requires new input/output layers that take staggering into account.\n\n> Are the physics conservation laws hard to implement? I wonder since it is not described too much in the method section but it is emphasized a lot in the experiments.\n\nWe have added further detail to the methods. Conservation laws are easy to implement by correcting or re-interpreting network outputs, without considering network architecture or equivariance. We will provide code if accepted.", "claims_and_evidence": "- The motivation of using staggered C-grid is clear and designing equivariant models for it is well-motivated.\n\n- However, I feel the presentation of the method is not very clear. Since I am not very familiar with the staggered C-grid. The technical challenges of extending group CNN to staggered C-grid are not sufficiently highlighted for me to get a clear understanding. Moreover, how to enforce the conservation laws is only briefly mentioned in the method section.", "methods_and_evaluation_criteria": "- The benchmarks make sense to me.", "theoretical_claims": "- They look good to me.", "experimental_designs_or_analyses": "- Section 4.1: \"Simulations in Fortran required 67 seconds on the CPU\", is the time for simulating one trajectory or the training set?\n\n- Table 3: from the result (especially at 25h) it looks like the baseline models barely work. Although I think symmetry is important but the effect is rather surprising.", "relation_to_broader_scientific_literature": "- They seem to be adequately discussed.", "essential_references_not_discussed": "- Not I am aware of.", "other_strengths_and_weaknesses": "- It is interesting to use closed boundary for SWEs, which make them more challenging.\n\n- The baseline models are reasonably chosen.", "other_comments_or_suggestions": "- The authors could consider moving figure 7b and description of staggered C-grid to the main text.\n\n- Minor: is this the right paper template? The footnote at first page and the running title seem to be missing.", "questions_for_authors": "- Assuming one has a good understanding of group equivariant CNNs, how would the authors explain the extension to staggered C-grid?\n\n- Are the physics conservation laws hard to implement? I wonder since it is not described too much in the method section but it is emphasized a lot in the experiments.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JU6ahPSFso", "reviewer": "Reviewer_YKe6", "summary": "The authors propose new input layers that can add inductive symmetry and conservation-law biases to neural PDE solvers to improve their performance in long-term rollouts. The main innovation of the work seems to be the ability to accommodate staggered grids commonly found in CFD. Other than this, the novelty component of the work is low. Its main contribution is a high-quality scientific computation study of tough CFD problems (the shallow water equation with close boundaries and decaying incompressible turbulence) using neural PDE solvers.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive assessment of our work and the recognition that we used challenging tasks.\n\nWe disagree, however, that the novelty component of the work is low overall. It is certainly true that the methods we introduce, equivariant input and output layers for staggered grids, are not revolutionary. However, we would argue that the major novelty and signficance of our work lies in our results themselves: it was known that both physical and symmetry constraints can improve long term accuracy of PDE surrogates, but to date there have been practically no results showing how well these approaches could be combined. Indeed, at the start of this project we were fully prepared for a negative result, in which one set of constraints made the other redundant. This would be, intuitively at least, consistent with the deep connection between symmetries and conservation laws in physics expressed by Noether's theorem (which admittedly does not hold on for discrete symmetry groups), and the fact that most numerical solvers are equivariant. In the end, we arrived at several important discoveries:\n* We can indeed fruitfully combine both constraint types.\n* We can go even further by combining them with other strategies for long-term accuracy such as pushforward training (Fig. 4g).\n* These results hold for multiple network sizes and architectures.\n* These results extend to real-world observational data.\n* The combination of symmetry and physical constraints improves generalization (Fig. 5).\n\nThese results are promising and of high relevance for PDE surrogate tasks requiring long-term accuracy or where training data is in short supply, such weather forecasting, climate projections and airfoil design. We now better emphasize the importance of these novel results in the introduction and discussion.\n\n> Would the use of FORTRAN be consistent with this kind of computationally-intensive data-driven work? I am just curious.\n\nThe constraints we have described are applicable in any programming language. ML in FORTRAN has several modern libraries and bindings, and can exchange data with python processes [1,2,3].\n\n> Can the authors define more clearly what they mean by \"equivariance\"? Would that be symmetry-group invariance?\n\nWe say a function is equivariant when it respects a set of symmetry constraints. That is we, refer to equivariance with respect to a group of symmetries specific to a PDE, as described in the \"symmetry equivariance\" paragraph of section 2 (see eq. 3, with specific examples in eq. 4-5). Following Cohen & Welling [4], we define equivariance for a PDE surrogate $\\mathcal M:w^t\\rightarrow w^{t+1}$ as the relation $\\mathcal T_g \\circ \\mathcal M(w^t) = \\mathcal M (\\mathcal T_g w^t), \\forall g\\in G$. Here $G$ is a group of symmetry transforamtions, and $\\mathcal T_g$ is the transformation on the set of PDE fields described by the group element $g$ (for example, a ninety-degree rotation). Equivariance means that every transformation of the inputs of $\\mathcal M$ leads to a corresponding transformation of the outputs. Equivalently, we say that $\\mathcal M$ respects the symmetries described by $G$. In general we try to endow our surrogates with the same symmetries as the numerical solvers and PDEs they are learning from.\n\nWe have added the clarifying sentence to the \"symmetry equivariance\" paragraph of sec. 2: \"That is, transforming the inputs of $f$ will transform its outputs correspondingly.\"\n\n[1] Brenowitz, Noah. Calling Python from Fortran (not the other way around). https://www.noahbrenowitz.com/post/calling-fortran-from-python/. 2022\n\n[2] Zhang, Tao, et al. \"A Fortran-Python Interface for Integrating Machine Learning Parameterization into Earth System Models.\" Geoscientific Model Development Discussions 2024 (2024): 1-26.\n\n[3] Arnold, Caroline, et al. \"Efficient and stable coupling of the SuperdropNet deep-learning-based cloud microphysics (v0. 1.0) with the ICON climate and weather model (v2. 6.5).\" Geoscientific Model Development 17.9 (2024): 4017-4029.\n\n[4] Cohen T, Welling M. Group equivariant convolutional networks. International conference on machine learning 2016 Jun 11 (pp. 2990-2999). PMLR.", "claims_and_evidence": "Yes, the claims are supported by high-quality experiments.", "methods_and_evaluation_criteria": "The authors evaluate their approach on two difficult CFD problems.", "theoretical_claims": "There aren't any theoretical claims.", "experimental_designs_or_analyses": "I checked briefly the experimental design and it seems adequate.", "relation_to_broader_scientific_literature": "The authors do a good job positioning the work with respect to the literature.", "essential_references_not_discussed": "None noted.", "other_strengths_and_weaknesses": "The paper is clearly written in a adequate technical style.", "other_comments_or_suggestions": "Nothing to add here.", "questions_for_authors": "Would the use of FORTRAN be consistent with this kind of computationally-intensive data-driven work? I am just curious.\n\nCan the authors define more clearly what they mean by \"equivariance\"? Would that be symmetry-group invariance?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UZTWNej0A1", "reviewer": "Reviewer_tGcM", "summary": "The paper explores how incorporating symmetric constraints and physical priors can improve predictions within the same base architecture. Specifically, it investigates the effects of integrating additional symmetry equivariance into convolutions—such as rotation and reflection—in combination with conservation laws (e.g., mass and momentum conservation) that align with the underlying dynamics. The study employs also complementary techniques like data augmentation and the push-forward trick, where backpropagation is restricted to the final step, to enhance performance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the careful evaluation and appreciate the positive assessments therein.\n\nWe have revised the manuscript to incorporate the real-world dataset from Wang et al. 2020, and added a new [figure](https://tinyurl.com/3u4a8wnu) and [table](https://tinyurl.com/mvz9u7f6). Similar to our results on simulation-based datasets, we find equivariant and physically constraints networks are more accurate than the same archictures without constraints and similar parameter counts. We achieved better accuracy than the equivariant network proposed in Wang et al. 2020, and also outperform it on our Naiver Stokes task, [updating fig. 4](https://tinyurl.com/2j8txw3n).", "claims_and_evidence": "- Comprehensive study on symmetries and prior knowledge: The paper systematically examines the extent to which different combinations of symmetries and physical priors influence prediction performance, both within the training trajectory horizon and in extrapolation beyond it.\n- Enhanced generalization: By integrating symmetry constraints and physical priors, the approach improves the model’s ability to generalize beyond the training data.", "methods_and_evaluation_criteria": "- Validity of the methods: The proposed methods are reasonable when the physical prior is known, ensuring that the imposed constraints align with the underlying dynamics.\n- Empirical evaluation: The study employs two synthetic datasets to systematically analyze the effects of different symmetry and prior combinations on prediction performance.", "theoretical_claims": "No.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "The key contributions of the paper is related to the integration of prior information for neural surrogate model learning.", "essential_references_not_discussed": "The related work is sufficient comprehensive.", "other_strengths_and_weaknesses": "Strengths\n- Compared to previous works, particularly Wang et al. (2020), this study conducts a more extensive analysis with finer-grained comparisons of different symmetry and prior combinations.  \n\nWeaknesses\n- Unlike Wang et al. (2020), the study does not include real-world data, limiting its direct applicability to practical scenarios.  \n\n\nReferences:\n- Wang et al. (2020), Incorporating symmetry into deep dynamics models for improved generalization.", "other_comments_or_suggestions": "- The study could benefit from validation on real-world datasets, such as sea surface temperature or atmospheric dynamics, to assess the generalizability of the proposed methods.", "questions_for_authors": "See suggestions above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oqbzn9pji4", "reviewer": "Reviewer_hyP8", "summary": "The authors propose a neural PDE surrogate solver that respects the rotation and reflection equivariance (via p4/p4m symmetry groups) and enforces physical conservation principles. Their approach is designed for scalar and vector field magnitudes on staggered grids, leveraging a modern U-Net architecture with group convolutions [Cohen & Welling, 2016]. Equivariance is enforced through custom input and output layers, while mass and momentum conservation are achieved by predicting a vector potential (for divergence-free conditions) and applying global mean corrections. The method is evaluated on two PDE systems: shallow water equations and incompressible decaying turbulence. Compared to baselines without these constraints or relying on data augmentation, the proposed model demonstrates improved accuracy and long-term stability.\n\n===========\n\nPost-rebuttal: I think the authors have addressed all my issues, and the work represents a neat investigation of equivariance for U-net architectures. There are really novel aspects here, and the approach seems useful. Hence, I fully support an accept of this paper for ICML.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the careful reading, positive assessments and constructive feedback.\n\n> comparisons to other equivariant models (e.g., [Wang et al., 2020]) are missing. \n\nWe now compare to the equivariant network of Wang et al., 2020 on our simulation-based INS task, [updating fig. 4](https://tinyurl.com/2j8txw3n). We also compare to it on the real-world ocean current dataset used in that paper, with a [new figure](https://tinyurl.com/3u4a8wnu) and [table](https://tinyurl.com/mvz9u7f6).\n\n> Additionally, this method is likely slower than a non-equivariant U-Net with the same number of weights. To provide a fairer assessment of practical trade-offs, the authors could compare their model against a non-equivariant counterpart with an equivalent inference time instead.\n\nWe now include inference speed of unconstrained and doubly-constrained networks in a new [table](https://tinyurl.com/44w2aubv). CPU implementations of equivariant networks were about half as fast, and inference time grew sublinearly with network size. For GPU implementations, no consistent relationship was observed, and equivariant networks were about 30% slower. These results suggest that overhead costs, kernel launches and memory transfers are likely bottlenecks, and further analysis and optimization would be required to fairly and quantitatively compare true inference speeds. We now discuss these issues further in our discussion section.\n\n> Unclear: “Boundary effects interfere with translation equivariance, so we provide a boundary mask input channel.”\n\nWe agree and have moved this sentence to the previous paragraph, and revised to: \"Since the time evolution of this SWE systems depends on the location of boundaries, we provide a binary boundary mask to the network as an additional input field with scalar values defined at grid cell centers. We note that this binary mask is invariant to rotations and reflections.\"\n \n> Unclear: “Momentum is not conserved due to reflection from closed boundaries.” It would be helpful if the authors could clarify these points.\n\nRevised for clarity: \"Momentum is not conserved in this SWE system, and a wave travelling eastward will reverse and head westwards after reflecting from a boundary. In reality, this momentum change would be compensated by a slight change in the momentum of the Earth itself, but this is not modeled in our simulation.\"\n\n> The advantage of using a staggered grid over the more common approach with CNNs is not clearly justified.\n\nStaggered grids are prevalent in atmospheric [1] and ocean models [2], largely due to the numerical advantages of finite volume approaches, especially for conservation laws [3]. We now discuss these advantages in greater detail in the introduction.\n\n> This work primarily focuses on p4m equivariance on grids, but a large body of research exists on E(2)/E(3) and SE(2)/SE(3) equivariance in unstructured grids. The following references are relevant but not discussed: ...\n\nWe agree and now included them all, along with several papers from 2024-2025.\n\n> Simple geometric test cases. Applying this method to more complex fluid domains (e.g., irregular obstacles) would better showcase its advantages.\n\nWe agree that more complex geometries, as well as the challenge of generalization to new geometries (c.f. Wandel et al., 2020), are important open questions. While beyond our current scope, we now discuss this explicitly in \"limitations and future work.\"\n\n> How does the use of staggered grids improve performance compared to standard CNN approaches?\n\nWe do not claim staggered grids are always better, and our networks do not use them for the internal representations of hidden layers. But when input data/target outputs are defined on a staggered grid, only input/output layers taking this into account can maintain equivariance, and standard libraries such as escnn will not do so. Thus, to the extent that equivariance improves performance, staggered grids must be accounted for.\n\n> Could momentum conservation be enforced pointwise instead of in an integral sense? Would that improve accuracy?\n\nWe suspect this would improve performance, especially when generalizing to new domain sizes. It would also ensure that surrogates learn causal relationship instead of memorizing statistical patterns. We expanded our discussion of this in appendix E, and now mention it in the discussion.\n\n> What is the computational cost of your method compared to a standard CNN with an equivalent number of parameters?\n\nWe address this in a [new table](https://tinyurl.com/44w2aubv) (see answer above).\n\n[1] Giorgetta, et al. \"ICON‐A, the atmosphere component of the ICON earth system model: I. Model description.\" Journal of Advances in Modeling Earth Systems, 2018.\n\n[2] NEMO Ocean Engine Reference Manual v 5.0. Madec et al., 2024.\n\n[3] Ferziger, Joel H., Milovan Perić, and Robert L. Street. Computational methods for fluid dynamics. springer, 2019.", "claims_and_evidence": "The authors claim that incorporating symmetry constraints (rotation/reflection equivariance) and physical conservation laws (mass/momentum conservation) improves the accuracy and stability of neural PDE surrogates. They further argue that combining both constraints leads to the best performance, even compared to strong baselines that use data augmentation or pushforward training.\n\nTheir empirical results support these claims, with particularly notable improvements in long-term stability.\n\nHowever, while the results are compelling, comparisons to other equivariant models (e.g., [Wang et al., 2020]) are missing. Additionally, this method is likely slower than a non-equivariant U-Net with the same number of weights. To provide a fairer assessment of practical trade-offs, the authors could compare their model against a non-equivariant counterpart with an equivalent inference time instead.", "methods_and_evaluation_criteria": "The method for guaranteeing p4m symmetry is based on an older approach [Cohen & Welling, 2016]. More recent methods ensure equivariance to continuous rotations (e.g., Tensor Field Networks), rather than just discrete ones, which may be more suitable for physics-based applications. Additionally, the advantage of using a staggered grid over the more common approach with CNNs is not clearly justified. \n\nIt should also be more explicitly stated that momentum conservation is enforced in an integral sense over the whole fluid domain, rather than pointwise. Achieving pointwise conservation might lead to even better results.\n\nThe experiments chosen to validate the method cover different physics and boundary conditions and are well explained. However, the fluid domains considered are geometrically simple. Since rotation equivariance is a key aspect of the method, it would be more relevant to test shallow water equations with walls or obstacles of varying geometries (e.g., [Simulating Surface Wave Dynamics with Convolutional Networks, Lino et al., 2020]). This would better assess how well the method generalizes to more complex real-world scenarios.", "theoretical_claims": "The design of the input layers is well justified and clearly explained in Appendix C. Overall, I find the theoretical claims to be well-founded, with the exception of two statements that seem unclear:\n- “Boundary effects interfere with translation equivariance, so we provide a boundary mask input channel.”\n- “Momentum is not conserved due to reflection from closed boundaries.”\nIt would be helpful if the authors could clarify these points.", "experimental_designs_or_analyses": "The selected test problems demonstrate the method’s ability to enforce symmetry and conservation laws, but more complex problems could have been chosen to illustrate the practical relevance of p4m symmetries. For example, testing on fluid systems with irregular geometries or obstacles would better highlight the benefits of rotation equivariance.", "relation_to_broader_scientific_literature": "The authors leverage discrete group convolutions in the hidden layers [Cohen & Welling, 2016] and modify the input layer to properly handle vector magnitudes. However, this is not a major innovation.\n\nIn fluid dynamics, Wang et al. (2020) have already incorporated symmetries and conservation principles using CNNs, and a broader body of literature has investigated rotation equivariance with Graph Neural Networks (GNNs). While the experimental results are strong, a key limitation is the lack of comparison with these prior methods, which would provide a clearer assessment of their contributions.", "essential_references_not_discussed": "This work primarily focuses on p4m equivariance on grids, but a large body of research exists on E(2)/E(3) and SE(2)/SE(3) equivariance in unstructured grids. The following references are relevant but not discussed:\n\t•\tGeneral equivariance methods:\n[1] Thomas et al. (2018) – Tensor Field Networks for rotation and translation-equivariant neural networks on 3D point clouds.\n[2] Gasteiger et al. (2020) – Directional message passing for molecular graphs.\n[3] Brandstetter et al. (2021) – Improving E(3)-equivariant message passing with geometric and physical constraints.\n\t•\tFluid dynamics applications:\n[4] Lino et al. (2022) – Multi-scale rotation-equivariant GNNs for unsteady Eulerian fluid dynamics.\n[5] Toshev et al. (2023) – E(3)-equivariant GNNs for particle-based fluid mechanics.", "other_strengths_and_weaknesses": "Strengths:\n\t•\tThe combination of equivariance and conservation laws leads to significant accuracy improvements.\n\t•\tThe results are strong, demonstrating better stability than many standard neural PDE solvers.\n\t•\tThe supplementary material is detailed.\n\n\nWeaknesses:\n\t•\tLack of comparisons to other equivariant models.\n\t•\tComputational efficiency is not analyzed—equivariant models are usually more expensive, and comparisons at equal inference time would provide a better practical assessment.\n\t•\tSimple geometric test cases. Applying this method to more complex fluid domains (e.g., irregular obstacles) would better showcase its advantages.", "other_comments_or_suggestions": "None.", "questions_for_authors": "1.\tHow does the use of staggered grids improve performance compared to standard CNN approaches?\n2.\tCould momentum conservation be enforced pointwise instead of in an integral sense? Would that improve accuracy?\n3.\tWhat is the computational cost of your method compared to a standard CNN with an equivalent number of parameters?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes novel input and output layers for neural PDE surrogates that enable the enforcement of physical constraints (conservation laws) and geometric constraints (symmetry equivariance) specifically on staggered grids, which are common in CFD. The authors systematically investigate the individual and combined effects of these constraints on two challenging problems: shallow water equations with closed boundaries and decaying incompressible turbulence. They demonstrate that combining both types of constraints improves accuracy, long-term stability, and generalization compared to strong baselines.

The authors' thorough rebuttal effectively addressed the main weaknesses identified by the reviewers. While some limitations remain the core findings on the power of combining constraints on staggered grids represent a solid contribution.
Decision: Accept (poster)