tsm-pde's Introduction

TSM-PDE

See "A Neural PDE Solver with Temporal Stencil Modeling" for the paper associated with this codebase.

Parts of the codebase is adapted from google/jax-cfd.

Setup

conda env create -f environment.yml
conda activate cfd
cd jax-cfd
pip install jaxlib
pip install -e ".[complete]"
cd ..

Data

Both the training and evaluation data can be deterministically generated. Please see the data_generation for more details.

Reproduction

Please check the reproducing_scripts for more details.

Pretrained Checkpoints

Please download the pretrained model checkpoints from here.

Reference

If you found this codebase useful, please consider citing the following papers:

Temporal Stencil Modeling:

@article{sun2023tsm,
  title={A Neural PDE Solver with Temporal Stencil Modeling},
  author={Sun, Zhiqing and Yang, Yiming and Yoo, Shinjae},
  journal={arXiv preprint arXiv:2302.08105},
  year={2023}
}

Learned Interpolation:

@article{Kochkov2021-ML-CFD,
  author = {Kochkov, Dmitrii and Smith, Jamie A. and Alieva, Ayya and Wang, Qing and Brenner, Michael P. and Hoyer, Stephan},
  title = {Machine learning{\textendash}accelerated computational fluid dynamics},
  volume = {118},
  number = {21},
  elocation-id = {e2101784118},
  year = {2021},
  doi = {10.1073/pnas.2101784118},
  publisher = {National Academy of Sciences},
  issn = {0027-8424},
  URL = {https://www.pnas.org/content/118/21/e2101784118},
  eprint = {https://www.pnas.org/content/118/21/e2101784118.full.pdf},
  journal = {Proceedings of the National Academy of Sciences}
}

tsm-pde's People

Contributors

Stargazers

Watchers

tsm-pde's Issues

Question on the data generation of 2D Kolmogorov flow

Hey Edward @Edward-Sun nice paper. I have two amateur questions on the data generation for the examples.

Why there is a $-0.1\mathbf{u}$ in the body force? 2D Kolmogorov flow should only have the first term.
Did you change any of the stencils on Jax cfd staggered grid? For example, how to compute the time derivative, or make the convection term unwinding, things like that.

Thanks,
Shuhao

Recommend Projects