atomicarchitects/symphony

[ICLR'24] Symphony: Symmetry-Equivariant Point-Centered Spherical Harmonics for Molecule Generation

Symphony: Symmetry-Equivariant Point-Centered Spherical Harmonics for Molecule Generation

This is the official code-release for the paper Symphony: Symmetry-Equivariant Point-Centered Spherical Harmonics for Molecule Generation, published at ICLR 2024.

Instructions

Clone the repository:

git clone git@github.com:atomicarchitects/symphony.git
cd symphony
git checkout iclr_2024_final

Since this repository is actively being developed, we recommend using the iclr_2024_final branch for the most stable version of the code.

Default Setup

Create and activate a virtual environment:

python -m venv .venv && source .venv/bin/activate

Install pip dependencies with:

pip install --upgrade pip && pip install -r requirements.txt

For GPU support, install JAX with CUDA support afterwards:

pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Some of the analysis scripts require openbabel==3.1.1. This can be installed through conda.

Checking Installation

Check that installation suceeded by running a short test:

python -m tests.train_test

Start a Training Run

Start training with a configuration defined under configs/:

python -m symphony \
    --config configs/qm9/e3schnet_and_nequip.py \
    --workdir ./workdirs

The --workdir flag specifies the directory where the model checkpoints, logs, and other artifacts will be saved.

Changing Hyperparameters

Since the configuration is defined using config_flags, you can override hyperparameters. For example, to change the number of training steps, and the batch size:

python -m symphony --config configs/qm9/e3schnet_and_nequip.py \
    --workdir ./workdirs \
    --config.num_train_steps=10 --config.max_n_graphs=16

For more extensive changes, directly edit the configuration files, or add your own.

Citation

Please cite our paper if you use this code!

@inproceedings{
    daigavane2024symphony,
    title={Symphony: Symmetry-Equivariant Point-Centered Spherical Harmonics for Molecule Generation},
    author={Ameya Daigavane and Song Eun Kim and Mario Geiger and Tess Smidt},
    booktitle={The Twelfth International Conference on Learning Representations},
    year={2024},
    url={https://openreview.net/forum?id=MIEnYtlGyv}
}