Space-Time Continuous PDE Forecasting using Equivariant Neural Fields

TL;DR

This work tackles data-driven forecasting of continuous PDE dynamics from irregular, sparse observations across diverse geometries. It introduces a space-time continuous framework that preserves known PDE symmetries by using Equivariant Neural Fields (ENFs) and a group-equivariant neural ODE to evolve latent representations, with meta-learning speeding up initialization of latent states. The key contributions are extending ENFs to operate under symmetries beyond SE(n), enforcing equivariance in both latent dynamics and decoding, and demonstrating improved generalization, data efficiency, and robustness across plane, torus, sphere, and ball geometries. The approach provides grid-agnostic, symmetry-informed PDE forecasting with strong performance advantages over prior NeF-based methods, particularly in non-grid and long-horizon settings relevant to geophysical and fluid-dynamics applications.

Abstract

Recently, Conditional Neural Fields (NeFs) have emerged as a powerful modelling paradigm for PDEs, by learning solutions as flows in the latent space of the Conditional NeF. Although benefiting from favourable properties of NeFs such as grid-agnosticity and space-time-continuous dynamics modelling, this approach limits the ability to impose known constraints of the PDE on the solutions -- e.g. symmetries or boundary conditions -- in favour of modelling flexibility. Instead, we propose a space-time continuous NeF-based solving framework that - by preserving geometric information in the latent space - respects known symmetries of the PDE. We show that modelling solutions as flows of pointclouds over the group of interest $G$ improves generalization and data-efficiency. We validated that our framework readily generalizes to unseen spatial and temporal locations, as well as geometric transformations of the initial conditions - where other NeF-based PDE forecasting methods fail - and improve over baselines in a number of challenging geometries.

Figures (9)

• Figure 1: We propose to solve an equivariant PDE in function space by solving an equivariant ODE in latent space. Through our proposed framework, which leverages Equivariant Neural Fields$f_\theta$, a field $\nu_t$ is represented by a set of latents $z^\nu_t = \{(p_{i}^\nu,\mathbf{c}_{i}^\nu)\}_{i=1}^N$ consisting of a pose$p_{i}$ and context vector $\mathbf{c}_{i}$. Using meta-learning, the initial latent $z^\nu_0$ is fit in only 3 SGD steps, after which an equivariant neural ODE $F_\psi$ models the solution as a latent flow.
• Figure 1: MSE $\downarrow$ for heat equation on $\mathbb{R}^2$.
• Figure 2: The proposed framework respects pre-defined symmetries of the PDE: a rotated solution $\mathcal{L}_g\nu_T$ may be obtained either by solving from latent $z^\nu_0$ (top-left) and transforming the solution $z_T^\nu$ (top-right) to $gz_T^\nu$ (bottom-right) or transforming $z^\nu_0$ to $gz^\nu_0$ (bottom-left) and solving this.
• Figure 3: We show the impact of meta-learning and equivariance on the latent space of the ENF when representing trajectories of PDE states. Fig. \ref{['fig:tsne-nonmeta']} shows a T-SNE plot of the latent space of $f_\theta$ when $z^\nu_t$ is optimized with autodecoding, and no weight sharing over bi-invariants is enforced. Fig. \ref{['fig:tsne-meta-nonequiv']} shows the latent space when meta-learning is used, but no weight sharing is enforced. Fig. \ref{['fig:tsne-meta-equiv']} shows the latent space when $z^\nu_t$ are obtained using meta-learning and $f_\theta$ shares weights over $\mathbf{a}^{\rm SE(n)}$.
• Figure 4: A train and test sample from the planar diffusion dataset. Initial conditions for train and test are spikes in disjoint subsets of $\mathbb{R}^2$.