DynamicsPLM: Learning Protein Representations with Conformational Dynamics

Dan Kalifa, Eric Horvitz and Kira Radinsky

The repository is the official implementation of the paper: "Learning Protein Representations with Conformational Dynamics". The paper was accepted for publication in the Bioinformatics journal.

DynamicsPLM - Overview

Some code was borrowed from SaProt repository, which is our base PLM used by DynamicsPLM.

Environment installation

Create a virtual environment

conda create -n DynamicsPLM python=3.10
conda activate DynamicsPLM

Install packages

bash environment.sh  

Prepare the DynamicsPLM model

Encoders checkpoints

The encoder weights are available in SaProt_650M_AF2 model. After downloading, place these files in the appropriate weights/PLMs/SaProt_650M_AF2 path. Then, in the config files, we will upload it using the config_path: weights/PLMs/SaProt_650M_AF2 keyword.

Experimental results

Some experimental results are listed below. For more details, please refer to our paper. For the fine-tuning tasks, the datasets were split based on 30% sequence identity.

Prepare dataset

Downstream tasks

We provide the dynamic datasets that are used in the paper. Datasets can be downloaded from here. Once downloaded, the datasets need to be decompressed and placed in the LMDB folder for supervised fine-tuning.

We provide an option to create the datasets from scratch by downloading the statics datasets from here, and then following the main.py to create the dynamics datasets. We use the conformation generator from RocketSHP.

Fair evaluation

We provide the full benchmark datasets. For each benchmark task, especially new ones, to ensure fair and rigorous evaluation, users must remove any sequence from the pre-training set that exhibits >30% Needleman–Wunsch (NW) similarity to any test sequence before training.

Train & Fine-tune DynamicsPLM

We provide a script for training and fine-tuning DynamicsPLM on the benchmark datasets. The following code shows how to fine-tune DynamicsPLM on specific downstream tasks. Before running the code, please make sure that the datasets are placed in the LMDB folder and the pre-trained weights are placed in the weights/ folder. Note that the default training setting is not the same as in the paper because of the hardware limitations for different users. We recommend that users modify the YAML file flexibly based on their conditions (i.e., batch_size, devices, and accumulate_grad_batches).

# Fine-tune DynamicsPLM on the HumanPPI task
python scripts/training.py -c config/HumanPPI/dynamic_plm.yaml

Test DynamicsPLM on the dynamics subsets of proteins

To test the DynamicsPLM on the dynamic subset of proteins, you can use the clean_clusters_codnasQ.csv file to get the proteins that meet the dynamics criteria. Then, to create a dynamic_test folder and to test the model using the following command:

# Testing the DynamicsPLM on the DeepLoc (Subcellular) task, on the dynamics proteins only
python scripts/training.py -c config/DeepLoc/cls10/dynamic_only.yaml

Record the training process (optional)

If you want to record the training process using wandb, you could modify the config file and set Trainer.logger = True, and then paste your wandb API key in the config key setting.os_environ.WANDB_API_KEY.

Additional Usages

Convert protein structure into structure-aware sequence

We provide a function to convert a protein structure into a structure-aware sequence. The function calls the foldseek binary file to encode the structure. You can download the binary file from here and place it in the bin folder . The following code shows how to use it.

from utils.foldseek_util import get_struc_seq
pdb_path = "example/8ac8.cif"

# Extract the "A" chain from the pdb file and encode it into a struc_seq
# pLDDT is used to mask low-confidence regions if "plddt_mask" is True. Please set it to True when
# use AF2 structures for best performance.
parsed_seqs = get_struc_seq("bin/foldseek", pdb_path, ["A"], plddt_mask=False)["A"]
seq, foldseek_seq, combined_seq = parsed_seqs

print(f"seq: {seq}")
print(f"foldseek_seq: {foldseek_seq}")
print(f"combined_seq: {combined_seq}")

Reproduce pre-trained embeddings or changing the base encoder

To train the SaProt model (the base encoders), you need the pre-training dataset, which can be downloaded from here.

Conformation generator

To train the conformation generator, follow the instruction here. This model use ESM-3 model for the structure predictions, therefore you should gain access to this model following instructions:

1. Get a token for access the model using the link: https://huggingface.co/settings/tokens, and the ESM-3 online web.
2. Then, in your env enter: huggingface-cli login

Bioemu Generator

For using other generators, such as Biomolecular Emulator (BioEmu), follow the instruction here. Here, you will find the required pip installation and the code required for generating conformations for each protein. Using our 'main.py' file you could switch the RocketSHP predictions with once from BioEmu.

Citation

If you find this repository useful, please cite our paper.