PyTorch Implementation for Proteus: Exploring Protein Structure Generation for Enhanced Designability and Efficiency.
Proteus is a novel deep diffusion network designed to generate protein backbones with enhanced designability and efficiency. Unlike RFDiffusion which relies on large pre-trained network RosettaFold for structure prediction, Proteus utilizes graph-based triangle methods and a multi-track interaction network, achieving state-of-the-art performance without the need for pre-training. Notably, the inference speed has been accelerated from 4x up to 10x compared to FrameDiff and RFdiffusion. Our model's capabilities have been validated through comprehensive in silico evaluations and experimental characterizations, demonstrating its potential to significantly advance the field of protein design.
![image](https://private-user-images.githubusercontent.com/59241275/335772105-9cd5d387-66c9-4f71-9fa8-6a27cd77a25b.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjM3OTg4MzAsIm5iZiI6MTcyMzc5ODUzMCwicGF0aCI6Ii81OTI0MTI3NS8zMzU3NzIxMDUtOWNkNWQzODctNjZjOS00ZjcxLTlmYTgtNmEyN2NkNzdhMjViLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA4MTYlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwODE2VDA4NTUzMFomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPTdlNjYwMzkzYWY5YTdjMjc4NTZlMWVmZDMyZTVjYjQ5NjkxMGU2N2I0Y2I1NDcyYjlmODE1MmEzYmNiNGZhZTcmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.KXGjN9IKBoS4aMzkemwNKK7hykEzlqqluf4Qn2HSfoc)
We recommend miniconda (or anaconda). Run the following to install a conda environment with the necessary dependencies. Using mamba if possible for better install speed.
# install
conda env create -f se3.yml
# optional : using mamba for faster environment installation
conda install mamba
mamba env create -f se3.yml
# activate environment
conda activate Proteus
# install this repo as a local package
pip install -e .
The checkpoint is avaiable at ./weights/paper_weights.pt
For the first time run, it might be a little slow because of downloading esmfold ckpt
weight_path=./weights/paper_weights.pt
python ./experiments/inference_se3_diffusion.py \
inference.output_dir=inference_outputs/monomer/ \
inference.weights_path=$weight_path \
inference.diffusion.samples.samples_lengths=[100,200,300,400,600,800] \
inference.diffusion.samples.samples_per_length=100 \
inference.diffusion.num_t=100
# config below is optional
# To disable esmfold prediction and mpnn design, add extra config
inference.mpnn.enable=False inference.esmfold.enable=False
# To disable esmfold prediction add extra config
inference.esmfold.enable=False
A self_consistency.csv will be generated in the inference_outputs/monomer/${timestap}/self_consistency.csv
, report all necessary metrics like dssp or sc-rmsd, etc.
baseline_weight_path=./weights/paper_weights.pt
python ./experiments/inference_se3_diffusion.py \
inference.output_dir=inference_outputs/oligomer/ \
inference.weights_path=$baseline_weight_path \
inference.diffusion.samples.contigs='60-80//60-80' \
inference.diffusion.samples.samples_per_length=100 \
inference.diffusion.num_t=100
Inference output wuold be like
inference_outputs
└── 12D_02M_2023Y_20h_46m_13s # Date time of inference.
├── mpnn.fasta # mpnn designed seuences.
├── self_consistency.csv # self consistency analysis, contains rmsd and tmscore between scaffold ans esmfold, mpnn score of sequence, scaffold path, esmf path etc.
├── diffusion # dir contains scaffold generated by proteus
│ ├── 100_1_sample.pdb
│ ├── 100_2_sample.pdb # {length}_{sample_id}_sample.pdb
| └── ...
├── trajctory # dir contains traj pdb, exists when inference.diffusion.option.save_trajactory=True
│ ├── 100_1_bb_traj.pdb
│ ├── 100_2_bb_traj.pdb # {length}_{sample_id}_traj.pdb
| └── ...
├── movie # dir contains full atom protein designed by mpnn, exists when inference.diffusion.option.plot.switch_on=True
│ ├── 100_1_rigid_movie.gif # movie of protein rigid at time t
│ ├── 100_1_rigid_0_movie.gif # movie of predict protein rigid at time 0 from time t
| └── ...
├── mpnn # dir exists when pyrosetta in installed and inference.mpnn.dump=True
│ ├── 100_0_sample_mpnn_0.pdb
│ ├── 100_0_sample_mpnn_1.pdb # {length}_{sample_id}_sample_mpnn_{sequence_id}.pdb
| └── ...
└── esmf # dir contians esmf predict strcture
├── 100_0_sample_esmf_0.pdb
├── 100_0_sample_esmf_0.pdb # {length}_{sample_id}_sample_esmf_{sequence_id}.pdb
└── ...
The local triangle attention is implemented below:
Line 245 in 71f9eaf
LICENSE: MIT
If you use our work then please cite
@article{wang2024proteus,
title={Proteus: exploring protein structure generation for enhanced designability and efficiency},
author={Wang, Chentong and Qu, Yannan and Peng, Zhangzhi and Wang, Yukai and Zhu, Hongli and Chen, Dachuan and Cao, Longxing},
journal={bioRxiv},
pages={2024--02},
year={2024},
publisher={Cold Spring Harbor Laboratory}
}
Proteus is built upon the following codebases, please give them a star if you enjoy Proteus :)