ProFSA - Self-supervised Pocket Pretraining via Protein Fragment-Surroundings Alignment

Official repository for "Self-supervised Pocket Pretraining via Protein Fragment-Surroundings Alignment".

Figure: An illustration of protein fragment-surroundings alignment framework.

Self-supervised Pocket Pretraining via Protein Fragment-Surroundings Alignment
Bowen Gao, Yinjun Jia, Yuanle Mo, Yuyan Ni, Wei-Ying Ma, Zhi-Ming Ma, Yanyan Lan
Published on The Twelfth International Conference on Learning Representations (ICLR 2024)

Description

Pocket representations play a vital role in various biomedical applications, such as druggability estimation, ligand affinity prediction, and de novo drug design. While existing geometric features and pretrained representations have demonstrated promising results, they usually treat pockets independent of ligands, neglecting the fundamental interactions between them. However, the limited pocket-ligand complex structures available in the PDB database (less than 100 thousand non-redundant pairs) hampers large-scale pretraining endeavors for interaction modeling. To address this constraint, we propose a novel pocket pretraining approach that leverages knowledge from high-resolution atomic protein structures, assisted by highly effective pretrained small molecule representations. By segmenting protein structures into drug-like fragments and their corresponding pockets, we obtain a reasonable simulation of ligand-receptor interactions, resulting in the generation of over 5 million complexes. Subsequently, the pocket encoder is trained in a contrastive manner to align with the representation of pseudo-ligand furnished by some pretrained small molecule encoders. Our method, named ProFSA, achieves state-of-the-art performance across various tasks, including pocket druggability prediction, pocket matching, and ligand binding affinity prediction. Notably, ProFSA surpasses other pretraining methods by a substantial margin. Moreover, our work opens up a new avenue for mitigating the scarcity of protein-ligand complex data through the utilization of high-quality and diverse protein structure databases.

If you find this code useful, please consider star this repo and cite us:

@inproceedings{gao2024self,
  title={Self-supervised Pocket Pretraining via Protein Fragment-Surroundings Alignment},
  author={Gao, Bowen and Jia, Yinjun and Mo, YuanLe and Ni, Yuyan and Ma, Wei-Ying and Ma, Zhi-Ming and Lan, Yanyan},
  booktitle={The Twelfth International Conference on Learning Representations},
  year={2024}
}

Dataset

The original dataset is located at Google Drive.

If you want to train or test the model on the downstream tasks, we also provide processed version of the datasets in lmdb format: profsa.tar.gz.

Environment Setups

1. Clone the repository

git clone https://github.com/bowen-gao/ProFSA.git
cd ProFSA

2. Prepare the dataset and pretrained molecular encoder weights

Download the profsa.tar.gz and decompress it under the data directory.

mkdir data
cd data
tar -zxvf profsa.tar.gz

The pretrained weight for molecular encoder mol_pre_no_h_220816.pt comes from Uni-Mol repo.

After decompress the package, the directory structure should be like this:

.
|-- dataset
|   |-- lba_identity_30
|   |   |-- dict_mol.txt
|   |   |-- dict_pkt.txt
|   |   |-- read_lmdb.py
|   |   |-- test.lmdb
|   |   |-- train.lmdb
|   |   `-- valid.lmdb
|   |-- lba_identity_60
|   |   |-- dict_mol.txt
|   |   |-- dict_pkt.txt
|   |   |-- test.lmdb
|   |   |-- train.lmdb
|   |   `-- valid.lmdb
|   |-- pocket_matching
|   |   |-- P5.pkl
|   |   |-- p5_pockets.lmdb
|   |   `-- test_ligand.lmdb
|   `-- train
|       |-- dict_mol.txt
|       |-- dict_pkt.txt
|       |-- train.lmdb
|       `-- valid.lmdb
|-- log
|   `-- train
|       `-- profsa
|           `-- profsa_release
|               |-- checkpoints
|               |   `-- last.ckpt
|               `-- config.yaml
`-- pretrain
    `-- mol_pre_no_h_220816.pt

3. Build the docker image and launch the container

make init

For the first time, it will prompt you the following configurations (enter for the default values):

Give a project name [profsa]:
Code root to be mounted at /project [.]:
Data root to be mounted at /data [./data]:
Log root to be mounted at /log [./data/log]:
directory to be mounted to xxx [container_home]:
`/home/hongxin/code/profsa/profsa_git/container_home` does not exist in your machine. Create? [yes]:

After Creating xxx ... done, the environment is ready. You can run the following command to go inside the container:

make in

Training

In the container, train a model by running:

zsh scripts/train/train_base.sh

The full command is:

python train.py experiment=base \
    seed=1 \
    trainer.precision="16-mixed" \
    logging.wandb.tags="[profsa]" \
    logging.wandb.name=profsa_fp16_4gpu \
    trainer.devices="[4,5,6,7]"

Note: This project is based on DeepCodebase template. You may need to learn some basic knowledge about Pytorch Lightning and Hydra to better understand the code.

Testing

We have already included our trained model weights in profsa.tar.gz, locating at data/log.

To test our trained model on pocket matching task, you can run:

zsh scripts/test/test_profsa_base_pm.sh

To test our trained model on ligand binding affinity prediction task, you can run:

zsh scripts/test/test_profsa_base_lba.sh

LICENSE

The code is licensed under the MIT license.

Citation

If you find our work useful, please cite our paper:

@inproceedings{gao2023self,
  title={Self-supervised Pocket Pretraining via Protein Fragment-Surroundings Alignment},
  author={Gao, Bowen and Jia, Yinjun and Mo, YuanLe and Ni, Yuyan and Ma, Wei-Ying and Ma, Zhi-Ming and Lan, Yanyan},
  booktitle={The Twelfth International Conference on Learning Representations},
  year={2023}
}