Skip to content

Learning to design protein-protein interactions with enhanced generalization (ICLR 2024)

License

Notifications You must be signed in to change notification settings

anton-bushuiev/PPIformer

Repository files navigation

PPIformer

arXiv badge Zenodo badge License: MIT Open in Spaces Open in Spaces

PPIformer is a state-of-the-art predictor of the effects of mutations on protein-protein interactions (PPIs), as quantified by the binding energy changes (ddG). The model was pre-trained on the PPIRef dataset via a coarse-grained structural masked modeling and fine-tuned on SKEMPI v2.0 via log odds. PPIformer was shown to successfully identify known favorable mutations of the staphylokinase thrombolytic and a human antibody against the SARS-CoV-2 spike protein. Please see more details in our paper.

Please do not hesitate to contact us or create an issue/PR if you have any questions or suggestions. ✌️

Web on Hugging Face Spaces 🤗

The preview of PPIformer is available via an interactive user interface on Hugging Face Spaces:

Model on HF

The Hugging Face Space above is running using Zero GPU, which is currently in beta. If you experience any issues, please try the CPU-only version below:

Model on HF

Installation

Step 1. To install PPIformer locally, clone this repository and install the environment (you may need to adjust the versions of the PyTorch-based packages in the script depending on your system):

conda create -n ppiformer python==3.10 -y
conda activate ppiformer
git clone https://github.com/anton-bushuiev/PPIformer && pip install -e PPIformer

Step 2. (Optional) After installation, you may need to adapt PyTorch to your system. Please see the official PyTorch installation guide for details. For example, if you are using AMD GPUs, you may need to install PyTorch for ROCm:

pip install -U torch --index-url https://download.pytorch.org/whl/rocm6.0

Step 3. (Optional) If you are planning to re-train the model or reproduce test results (see Training and testing below), please clone and install the ppiref and mutils packages locally to download the necessary data (i.e., data split files and ddG-labeled datasets):

git clone https://github.com/anton-bushuiev/PPIRef && pip install -e PPIRef
git clone https://github.com/anton-bushuiev/mutils && pip install -e mutils

Inference

import torch
from ppiformer.tasks.node import DDGPPIformer
from ppiformer.utils.api import download_from_zenodo, predict_ddg, embed
from ppiformer.definitions import PPIFORMER_WEIGHTS_DIR, PPIFORMER_TEST_DATA_DIR

# Download the weights
download_from_zenodo('weights.zip')

Predict ddG for a PPI upon mutation

PPIformer was fine-tuned on the SKEMPI v2.0 dataset via log odds. The fine-tuned models can be used to predict the binding energy changes (ddG) for a PPI upon mutation.

# Load the ensamble of fine-tuned models
device = 'cuda' if torch.cuda.is_available() else 'cpu'
models = [DDGPPIformer.load_from_checkpoint(PPIFORMER_WEIGHTS_DIR / f'ddg_regression/{i}.ckpt', map_location=torch.device('cpu')).eval() for i in range(3)]
models = [model.to(device) for model in models]

# Specify input
ppi_path = PPIFORMER_TEST_DATA_DIR / '1bui_A_C.pdb'  # PDB or PPIRef file (see https://ppiref.readthedocs.io/en/latest/extracting_ppis.html)
muts = ['SC16A', 'FC47A', 'SC16A,FC47A']  # List of single- or multi-point mutations

# Predict
ddg = predict_ddg(models, ppi_path, muts)
ddg
> tensor([-0.3708,  1.5188,  1.1482])

Embed PPI

PPIformer was pre-trained using structural masked modeling. The pre-trained model can be used to obtain PPI embeddings, similar to BERT embeddings in natural language processing.

# Load the pre-trained model
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = PPIformer.load_from_checkpoint(PPIFORMER_WEIGHTS_DIR / 'masked_modeling.ckpt', map_location=torch.device('cpu'))
model = model.to(device).eval()

# Specify input
ppi_path = PPIFORMER_TEST_DATA_DIR / '1bui_A_C.pdb'  # PDB or PPIRef file (see https://ppiref.readthedocs.io/en/latest/extracting_ppis.html)

# Embed (get the final type-0 features). Here, 128-dimensional embedding for each of 124 amino acids in the PPI
embedding = embed(model, ppi_path)
embedding.shape
> torch.Size([124, 128])

Training and testing

To train and validate PPIformer, please see PPIformer/scripts/README.md. To test the model and reproduce the results from the paper, please see PPIformer/notebooks/test.ipynb.

How it works

The model was pre-trained on the PPIRef dataset via a coarse-grained structural masked modeling and fine-tuned on the SKEMPI v2.0 dataset via log odds.

A single pre-training step starts with randomly sampling a protein-protein interaction (in this example figure, the staphylokinase dimer A-B from the PDB entry 1C78) from PPIRef. Next, randomly selected residues are masked to obtain the masked interaction . After that, the interaction is converted into a graph representation with masked nodes (black circles). The model subsequently learns to classify the types of masked amino acids by acquiring -invariant hidden representation of the whole interface via the encoder and classifier (red arrows). On the downstream task of ddG prediction, mutated amino acids are masked, and the probabilities of possible substitutions are jointly inferred with the pre-trained model. Finally, the estimate is obtained using the predicted probabilities of the wild-type and the mutant amino acids via log odds (blue arrows).

TODO

  • Pre-training and fine-tuning examples with scripts/run.py
  • Installation script examples for AMD GPUs and NVIDIA GPUs
  • SSL-pretrained weights (without fine-tuning)

References

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

@article{
  bushuiev2024learning,
  title={Learning to design protein-protein interactions with enhanced generalization},
  author={Anton Bushuiev and Roman Bushuiev and Petr Kouba and Anatolii Filkin and Marketa Gabrielova and Michal Gabriel and Jiri Sedlar and Tomas Pluskal and Jiri Damborsky and Stanislav Mazurenko and Josef Sivic},
  booktitle={The Twelfth International Conference on Learning Representations},
  year={2024}
}