0、推断时，修改diffdock_config.yaml文件。输入可以是csv文件（protein_ligand_csv参数），第一列是pdb文件绝对路径，第二列是分子（smiles，sdf，mol2均可）；也可以输入protein_path参数和ligand参数。

1、获取蛋白质pdb文件，记得在pymol中去除水和其他杂质分子、去离子。注意���项目用BioPython读取PDB文件，会自动将pdb中occupancy不确定的原子选取最高的那个进行保存
2、为了避免因为PDB网站的pdb文件有各种问题，比如侧链不完整、原子occupancy不高、residue缺失等导致最后的结果有影响，其实不如直接按照UniprotID获取alphafold2的预测结构来用

3、分子可以提供smiles字符串、也可以提供sdf文件或mol2文件。注意sdf文件需要提前补H。直接提供smiles字符串是比较方便的

4、目前提取esm2特征的步骤已集成在inference代码中自动进行，对应的esm代码也一起放在了项目里
注意：在对接前需要预计算蛋白质的ESM2特征，注意参数--esm_embeddings_path。每个AA的特征都要保存，结果为N*M，N为AA的数目,M为ESM2特征维度。ESM2的使用参考README。
python scripts/extract.py esm2_t33_650M_UR50D pdbbind_sequences.fasta embeddings_output --repr_layers 33 --include per_tok

5、该步骤已经集成在inference中自动完成
根据PDB文件可以获取对应的氨基酸序列，可以使用biopython中的PPBuilder来获取。参考get_seq_from_pdb.py文件

6、Score_model的model_parameters如果将remove_hs修改为False就导致全部对接失败。暂时不清楚原因。但该参数会影响lig graph的建立

7. 对接完成后获得ligand的sdf文件和蛋白质的pdb文件，然后可以使用get_AAs_around_ligand.py来提取ligand周围（radius指定范围）的AAs
   

8. score_model.py是score模型的主函数，其关键在于GFN网络的应用。代码注释里我给出了两个GFN网络的小例子，要跑一下了解了GFN再看其他的
   

9. sampling.py是逆向SDE或ODE的代码，照着文献看，还是蛮容易理解的

Implementation of DiffDock, state-of-the-art method for molecular docking, by Gabriele Corso*, Hannes Stark*, Bowen Jing*, Regina Barzilay and Tommi Jaakkola. This repository contains all code, instructions and model weights necessary to run the method or to retrain a model. If you have any question, feel free to open an issue or reach out to us: gcorso@mit.edu, hstark@mit.edu, bjing@mit.edu.

The repository also contains all the scripts to run the baselines and generate the figures. Additionally, there are visualization videos in visualizations.

You might also be interested in this Google Colab notebook to run DiffDock by Brian Naughton.

The files in data contain the names for the time-based data split.

If you want to train one of our models with the data then:

1. download it from zenodo
2. unzip the directory and place it into data such that you have the path data/PDBBind_processed

We will set up the environment using Anaconda. Clone the current repo

git clone https://github.com/gcorso/DiffDock.git

This is an example for how to set up a working conda environment to run the code (but make sure to use the correct pytorch, pytorch-geometric, cuda versions or cpu only versions):

conda create --name diffdock python=3.9
conda activate diffdock
conda install pytorch==1.11.0 pytorch-cuda=11.7 -c pytorch -c nvidia
pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric==2.0.4 -f https://data.pyg.org/whl/torch-1.11.0+cu117.html
python -m pip install PyYAML scipy "networkx[default]" biopython rdkit-pypi e3nn spyrmsd pandas biopandas

Then you need to install ESM that we use both for protein sequence embeddings and for the protein structure prediction in case you only have the sequence of your target. Note that OpenFold (and so ESMFold) requires a GPU. If you don't have a GPU, you can still use DiffDock with existing protein structures.

pip install "fair-esm[esmfold]"
pip install 'dllogger @ git+https://github.com/NVIDIA/dllogger.git'
pip install 'openfold @ git+https://github.com/aqlaboratory/openfold.git@4b41059694619831a7db195b7e0988fc4ff3a307'

We support multiple input formats depending on whether you only want to make predictions for a single complex or for many at once.
The protein inputs need to be .pdb files or sequences that will be folded with ESMFold. The ligand input can either be a SMILES string or a filetype that RDKit can read like .sdf or .mol2.

For a single complex: specify the protein with --protein_path protein.pdb or --protein_sequence GIQSYCTPPYSVLQDPPQPVV and the ligand with --ligand ligand.sdf or --ligand "COc(cc1)ccc1C#N"

For many complexes: create a csv file with paths to proteins and ligand files or SMILES. It contains as columns complex_name (name used to save predictions, can be left empty), protein_path (path to .pdb file, if empty uses sequence), ligand_description (SMILE or file path) and protein_sequence (to fold with ESMFold in case the protein_path is empty). An example .csv is at data/protein_ligand_example_csv.csv and you would use it with --protein_ligand_csv protein_ligand_example_csv.csv.

And you are ready to run inference:

python -m inference --protein_ligand_csv data/protein_ligand_example_csv.csv --out_dir results/user_predictions_small --inference_steps 20 --samples_per_complex 40 --batch_size 10 --actual_steps 18 --no_final_step_noise

