This is the official Pytorch implementation of the paper: SEGNO: Generalizing Equivariant Graph Neural Networks with Physical Inductive Biases in ICLR 2024. [OpenReview] [Paper] [Poster]
- Python 3.8.12
- For the other packages, please refer to the
environment.yml. You can just execute following command to create the conda environment.
conda env create -f environment.yml
Use the script generate_dataset.py in nbody/dataset to generate the n-body simulation data used in experiments.
To generate the Charged dataset:
python N-body/nbody/dataset/generate_dataset.py --simulation=charged --num-train 10000 --seed 43 --suffix small --save_path N-body/nbody/dataset
To generate the Gravity dataset:
python N-body/nbody/dataset/generate_dataset.py --simulation=gravity --num-train 10000 --seed 43 --suffix small --save_path N-body/nbody/dataset
We also provide the download link of the generated dataset: N-body dataset.
We use the --target option to change the predicted interval, where short/medium/long denote the 1000/1500/2000 ts.
To predict the charged N-body system after 1000ts:
python N-body/main.py --epochs=500 --max_samples=3000 --layers=8 --hidden_features=64 --norm=none --batch_size=100 --gpu=1 --weight_decay=1e-12 --target short --dataset_dir dataset --dataset charged
To predict the charged N-body system after 1500ts:
python N-body/main.py --epochs=500 --max_samples=3000 --layers=8 --hidden_features=64 --norm=none --batch_size=100 --gpu=1 --weight_decay=1e-12 --target medium --dataset_dir dataset_long
To predict the gravitational N-body system:
python N-body/main.py --epochs=500 --max_samples=3000 --layers=8 --hidden_features=64 --norm=none --batch_size=100 --gpu=1 --weight_decay=1e-12 --target short --dataset_dir datset_gravity
The MD22 dataset can be downloaded from MD22. The splits are provided in the MD22 folder.
python -u MD22/spatial_graph/main_md22.py --data_dir MD22/spatial_graph/md22 --n_layers 5 --lr 1e-3 --mol <mol>
The value of <mol> can be:
nhme, dha, atat, stachyose, atatcgcg, buckyball, nanotube.
The raw data were obtained via CMU Motion Capture Database. The preprocessed dataset as well as the splits are provided in spatial_graph/motion folder.
python -u spatial_graph/main_motion.py --config_by_file
Please kindly cite our paper if you find this paper and the codes helpful. :)
@inproceedings{
yang2024improving,
title={Improving Generalization in Equivariant Graph Neural Networks with Physical Inductive Biases},
author={Yang Liu and Jiashun Cheng and Haihong Zhao and Tingyang Xu and Peilin Zhao and Fugee Tsung and Jia Li and Yu Rong},
booktitle={International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=Hkx1qkrKPr](https://openreview.net/forum?id=3oTPsORaDH}
}
We use the code from the following repository: