Geometric and Physical Constraints Synergistically Enhance Neural PDE Surrogates

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.

Partial differential equation (PDE) solvers use scientific knowledge to predict how physical systems evolve: they are accurate but often slow. Neural networks can learn to predict the same results more quickly, but can be inaccurate or unstable when predicting too far into the future. To get the best of both worlds, these networks can be cleverly constructed to mimic some aspects the physical systems they predict.We focus on two such aspects. First, we impose the conservation of physical quantities such as mass or momentum, which for certain PDEs should never be created or destroyed. Second, we force the network to respect the geometric symmetries of the PDE: for example, rotating or otherwise transforming the initial conditions of should transform the predicted solution. Both of these aspects have been use to improve neural PDE predictions before, but we carefully evaluate their combination for the first time. To do this, we created new network components to for symmetries on special grids used in fluid dynamics, with staggered spacing between variables.Overall, we were able to confirm the usefulness of combined physical and symmetry constraints. Our results are promising for ML-based prediction of weather, climate and other fluid dynamical systems.