使用ESM模型进行蛋白质接触预测教程

概述

本教程将详细介绍如何使用ESM(Evolutionary Scale Modeling)模型进行蛋白质接触预测。ESM是由Facebook Research开发的蛋白质语言模型，能够从蛋白质序列中提取丰富的结构信息。接触预测是蛋白质结构预测中的关键步骤，能够帮助理解蛋白质的三维结构。

环境准备

依赖安装

首先需要安装必要的Python依赖包：

pip install biopython biotite
pip install fair-esm

对于Colab用户，建议安装aria2以加速模型权重下载：

apt-get install aria2

数据下载

需要下载示例MSA(多序列比对)文件，这些文件包含蛋白质序列的进化信息：

# 示例MSA文件
1a3a_1_A.a3m
5ahw_1_A.a3m 
1xcr_1_A.a3m

模型权重下载

ESM会自动下载模型权重，但手动下载速度更快：

mkdir -p /root/.cache/torch/hub/checkpoints
aria2c --dir=/root/.cache/torch/hub/checkpoints --continue --split 8 \
    https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t33_650M_UR50S.pt
aria2c --dir=/root/.cache/torch/hub/checkpoints --continue --split 8 \
    https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t33_650M_UR50S-contact-regression.pt

核心功能实现

MSA文件处理

def read_sequence(filename: str) -> Tuple[str, str]:
    """读取fasta或MSA文件中的第一个(参考)序列"""
    record = next(SeqIO.parse(filename, "fasta"))
    return record.description, str(record.seq)

def remove_insertions(sequence: str) -> str:
    """移除序列中的插入字符"""
    translation = str.maketrans(dict.fromkeys(string.ascii_lowercase + ".*"))
    return sequence.translate(translation)

def read_msa(filename: str) -> List[Tuple[str, str]]:
    """读取MSA文件中的所有序列，自动移除插入字符"""
    return [(record.description, remove_insertions(str(record.seq))) 
            for record in SeqIO.parse(filename, "fasta")]

结构到接触矩阵转换

def contacts_from_pdb(structure: bs.AtomArray, 
                    distance_threshold: float = 8.0,
                    chain: Optional[str] = None) -> np.ndarray:
    """从PDB结构中计算接触矩阵"""
    # 计算Cβ原子坐标
    N = structure.coord[mask & (structure.atom_name == "N")]
    CA = structure.coord[mask & (structure.atom_name == "CA")] 
    C = structure.coord[mask & (structure.atom_name == "C")]
    Cbeta = extend(C, N, CA, 1.522, 1.927, -2.143)
    
    # 计算距离矩阵并转换为接触矩阵
    dist = squareform(pdist(Cbeta))
    contacts = (dist < distance_threshold).astype(np.int64)
    contacts[np.isnan(dist)] = -1
    return contacts

MSA子采样

def greedy_select(msa: List[Tuple[str, str]], 
                 num_seqs: int, 
                 mode: str = "max") -> List[Tuple[str, str]]:
    """贪心算法选择MSA中的代表性序列"""
    array = np.array([list(seq) for _, seq in msa], dtype=np.bytes_).view(np.uint8)
    indices = [0]
    for _ in range(num_seqs - 1):
        dist = cdist(array[indices[-1:]], array, "hamming")
        new_index = optfunc(np.mean(pairwise_distances, axis=0) + dist)
        indices.append(new_index)
    return [msa[idx] for idx in indices]

实际应用

1. 加载MSA数据：使用read_msa函数读取多序列比对文件
2. 子采样处理：使用greedy_select选择代表性序列
3. 模型预测：将处理后的序列输入ESM模型获取接触预测
4. 结果可视化：将预测结果与真实结构接触矩阵比较

技术要点

• 接触定义：本教程使用Cβ原子间距离小于8Å作为接触标准
• MSA处理：需要移除插入字符和小写字母(代表非保守区域)
• 高效计算：使用向量化操作加速距离矩阵计算
• 模型选择：ESM提供多种预训练模型，可根据需求选择

通过本教程，读者可以掌握使用ESM模型进行蛋白质接触预测的完整流程，从数据准备到结果分析。这种方法在蛋白质结构预测和功能研究中具有重要应用价值。