This repository provides the official code for "Equivariant Diffusion for Molecule Generation in 3D," an E(3) equivariant diffusion model designed for generating 3D molecular structures. It targets researchers and practitioners in computational chemistry and drug discovery seeking advanced generative models for molecular design. The primary benefit is the generation of physically plausible 3D molecular conformations.

How It Works

The model employs an E(3) equivariant graph neural network (EGNN) within a diffusion framework. This architecture ensures that predictions are invariant to rotations and translations of the input, crucial for molecular representations. The diffusion process iteratively refines a noisy molecular structure into a coherent 3D conformation, guided by the learned equivariant dynamics.

Quick Start & Requirements

• Install: conda create -c conda-forge -n my-rdkit-env rdkit followed by installing other requirements.
• Prerequisites: Python, CUDA (implied by GPU usage), RDKit (optional but recommended).
• Training: python main_qm9.py --n_epochs 3000 --exp_name edm_qm9 ... (see README for full command).
• Analysis: python eval_analyze.py --model_path outputs/edm_qm9 ...
• Sampling: python eval_sample.py --model_path outputs/edm_qm9 ...
• Docs: README

Highlighted Details

• Supports generation for QM9 and GEOM-Drugs datasets.
• Enables conditional generation based on molecular properties (e.g., alpha, gap, HOMO, LUMO).
• Includes functionality to train separate property classifiers for evaluating generated molecules.
• EGNN's fully connected message passing can be memory-intensive on smaller GPUs.

Maintenance & Community

The project is associated with the work of E. Hoogeboom and collaborators. Further community engagement channels are not explicitly listed in the README.

Licensing & Compatibility

The repository does not explicitly state a license. Users should verify licensing terms for commercial use or integration into closed-source projects.

Limitations & Caveats

The README notes that the GPUs used for experiments were large, suggesting potential memory constraints for users with less powerful hardware due to the EGNN's memory requirements.