Equivariant Neural Simulators for Stochastic Spatiotemporal Dynamics

This is the code implementation of the paper Equivariant Neural Simulators for Stochastic Spatiotemporal Dynamics.

Requirements

To install requirements:

pip install -r requirements.txt

Additionally, CC3D version 4.3.0 is use to generate the cellular dynamics data. Please see this link for instructions.

Training

To train the EPNS or PNS models, run this command, where config is a config file in the configs directory:

python train_model.py <config>

For example, for training EPNS on celestial dynamics, run:

python train_model.py n_body_dynamics_EPNS

In addition, you can easily overwrite parameters in the configuration files from the command line. For example, if you want to train EPNS with one-step training on the cellular dynamics data, simply run:

python train_model.py cell_dynamics_EPNS --training_strategy=one-step

Please consult the config files for the relevant parameters.

Evaluation

To evaluate EPNS on the celestial dynamics data, run:

python evaluation/run_experiments_n_body.py

For cellular dynamics, run:

python evaluation/run_experiments_cell.py

Generating data

The data used for training and experiments can be downloaded here. https://drive.google.com/drive/folders/1n2RTTitoUqVhw8NIiVn_hvz8dz9u4OlW?usp=sharing

Code for generating the datasets will be added to this repository soon.

---

If you found our work useful, consider citing:

@inproceedings{
minartz2023,
title={Equivariant Neural Simulators for Stochastic Spatiotemporal Dynamics},
author={Koen Minartz and Yoeri Poels and Simon Martinus Koop and Vlado Menkovski},
booktitle={Thirty-seventh Conference on Neural Information Processing Systems},
year={2023},
url={https://openreview.net/forum?id=CCVsGbhFdj}
}