This implementation is based on DiGress's excellent work.

The code has been tested with:

• PyTorch 2.2
• CUDA 11.8
• PyTorch Geometric 2.3.1

conda create -n THGD python=3.10 -y
pip install torch==2.2 --index-url https://download.pytorch.org/whl/cu121
pip install -r requirements.txt
pip install torch-scatter==2.1.2
pip install -e .

All executions are launched via python3 main.py. Refer to Hydra documentation for parameter customization.

The model training consists of two sequential phases:

1. Coarse Model Training (learns high-level molecular topologies)
2. Refinement Model Training (recovers atomic-level details)

Pre-configured training files are provided in configs/. Example for GuacaMol:

1. Train coarse model:

cd src
python main.py +guacamol_exp=coarse

2. Train refinement model:

python main.py +guacamol_exp=refine

Notes:

• Dataset preprocessing (feature extraction for coarse/expanded graphs) runs on CPU and typically takes ~30 minutes (device-dependent). It is recommended to use our preprocessed dataset here:
  • https://drive.google.com/file/d/1Rr5EOMLNOdOYCj6L5VzQaS95NooPDbyC/view?usp=sharing
• This preprocessing must complete before training begins.

Sampling requires:

1. Trained coarse model checkpoint
2. Trained refinement model checkpoint
3. Precomputed optimal prior distribution tensor (obtain by preprocessing)

Use pre-configured sampling profiles:

python sample.py sample=guacamol

❗ Configuration Alignment:
Ensure the checkpoint's training config matches your sampling config (e.g., coarse_cfg in YAML must correspond to the checkpoint's original training config).

Modify the scaffold field in the sampling config.

Available checkpoint for three datasets (place downloaded files and place in checkpoints/ folder):

https://drive.google.com/file/d/1nOrOq6Jf7adqG0vbitnbFPJ9_-AzdHcX/view?usp=sharing

We provide few example outputs in sample-results/.