Implementation of COSMIC by Maksim Kuznetsov*, Fedor Ryabov*, Roman Schutski*, Rim Shayakhmetov, Yen-Chu Lin, Alex Aliper, and Daniil Polykovskiy.
To build Docker image and install requirements, run:
docker build -t {IMAGE_NAME} ./docker/
docker run -v ./:/COSMIC -shm-size 16G --memory=16G -it --name {CONTAINER_NAME} {IMAGE_NAME}
Also, download the GEOM dataset using the following link and preprocess it with the following command
GEOMDS=/path/to/geomdataset
mkdir $GEOMDS/rdkit_folder/drugs_sdf/
python ./COSMIC/scripts/preprocess_dataset.py --root=$GEOMDS/rdkit_folder --summary_file=summary_drugs.json --new_summary_file=summary_preprocessed.json
mkdir $pathtogeomdataset/qm9_sdf/
python ./COSMIC/scripts/preprocess_dataset.py --root=$GEOMDS/rdkit_folder --summary_file=summary_qm9.json --new_summary_file=summary_qm9_preprocessed.json
To train the models in the paper, run this commands:
GEOM-Drugs dataset experiments
python COSMIC/scripts/train.py --root=$GEOMDS/rdkit_folder --summary_path=$GEOMDS/rdkit_folder/summary_preprocessed.json --num_workers=8 --num_epochs=6 --gpu=... --conditions=none --model=aae_wgan --batch_size=256
python COSMIC/scripts/train.py --root=$GEOMDS/rdkit_folder --summary_path=$GEOMDS/rdkit_folder/summary_preprocessed.json --num_workers=8 --num_epochs=6 --gpu=... --conditions=descriptors3d --model=aae_wgan --batch_size=256
GEOM-QM9 dataset experiments
python COSMIC/scripts/train.py --root=$GEOMDS/rdkit_folder --summary_path=$GEOMDS/rdkit_folder/summary_qm9_preprocessed.json --num_workers=8 --num_epochs=60 --gpu=... --conditions=none --model=aae_wgan --batch_size=256 --dataset=qm9
To track training metrics, you can run Tensorboard:
tensorboard --logdir=logs/ --bind_all
To sample conformations and compute metrics, please, use notebooks in notebooks folder.
@article{kuznetsov2024cosmic,
author = {Kuznetsov, Maksim and Ryabov, Fedor and Schutski, Roman and Shayakhmetov, Rim and Lin, Yen-Chu and Aliper, Alex and Polykovskiy, Daniil},
title = {COSMIC: Molecular Conformation Space Modeling in Internal Coordinates with an Adversarial Framework},
journal = {Journal of Chemical Information and Modeling},
doi = {10.1021/acs.jcim.3c00989},
note ={PMID: 38668753},
URL = {https://doi.org/10.1021/acs.jcim.3c00989},
eprint = {https://doi.org/10.1021/acs.jcim.3c00989}
}