This example uses recurrent neural networks for spatio-temporal prediction of the Navier Stokes flow. The example uses architecture that is inspired from Physics Informed RNN-DCT Networks for Time-Dependent Partial Differential Equations paper.
Time-series prediction is a key task in many domains. The application of deep learning architectures—particularly RNNs, long short-term memory networks (LSTMs), and similar networks has significantly enhanced the predictive capabilities. These models are unique in their ability to capture temporal dependencies and learn complex patterns over time, making them well suited for forecasting time varying relationships. In physics-ML, these models are critical in predicting dynamic physical systems’ evolution, enabling better simulations, understanding of complex natural phenomena, and aiding in discoveries.
This problem involves predicting the next timesteps of a 2D Navier Stokes flow given input of the initial condition or multiple input timesteps.
This example relies on the Dataset used in Fourier Neural Operator paper which solves the Navier Stokes equations in the vorticity form on a unit torus. The different samples are generated by using different initial conditions for the simulation The example uses 1000 training samples and 10 test samples.
The model uses Convolutional GRU layers for the RNN propagation and use a ResNet type architecture for spatial encoding. The model has two variants- one in which the subsequent predictions can be generated using only a single input time step (one-to-many) and other where multiple input time-steps are used to produce multiple output time-steps (many-to-many/seq-to-seq).
The example script contains code to download and do any pre-processing
for the dataset. To download the dataset, gdown package is required,
which can be installed using
pip install gdown
To get started, simply run
python navier_stokes_rnn.py
To run the seq2seq variant of the example, simply run
python navier_stokes_rnn.py model_type=seq2seq