ByProt: 功能强大的蛋白质设计工具包

ByProt简介

ByProt是一个专为蛋白质研究领域设计的多功能工具包。它主要聚焦于基于结构的序列设计(又称固定骨架设计),为研究人员提供了强大而灵活的工具。

LM-Design示意图

ByProt的主要特点

高效的非自回归ProteinMPNN变体: ByProt提供了ProteinMPNN的一个高效且有效的非自回归变体。ProteinMPNN是一个强大的蛋白质固定骨架序列设计工具。
LM-Design的官方实现: ByProt作为LM-Design模型的官方实现,这是一个最先进的蛋白质序列设计模型。该模型源自论文《Structure-informed Language Models Are Protein Designers》,并在ICML 2023会议上进行了口头报告。
持续扩展的功能: ByProt团队正在不断扩展其功能,以涵盖更广泛的任务和特性。这使得ByProt能够为蛋白质研究提供更加全面的工具包。

安装ByProt

要安装ByProt,请按照以下步骤操作:

克隆项目仓库:

git clone --recursive https://url/to/this/repo/ByProt.git
cd ByProt

创建conda虚拟环境:

env_name=ByProt
conda create -n ${env_name} python=3.7 pip
conda activate ${env_name}

自动安装所有依赖:

bash install.sh

基于结构的蛋白质序列设计(反向折叠)

ByProt提供了多个预训练模型权重,可用于不同的蛋白质设计任务。这些模型权重可以从Zenodo下载。

数据准备

为了使用ByProt进行蛋白质设计,需要准备适当的数据集。ByProt支持以下数据集:

CATH数据集:
- CATH 4.2数据集由论文《Generative Models for Graph-Based Protein Design》提供。
- CATH 4.3数据集由论文《Learning inverse folding from millions of predicted structures》提供。
可以使用以下命令下载CATH数据集:
```
bash scripts/download_cath.sh
```
PDB复合物数据(多链): 这个数据集包含了从蛋白质数据库(PDB)中精选的蛋白质(多链)复合物。可以使用以下命令下载:
```
bash scripts/download_multichain.sh
```

下载完数据后,需要在相应的配置文件中设置正确的数据路径。

模型训练

ByProt支持多种模型的训练。以下是两个训练示例:

示例1: 非自回归(NAR) ProteinMPNN基线

使用条件掩码语言建模(CMLM)训练NAR ProteinMPNN:

export CUDA_VISIBLE_DEVICES=0
exp=fixedbb/protein_mpnn_cmlm  
dataset=cath_4.2
name=fixedbb/${dataset}/protein_mpnn_cmlm

python ./train.py \
    experiment=${exp} datamodule=${dataset} name=${name} \
    logger=tensorboard trainer=ddp_fp16

示例2: LM-Design

在ESM-1b 650M基础上训练LM-Design:

exp=fixedbb/lm_design_esm1b_650m
dataset=cath_4.2
name=fixedbb/${dataset}/lm_design_esm1b_650m

./train.py \
    experiment=${exp} datamodule=${dataset} name=${name} \
    logger=tensorboard trainer=ddp_fp16

评估和推理

训练完成后,可以在验证集或测试集上进行评估和推理:

dataset=cath_4.2
name=fixedbb/${dataset}/lm_design_esm1b_650m
exp_path=/root/research/projects/ByProt/run/logs/${name}

python ./test.py \
    experiment_path=${exp_path} \
    data_split=test ckpt_path=best.ckpt mode=predict \
    task.generator.max_iter=5

在Notebook中使用训练好的模型设计序列

ByProt提供了方便的接口,允许用户在Jupyter Notebook中使用训练好的模型进行序列设计。以下是两个示例:

示例1: 使用ProteinMPNN-CMLM

from byprot.utils.config import compose_config as Cfg
from byprot.tasks.fixedbb.designer import Designer

# 1. 初始化设计器
exp_path = "/root/research/projects/ByProt/run/logs/fixedbb/cath_4.2/protein_mpnn_cmlm"
cfg = Cfg(
    cuda=True,
    generator=Cfg(
        max_iter=1,
        strategy='mask_predict',
        temperature=0,
        eval_sc=False,  
    )
)
designer = Designer(experiment_path=exp_path, cfg=cfg)

# 2. 从PDB文件加载结构
pdb_path = "/root/research/projects/ByProt/data/3uat_variants/3uat_GK.pdb"
designer.set_structure(pdb_path)

# 3. 从给定结构生成序列
designer.generate()

# 4. 计算评估指标
designer.calculate_metrics()

示例2: 使用LM-Design

from byprot.utils.config import compose_config as Cfg
from byprot.tasks.fixedbb.designer import Designer

# 1. 初始化设计器
exp_path = "/root/research/projects/ByProt/run/logs/fixedbb/cath_4.2/lm_design_esm2_650m"
cfg = Cfg(
    cuda=True,
    generator=Cfg(
        max_iter=5,
        strategy='denoise', 
        temperature=0,
        eval_sc=False,  
    )
)
designer = Designer(experiment_path=exp_path, cfg=cfg)

# 2. 从PDB文件加载结构
pdb_path = "/root/research/projects/ByProt/data/3uat_variants/3uat_GK.pdb"
designer.set_structure(pdb_path)

# 3. 生成序列
designer.generate()

# 4. 计算评估指标
designer.calculate_metrics()

示例3: 局部序列设计(Inpainting)

ByProt还支持对蛋白质的特定区段进行局部序列设计,这在设计抗体CDR等应用中非常有用:

pdb_path = "/root/research/projects/ByProt/data/pdb_samples/5izu_proc.pdb"
designer.set_structure(pdb_path)

start_ids = [1, 50]
end_ids = [10, 100]

for i in range(5):
    out, ori_seg, designed_seg = designer.inpaint(
        start_ids=start_ids, end_ids=end_ids, 
        generator_args={'temperature': 1.0}
    )
    print(designed_seg)
print('Original Segments:')
print(ori_seg)

致谢

ByProt项目得益于多个开源项目和个人的贡献:

PyTorch Lightning和lightning-hydra-template为开发过程提供了强大的基础。
jingraham/neurips19-graph-protein-design提供了预处理的CATH数据集和数据管道实现。
facebook/esm提供了ESM实现、预训练模型权重和数据管道组件。
dauparas/ProteinMPNN提供了ProteinMPNN实现和多链数据集。
A4Bio/PiFold提供了PiFold实现。
jasonkyuyim/se3_diffusion提供了自一致性结构评估实现。

ByProt团队对这些仓库的作者表示诚挚的感谢,他们的贡献对ByProt的开发至关重要。

总结

ByProt作为一个功能强大的蛋白质设计工具包,为研究人员提供了丰富的资源和灵活的接口。无论是进行基础研究还是应用开发,ByProt都能满足各种蛋白质设计需求。随着功能的不断扩展,ByProt有望成为蛋白质研究领域的重要工具之一。

如果您在研究中使用了ByProt,请引用以下论文:

@inproceedings{zheng2023lm_design,
    title={Structure-informed Language Models Are Protein Designers},
    author={Zheng, Zaixiang and Deng, Yifan and Xue, Dongyu and Zhou, Yi and YE, Fei and Gu, Quanquan},
    booktitle={International Conference on Machine Learning},
    year={2023}
}

通过使用ByProt,研究人员可以更高效地进行蛋白质序列设计,推动蛋白质工程和生物技术领域的发展。随着人工智能和机器学习技术在生物学领域的不断应用,像ByProt这样的工具将在未来的蛋白质研究中发挥越来越重要的作用。