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

TL;DR

This paper introduces a flexible SE(2) equivariant graph neural network architecture for PDE surrogates on non-gridded domains, achieving improved data efficiency and accuracy in fluid flow simulations.

Contribution

It proposes a novel SE(2) equivariant GNN that aligns representations with the principal axis, enabling better PDE surrogate modeling on irregular domains.

Findings

Significant improvements in data efficiency over non-equivariant models

Enhanced accuracy in fluid flow simulation benchmarks

Effective handling of non-gridded domain data

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.