Solving the Discretised Shallow Water Equations Using Neural Networks

AI technology is now ubiquitous, with specialised hardware and software playing crucial roles in the development and deployment of artificial intelligence models. Popular frameworks like TensorFlow and PyTorch run seamlessly on multiple types of hardware and Graphics Processing Units (GPUs) and Tensor Processing Units (TPUs) are extensively employed to solve computational fluid dynamics problems.

In a recent paper co-authored by Boyang Chen, Amin Nadimy, and colleagues, we demonstrate how convolutional layers commonly used in neural networks can help solve complex partial differential equations. The paper, titled “Solving the Discretised Shallow Water Equations Using Neural Networks“, showcases this innovative approach. The NN4PDEs framework, developed by the Applied Modelling and Computational Group at Imperial College London, enables these use cases. In this method, the convolutional layer provides a discretisation scheme to solve the equations. Though conceptually simple, the paper illustrates the practical implementation for complex equations like the shallow water equations.

Importantly, these tools can effectively model flood risk with the accuracy and spatial resolution needed to represent local water depth during significant floods. We applied this to simulate the flood event impacting Carlisle, UK, in January 2005. The figure below, from the paper, depicts the problem setup and the solution provided by this approach.

While this early work demonstrates that the NN4PDEs approach achieves results and computational performances comparable to traditional computational fluid dynamics tools, it also paves the way for further applications. These applications can seamlessly combine AI models and physics on different scales or allow scaling up hardware to solve complex fluid dynamics problems over large domains.