AdsorbDiff: Adsorbate Placement via Conditional Denoising Diffusion [ICML 2024]

This is the official code repository for the paper AdsorbDiff: Adsorbate Placement via Conditional Denoising Diffusion, accepted at International Conference on Machine Learning, 2024.

If you have any questions, concerns, or feature requests, please feel free to email me.

Installation

conda create -n adsorbdiff python=3.10
conda activate adsorbdiff
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.2.0+cu118.html

git clone https://github.com/AdeeshKolluru/AdsorbDiff.git
cd AdsorbDiff
pip install -r requirements.txt
pip install -e .

Datasets

Creating lmdbs for training a diffusion model -

• Download the ASE trajectories of OC20-Dense dataset (link) as well as mappings.
• Download OC20-Dense subset of data in lmdb format with 244 adsorbate surface combinations which is utilized for Open Catalyst Challenge 2023 - Link
• Find unique ids (sids) by iterating across the lmdb above and store it using this script.
• We can then create the train and val ID lmdbs using this script.
• For generating lmdbs for conditional training, this script can be used.

Or you can directly download the lmdbs for conditional diffusion training here.

Creating lmdbs for sampling -

• Val OOD ASE trajectories - download.
• For generating val OOD lmdbs from val OOD trajs, this script can be used.

Val OOD subset lmdb - download

• To generate lmdbs for AdsorbML baselines, this script can be used.

Val ID lmdb - download

Pretraining on OC20 -

• OC20 IS2RE lmdbs - download.

Checkpoints

Diffusion

PaiNN - download - recommended due to faster inference

EquiformerV2 - download

MLFF optimization

All pre-trained OCP checkpoints can be downloaded here for MLFF optimization.

Training, sampling and relaxation

The bash script to generate commands for training, sampling and relaxations is in run.py. Different commands in this file can be used for different cases.

Citation

Please consider citing our paper if you find it helpful. Thank you!

@misc{kolluru2024adsorbdiff,
      title={AdsorbDiff: Adsorbate Placement via Conditional Denoising Diffusion}, 
      author={Adeesh Kolluru and John R Kitchin},
      year={2024},
      eprint={2405.03962},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

Acknowledgements

This codebase was built on top of ocp, Open-Catalyst-Dataset repositories as well as adapting code from DiffDock and AdsorbML. Please cite these works as well!