E(3) Equivariant Diffusion Model (EDM)

After training a model on QM9, use the following scripts to evaluate quality and visualize results:

Analyze sample quality:

python eval_analyze.py --model_path outputs/edm_qm9 --n_samples 10_000

Visualize molecules:

python eval_sample.py --model_path outputs/edm_qm9 --n_samples 10_000

To work with GEOM-Drugs, first follow the data setup instructions in data/geom/README.md.

Training: Use main_geom_drugs.py. Note that the EGNN model uses fully connected message passing, which is memory intensive; if you encounter OOM errors, try reducing the batch size.

python main_geom_drugs.py --n_epochs 3000 --exp_name edm_geom_drugs --n_stability_samples 500 --diffusion_noise_schedule polynomial_2 --diffusion_steps 1000 --diffusion_noise_precision 1e-5 --diffusion_loss_type l2 --batch_size 64 --nf 256 --n_layers 4 --lr 1e-4 --normalize_factors [1,4,10] --test_epochs 1 --ema_decay 0.9999 --normalization_factor 1 --model egnn_dynamics --visualize_every_batch 10000

Analyze and Sample:

python eval_analyze.py --model_path outputs/edm_geom_drugs --n_samples 10_000
python eval_sample.py --model_path outputs/edm_geom_drugs

To train a model that generates molecules conditioned on specific properties, use main_qm9.py with the --conditioning flag.

Supported properties for --conditioning include:

• alpha
• gap
• homo
• lumo
• mu
• Cv

Example training command:

python main_qm9.py --exp_name exp_cond_alpha  --model egnn_dynamics --lr 1e-4  --nf 192 --n_layers 9 --save_model True --diffusion_steps 1000 --sin_embedding False --n_epochs 3000 --n_stability_samples 500 --diffusion_noise_schedule polynomial_2 --diffusion_noise_precision 1e-5 --dequantization deterministic --include_charges False --diffusion_loss_type l2 --batch_size 64 --normalize_factors [1,8,1] --conditioning alpha --dataset qm9_second_half

To train a classifier to predict properties (e.g., for evaluating the EDM), navigate to the property prediction directory and use main_qm9_prop.py.

Training with EGNN:

cd qm9/property_prediction
python main_qm9_prop.py --num_workers 2 --lr 5e-4 --property alpha --exp_name exp_class_alpha --model_name egnn

Training with Node Count baseline: To train a baseline that only classifies based on the number of nodes, change --model_name to numnodes:

python main_qm9_prop.py --num_workers 2 --lr 5e-4 --property alpha --exp_name exp_class_alpha --model_name numnodes

To train the E(3) Equivariant Diffusion Model on the QM9 dataset, use the main_qm9.py script. You can configure hyperparameters such as epochs, noise schedules, and diffusion steps via CLI flags.

python main_qm9.py --n_epochs 3000 --exp_name edm_qm9 --n_stability_samples 1000 --diffusion_noise_schedule polynomial_2 --diffusion_noise_precision 1e-5 --diffusion_steps 1000 --diffusion_loss_type l2 --batch_size 64 --nf 256 --n_layers 9 --lr 1e-4 --normalize_factors [1,4,10] --test_epochs 20 --ema_decay 0.9999

To evaluate how well a trained property classifier performs on samples generated by a trained EDM model, use eval_conditional_qm9.py with both --generators_path and --classifiers_path.

Example command:

python eval_conditional_qm9.py --generators_path outputs/exp_cond_alpha --classifiers_path qm9/property_prediction/outputs/exp_class_alpha --property alpha  --iterations 100  --batch_size 100 --task edm

Predict molecular properties on the QM9 dataset using main_qm9.py.

Available Properties (--property): alpha, gap, homo, lumo, mu, Cv, G, H, r2, U, U0, zpve.

Recommended Learning Rates (--lr):

• Use 1e-3 for: gap, homo, lumo.
• Use 5e-4 for all other properties (e.g., alpha, mu, Cv, G, H, r2, U, U0, zpve).

Common Arguments:

• --num_workers: Number of worker threads.
• --exp_name: Unique name for the experiment.

# Predict gap (requires 1e-3 lr)
python -u main_qm9.py --num_workers 2 --lr 1e-3 --property gap --exp_name exp_1_gap

# Predict alpha (requires 5e-4 lr)
python -u main_qm9.py --num_workers 2 --lr 5e-4 --property alpha --exp_name exp_1_alpha

While the code can run without RDKit, it is recommended to set up a dedicated Conda environment for RDKit to ensure proper molecular processing. You can create an environment named my-rdkit-env using the following command:

conda create -c conda-forge -n my-rdkit-env rdkit

Use eval_conditional_qm9.py to generate samples based on a trained conditional generator. You must specify the --property and the path to your generators.

Example command:

python eval_conditional_qm9.py --generators_path outputs/exp_cond_alpha --property alpha --n_sweeps 10 --task qualitative

To prepare the GEOM-DRUGS dataset for use with EDM, follow these steps to download, extract, and process the raw data:

1. Download the dataset: Use wget to download the file from Dataverse. Note that the file is approximately 50GB.
2. Extract and Organize: Untar the downloaded file and move the contents into the data/geom/ directory of your project.
3. Install Dependencies: Ensure msgpack is installed in your Python environment.
4. Run Build Script: Execute the build_geom_dataset.py script to process the raw files into the required format.

# 1. Download
wget https://dataverse.harvard.edu/api/access/datafile/4360331

# 2. Untar and move to data/geom/
tar -xzvf 4360331

# 3. Install dependency
pip install msgpack

# 4. Build dataset
python3 build_geom_dataset.py