Space-Time Continuous PDE Forecasting
using Equivariant Neural Fields
this paper on ArXiv 
code for experiments Summary
Building on recent work that uses Conditional Neural Fields as downstream representations, we propose a framework for equivariant spatio-temporally continuous PDE solving to leverage the symmetries that often occur in physical dynamical systems. Using an Equivariant Neural Field (ENFs) backbone $f_\theta$, we fit a latent $z^{\nu}_0$ to the initial state $\nu_0$ of a PDE, and subsequently solve the PDE in latent space using an equivariant graph-based neural ODE $F_\psi$. We show that this approach is able to solve PDEs with high accuracy and robustness to sparse initial conditions, and that it can readily be applied to PDEs defined over a variety of geometries, with different symmetries.
.png)
Fig. 1) Method overview. We propose to solve an equivariant PDE in function space by solving an equivariant ODE in latent space. Through our proposed framework, which leverages ENFs $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 (left), after which an equivariant neural ODE $F_\psi$ models the solution as a latent flow (right).
Main findings
Equivariance and Meta-Learning structure the ENF latent space
We apply T-SNE on the latents $z^\nu_t$ for $t=1, ..., t=T=10$ fit to timesteps from a planar heat diffusion equation.
Results show that the latent space of an ENF trained without meta-learning or equivariance constraints (a) is less structured than the latent space of an ENF trained with meta-learning (b) which in turn is less structured than an ENF trained with both meta-learning and equivariance constraints (c). Further on, we show that this increased structure improves the ability of the Neural ODE $F_\psi$ to model the latent flow.
Improved stability to sparse observations of initial conditions
We fit our framework to Navier-Stokes equations in 2D with periodic boundary conditions. We show that our model achieves improved performance over baselines when initial conditions are sparsely observed.
Solving PDEs on challenging geometries
We showcase the ability of our model to be applied in challenging geometries, improving performance in zero-shot super-resolution on the shallow-water equations, and allowing for solving PDEs over geometries where previous approaches fail, such as internally heated convection in the ball.
Fig. 2) Shallow-water zero-shot super-resolution. We showcase the ability of our model to natively handle zero-shot super-resolution tasks, training the model on a set of low-resolution solutions (top) to the shallow-water equations on the sphere. The model is then able to predict solutions (bottom) for for high-resolution unseen initial conditions (middle), without any additional training. We showcase performance for a version of our framework that respects the longitudinal symmetries of the shallow water equations resulting from the Coriolis force.
.png)
Fig. 3) Internally heated convection in the ball. We show that the model is able to solve PDEs over geometries where previous approaches fail, such as internally heated convection in the ball with 3D rotational symmetries. This experiment also nicely showcases the strenght of using Neural Fields to solve PDEs, making it almost trivial to adapt the model to different geometries.
Learn more
If you would like to learn more about Equivariant Neural Fields, check out the original paper on ArXiv. Looking to use Equivariant Neural Fields in your own projects? We wrote a minimally bloated, hopefully easy-to-understand, maximally extensible implementation of ENFs in JAX (PyTorch coming soon), with accompanying explainer notebook.
Any other questions, comments or criticisms? Feel free to reach out to us over email (d.m.knigge@uva.nl, d.r.wessels@uva.nl) or on Twitter (@davidmknigge, @dafidofff) we're always down to chat about equivariance, neural fields, or anything else that piques your interest.
Citing
If you found this work useful, and should you like to reference it, please use the following citation.
@article{knigge2024space,
title={Space-Time Continuous PDE Forecasting using Equivariant Neural Fields},
author={Knigge, David M and Wessels, David R and Valperga, Riccardo and
Papa, Samuele and Sonke, Jan-Jakob and
Gavves, Efstratios and Bekkers, Erik J},
journal={arXiv preprint arXiv:2406.06660},
year={2024}
}