The data used for this code can either be generated directly or is already included in the code/packages installed. For QM9, the data is already included by default and is loaded automatically when running the script (after the environment has been setup). Meanwhile, for N-body, run the following script after installing the environment:

python n_body/dataset/generate_dataset.py --initial_vel 1 --num-train 3000 --length 1000 --length_test 1000 --sufix "small"

Both a .txt and .yaml file have been provided for convenience.

1. Create a conda environment and install the packages:

conda env create -f environment.yaml

Note: Be sure to use the enyironment file that corresponds to your device.

2. Activate the environment:

conda activate egnn-transformer

Example notebooks with the runs for our experiments can be found in /notebooks/ To perform the experiments, simply run the following lines:

1. For QM9:

python main_qm9.py
python transformer_trainer.py

You can adjust the process using the following parameters:

2. For N-Body:

python nbody_egnn_trainer.py
python nbody_transformer_trainer.py

You can adjust the process using the following parameters:

Alternatively, the experiments can be run here.

Some of the code has been adapted from the following sources:

N-body: https://github.com/vgsatorras/egnn

The above can be cited through the following:

@article{Satorras2021EnEG,
  title={E(n) Equivariant Graph Neural Networks},
  author={Victor Garcia Satorras and Emiel Hoogeboom and Max Welling},
  journal={ArXiv},
  year={2021},
  volume={abs/2102.09844},
  url={https://api.semanticscholar.org/CorpusID:231979049}
}

Transformer: https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/JAX/tutorial6/Transformers_and_MHAttention.html