SineNet

SineNet: Learning Temporal Dynamics in Time-Dependent Partial Differential Equations

Code for our ICLR 2024 paper on deep-learning-based fluid dynamics simulation. Our code is based on the PDEArena library [1].

Incompressible Navier-Stokes	Compressible Navier-Stokes	Shallow Water Equations

Model

We employ a multi-stage UNet model. Check out our paper for details.

Data

For INS and SWE, please download from PDEArena here. The SWE data are converted from .nc to .h5 using the h5_conv.py script.

The CNS data were generated using PDEBench [2] here. This data can be generated using their code with the modified files and data generation script PDEBench_gen.sh in the PDEbench folder of this repo as:

bash data_gen.sh --mode train --nsamples 5600 --batch_size 50 --run && bash data_gen.sh --mode valid --nsamples 1400 --batch_size 50 --run && bash data_gen.sh --mode test --nsamples 1400 --batch_size 50 --run

Note that the solver rarely but consistently exhibits instability resulting in trajectories of all 0 which are removed following data generation, which is why the split we presented in the paper is 5400/1300/1300.

Setup

cd pdearena
source setup.sh

This will create a new conda environment named pdearena, install this code base locally, and install other necessary packages.

Training

python scripts/train.py -c configs/<config>.yaml \
            --data.data_dir=<data_dir> \
            --data.num_workers=8 \
            --data.batch_size=32 \
            --model.name=<model_name> \
            --model.lr=2e-4 --optimizer.lr=2e-4

where <config> can be navierstokes2d, cfd or shallowwater2d_2day. Valid <model_name>'s can be found in pdearena/models/registry.py, e.g., sinenet8-dual.

Testing

python scripts/test.py test -c configs/<config>.yaml \
            --data.data_dir=<data_dir> \
            --trainer.devices=1 \
            --data.num_workers=8 \
            --data.batch_size=32 \
            --model.name=<model_name> \
            --ckpt_path=<ckpt_path>

Conditional Training

For training on the conditional Navier-Stokes data from [1], use the following command:

python scripts/cond_train.py -c configs/cond_navierstokes2d.yaml 
            --data.data_dir=<data dir> 
            --trainer.devices=1 
            --data.num_workers=8 
            --data.valid_limit_trajectories=5 
            --data.batch_size=32 
            --model.name=<model name> \
            --model.lr=2e-4 --optimizer.lr=2e-4

Valid <model_name>'s can be found in COND_MODEL_REGISTRY from pdearena/models/registry.py, e.g., sinenet8-adagn.

[1] Gupta, Jayesh K., and Johannes Brandstetter. "Towards multi-spatiotemporal-scale generalized pde modeling." arXiv preprint arXiv:2209.15616 (2022).

[2] Takamoto, Makoto, et al. "PDEBench: An extensive benchmark for scientific machine learning." Advances in Neural Information Processing Systems 35 (2022): 1596-1611.

Citation

@inproceedings{zhang2024sinenet,
    title={SineNet: Learning Temporal Dynamics in Time-Dependent Partial Differential Equations},
    author={Xuan Zhang and Jacob Helwig and Yuchao Lin and Yaochen Xie and Cong Fu and Stephan Wojtowytsch and Shuiwang Ji},
    booktitle={The Twelfth International Conference on Learning Representations},
    year={2024},
    url={https://openreview.net/forum?id=LSYhE2hLWG}
}

Acknowledgements

This work was supported in part by National Science Foundation grants IIS-2243850 and IIS-2006861.