When providing the .pdb files you can run DiffDock also on CPU, however, if possible, we recommend using a GPU as the model runs significantly faster. Note that the first time you run DiffDock on a device the program will precompute and store in cache look-up tables for SO(2) and SO(3) distributions (typically takes a couple of minutes), this won't be repeated in following runs.

Download the data and place it as described in the "Dataset" section above.

First run:

python datasets/pdbbind_lm_embedding_preparation.py

Use the generated file data/pdbbind_sequences.fasta to generate the ESM2 language model embeddings using the library https://github.com/facebookresearch/esm by installing their repository and executing the following in their repository:

python scripts/extract.py esm2_t33_650M_UR50D pdbbind_sequences.fasta embeddings_output --repr_layers 33 --include per_tok --truncation_seq_length 4096

This generates the embeddings_output directory which you have to copy into the data folder of our repository to have data/embeddings_output. Then run the command:

python datasets/esm_embeddings_to_pt.py

We first generate the language model embeddings for the testset, then run inference with DiffDock, and then evaluate the files that DiffDock produced:

python datasets/esm_embedding_preparation.py --protein_ligand_csv data/testset_csv.csv --out_file data/prepared_for_esm_testset.fasta
git clone https://github.com/facebookresearch/esm 
cd esm
pip install -e .
cd ..
HOME=esm/model_weights python esm/scripts/extract.py esm2_t33_650M_UR50D data/prepared_for_esm_testset.fasta data/esm2_output --repr_layers 33 --include per_tok
python -m inference --protein_ligand_csv data/testset_csv.csv --out_dir results/user_predictions_testset --inference_steps 20 --samples_per_complex 40 --batch_size 10 --actual_steps 18 --no_final_step_noise
python evaluate_files.py --results_path results/user_predictions_testset --file_to_exclude rank1.sdf --num_predictions 40

Train the large score model:

python -m train --run_name big_score_model --test_sigma_intervals --esm_embeddings_path data/esm2_3billion_embeddings.pt --log_dir workdir --lr 1e-3 --tr_sigma_min 0.1 --tr_sigma_max 19 --rot_sigma_min 0.03 --rot_sigma_max 1.55 --batch_size 16 --ns 48 --nv 10 --num_conv_layers 6 --dynamic_max_cross --scheduler plateau --scale_by_sigma --dropout 0.1 --remove_hs --c_alpha_max_neighbors 24 --receptor_radius 15 --num_dataloader_workers 1 --cudnn_benchmark --val_inference_freq 5 --num_inference_complexes 500 --use_ema --distance_embed_dim 64 --cross_distance_embed_dim 64 --sigma_embed_dim 64 --scheduler_patience 30 --n_epochs 850

The model weights are saved in the workdir directory.

Train a small score model with higher maximum translation sigma that will be used to generate the samples for training the confidence model:

python -m train --run_name small_score_model --test_sigma_intervals --esm_embeddings_path data/esm2_3billion_embeddings.pt --log_dir workdir --lr 1e-3 --tr_sigma_min 0.1 --tr_sigma_max 34 --rot_sigma_min 0.03 --rot_sigma_max 1.55 --batch_size 16 --ns 24 --nv 6 --num_conv_layers 5 --dynamic_max_cross --scheduler plateau --scale_by_sigma --dropout 0.1 --remove_hs --c_alpha_max_neighbors 24 --receptor_radius 15 --num_dataloader_workers 1 --cudnn_benchmark --val_inference_freq 5 --num_inference_complexes 500 --use_ema --scheduler_patience 30 --n_epochs 300

In practice, you could also likely achieve the same or better results by using the first score model for creating the samples to train the confidence model, but this is what we did in the paper.
The score model used to generate the samples to train the confidence model does not have to be the same as the score model that is used with that confidence model during inference.

Train the confidence model by running the following:

python -m confidence.confidence_train --original_model_dir workdir/small_score_model --run_name confidence_model --inference_steps 20 --samples_per_complex 7 --batch_size 16 --n_epochs 100 --lr 3e-4 --scheduler_patience 50 --ns 24 --nv 6 --num_conv_layers 5 --dynamic_max_cross --scale_by_sigma --dropout 0.1 --all_atoms --remove_hs --c_alpha_max_neighbors 24 --receptor_radius 15 --esm_embeddings_path data/esm2_3billion_embeddings.pt --main_metric loss --main_metric_goal min --best_model_save_frequency 5 --rmsd_classification_cutoff 2 --cache_creation_id 1 --cache_ids_to_combine 1 2 3 4

first with --cache_creation_id 1 then --cache_creation_id 2 etc. up to 4

Now everything is trained and you can run inference with:

python -m evaluate --model_dir workdir/big_score_model --ckpt best_ema_inference_epoch_model.pt --confidence_ckpt best_model_epoch75.pt --confidence_model_dir workdir/confidence_model --run_name DiffDockInference --inference_steps 20 --split_path data/splits/timesplit_test --samples_per_complex 40 --batch_size 10 --actual_steps 18 --no_final_step_noise

Note: the notebook data/apo_alignment.ipynb contains the code used to align the ESMFold-generated apo-structures to the holo-structures.

@article{corso2023diffdock,
      title={DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking}, 
      author = {Corso, Gabriele and Stärk, Hannes and Jing, Bowen and Barzilay, Regina and Jaakkola, Tommi},
      journal={International Conference on Learning Representations (ICLR)},
      year={2023}
}

MIT

We thank Wei Lu and Rachel Wu for pointing out some issues with the code.