Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Flexible SE(2) graph neural networks with applications to PDE surrogates

Maria Bånkestad, Olof Mogren, Aleksis Pirinen

TL;DR

This work introduces a $SE(2)$-equivariant graph neural network for surrogate modeling of PDEs on irregular 2D domains, notably applied to Navier–Stokes fluid flows. By aligning node representations to a principal axis, the model reduces the symmetry to $SO(1)$, enabling flexible nonlinear processing while preserving equivariance under rotations and translations. The method demonstrates superior data efficiency and accuracy on a 2D Tetris equivariance benchmark and on Navier–Stokes surrogates with varied buoyancy and obstacles, outperforming non-equivariant baselines in both one-step and rollout errors. This approach enables accurate PDE surrogates on non-gridded domains and offers practical utility for fast fluid simulations with irregular geometries, with code released for reproducibility.

Abstract

This paper presents a novel approach for constructing graph neural networks equivariant to 2D rotations and translations and leveraging them as PDE surrogates on non-gridded domains. We show that aligning the representations with the principal axis allows us to sidestep many constraints while preserving SE(2) equivariance. By applying our model as a surrogate for fluid flow simulations and conducting thorough benchmarks against non-equivariant models, we demonstrate significant gains in terms of both data efficiency and accuracy.

Flexible SE(2) graph neural networks with applications to PDE surrogates

TL;DR

This work introduces a -equivariant graph neural network for surrogate modeling of PDEs on irregular 2D domains, notably applied to Navier–Stokes fluid flows. By aligning node representations to a principal axis, the model reduces the symmetry to , enabling flexible nonlinear processing while preserving equivariance under rotations and translations. The method demonstrates superior data efficiency and accuracy on a 2D Tetris equivariance benchmark and on Navier–Stokes surrogates with varied buoyancy and obstacles, outperforming non-equivariant baselines in both one-step and rollout errors. This approach enables accurate PDE surrogates on non-gridded domains and offers practical utility for fast fluid simulations with irregular geometries, with code released for reproducibility.

Abstract

This paper presents a novel approach for constructing graph neural networks equivariant to 2D rotations and translations and leveraging them as PDE surrogates on non-gridded domains. We show that aligning the representations with the principal axis allows us to sidestep many constraints while preserving SE(2) equivariance. By applying our model as a surrogate for fluid flow simulations and conducting thorough benchmarks against non-equivariant models, we demonstrate significant gains in terms of both data efficiency and accuracy.
Paper Structure (26 sections, 1 theorem, 37 equations, 20 figures, 3 tables)

This paper contains 26 sections, 1 theorem, 37 equations, 20 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Equivariance
  4. Graph neural networks
  5. Navier-Stokes equations
  6. PDE surrogate models
  7. SE(2) graph neural network
  8. Experiments
  9. Equivariance evaluation using 2D Tetris classification
  10. Navier-Stokes simulation
  11. Navier-Stokes simulation with obstacle
  12. Related work
  13. Conclusions
  14. Method: Additional details
  15. An equivariance proof
  16. ...and 11 more sections

Key Result

Proposition 3.1

If ${\bm{R}}_{\alpha_i}$ is a rotation to align ${\bm{x}}_i, {\bm{r}}_i$ to the $x$-axis, where $\alpha_i$ is the rotation angle, and ${\bm{R}}_{-\alpha_i}$ is the rotation matrix inverse. Then, for a function where $g$ is any nonlinear function, it holds that Thus, $f$ is equivariant under the SO(2) group.

Figures (20)

  • Figure 1: A snapshot from two simulations (with an upwards facing force) of smoke flowing around an obstacle, where the location of the smoke inflow differs.
  • Figure 2: The message-passing procedure. We rotate the node features by the angles $\theta_i$ so that they all align with the $x$-coordinate axis. The scalar features $\tilde{{\bm{x}}}_i$ and the rotation feature $\hat{{\bm{x}}}_i$ can now be concatenated and inputted to any arbitrary message-passing function $f^m$. The messages are then rotated back and aggregated by arbitrary aggregation function $f^a$.
  • Figure 3: An outline of our model. The top green blocks denote the input embedding of the data, followed by the main block that comprises a message-passing and a feed-forward layer. Finally, the output layer aligns with what we aim to predict.
  • Figure 4: The seven shapes in the Tetris training dataset. The tests will of these shapes rotated at random angles $\theta$ (see Figure \ref{['fig:tetris_test_blocks']} in the appendix).
  • Figure 5: Our SE(2) model obtains perfect accuracy despite being trained only a single copy of each Tetris shape, whereas the regular message-passing network obtains similar results only after being exposed to eight copies of each shape.
  • ...and 15 more figures

Theorems & Definitions (3)

  • Proposition 3.1
  • proof
  • proof