atomicarchitects / symphony

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 [email protected]: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}
}

symphony's Issues

Structure codebase

Follow https://github.com/google/flax/tree/main/examples/ogbg_molpcba

Settings for analysis with xyz2mol

Dear authors,

first of all, congrats for your insightful work and its publication at ICLR 2024. I really enjoyed reading the paper. The thorough analysis of your model and the baselines is very interesting and well executed. Thank you!

I wanted to ask which version and settings of xyz2mol you use for your analysis. You report that 99.99% of the molecules in QM9 are validated. I recently implemented a script that uses the implementation of xyz2mol in RDKit (rdkit.Chem.rdDetermineBonds) for our schnetpack-gschnet package and it only finds SMILES for about 95% (~124k) of the molecules.
Did you use the version from RDKit or the script? And which settings/arguments did you use?

Thanks for your help and best regards,
Niklas