Equivariant Spatio-Temporal Attentive Graph Networks for Physical Dynamics Simulation

Initialize Environment

It is suggested to use conda to manage the python environment, and you can install the required packages from the provided environment.yml directly.

conda env create --file environment.yml

Data Preparation

We provide three datasets in Google Drive

1. MD17

The MD17 dataset can be downloaded from MD17.

2. Motion Capture

The raw data were obtained via CMU Motion Capture Database. The preprocessed dataset as well as the splits are provided in motion folder.

3. Protein MD

We provide the data preprocessing code in mdanalysis/preprocess.py. One can simply run

python mdanalysis/preprocess.py

Training and Evaluation

1. MD17

python main_md.py --exp_name='exp_1' --model='estag' --mol='aspirin' --n_layers=2 --fft=True --eat=True --with_mask

2. Protein MD

python main_mdanalysis.py --exp_name='exp_2' --model='estag' --n_layers=2 --fft=True --eat=True --with_mask

3. Motion Capture

python main_motion.py --exp_name='exp_3' --model='estag' --n_layers=2 --fft=True --eat=True --with_mask

Rollout

bash rollout/rollout.sh

Visualization

Here we demonstrate with MD17 as an instance, and the same procedure can be employed for Protein and Motion.

Predict the states (coordinates) of the next frame according to a selected trajectory
```
python rollout/md17_pred.py
```
Based on the predicted coordinates, launch the file visualization/vis_md.ipynb to display the molecule.

Citation

If our work is helpful for your research, you can cite it as follows:

@article{limi_2024_estag,
  title={Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics},
  author={Wu, Liming and Hou, Zhichao and Yuan, Jirui and Rong, Yu and Huang, Wenbing},
  journal={Advances in Neural Information Processing Systems},
  volume={36},
  year={2024}
}