蛋白质功能注释与分析🔗

import torch
import esm

# ---------- 第一步：加载模型和分词器 ----------
# 首次运行会自动下载模型权重（650M 约 2.5 GB）
model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
batch_converter = alphabet.get_batch_converter()
model.eval()  # 设为推理模式

# 如果有 GPU，将模型移到 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# ---------- 第二步：准备输入序列 ----------
# 格式：[(名称, 序列), ...]
data = [
    ("protein_1", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"),
    ("protein_2", "KALTARQQEVFDLIRDHISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE"),
]

batch_labels, batch_strs, batch_tokens = batch_converter(data)
batch_tokens = batch_tokens.to(device)

# ---------- 第三步：提取 embedding ----------
with torch.no_grad():
    results = model(batch_tokens, repr_layers=[33], return_contacts=True)

# 获取最后一层的 token embedding
# shape: (batch_size, seq_len + 2, 1280)  注意：+2 是因为有 <cls> 和 <eos> 特殊标记
token_embeddings = results["representations"][33]

# ---------- 第四步：生成序列级别的 embedding ----------
# 方法 1：取所有残基 embedding 的平均值（最常用）
# 注意要去掉首尾的特殊标记 <cls> 和 <eos>
sequence_embeddings = []
for i, (_, seq) in enumerate(data):
    seq_len = len(seq)
    # token_embeddings[i, 0] 是 <cls>，token_embeddings[i, seq_len+1] 是 <eos>
    emb = token_embeddings[i, 1:seq_len + 1].mean(0)  # shape: (1280,)
    sequence_embeddings.append(emb)
    print(f"{data[i][0]}: embedding shape = {emb.shape}")

# 方法 2：提取 <cls> token 的表示（类似 BERT 的 [CLS]）
cls_embeddings = token_embeddings[:, 0, :]  # shape: (batch_size, 1280)

# ---------- 第五步：获取接触图预测（附加功能）----------
# ESM-2 的注意力矩阵可以用来预测残基接触图
contact_map = results["contacts"][0].cpu().numpy()  # shape: (seq_len, seq_len)
print(f"Contact map shape: {contact_map.shape}")

# ESM-3 目前通过 EvolutionaryScale 的 Forge API 使用
# 开源版本为 esm3-open-small（1.4B 参数），可本地运行
# 完整版（98B）需要申请 API 访问

from esm.models.esm3 import ESM3
from esm.sdk.api import ESMProtein, GenerationConfig

# 加载开源小模型
model = ESM3.from_pretrained("esm3_sm_open_v1")

# 示例 1：从序列预测结构和功能
protein = ESMProtein(sequence="MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG")

# 生成结构
structure_prediction = model.generate(
    protein,
    GenerationConfig(track="structure", num_steps=8)
)

# 示例 2：给定功能约束生成蛋白质
# 使用功能关键词作为条件
protein_prompt = ESMProtein(
    function_annotations=["ATP binding", "kinase activity"]
)
generated_protein = model.generate(
    protein_prompt,
    GenerationConfig(track="sequence", num_steps=32)
)
print(f"生成的序列: {generated_protein.sequence}")

from transformers import T5Tokenizer, T5EncoderModel
import torch

# 加载模型（首次运行会下载约 11 GB）
tokenizer = T5Tokenizer.from_pretrained(
    "Rostlab/prot_t5_xl_uniref50", do_lower_case=False
)
model = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_uniref50")
model.eval()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# 准备序列（注意：ProtT5 要求氨基酸之间用空格分隔）
sequence = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"
sequence_spaced = " ".join(list(sequence))

# Tokenize
inputs = tokenizer(sequence_spaced, return_tensors="pt", add_special_tokens=True)
inputs = {k: v.to(device) for k, v in inputs.items()}

# 提取 embedding
with torch.no_grad():
    outputs = model(**inputs)

# outputs.last_hidden_state shape: (1, seq_len, 1024)
embeddings = outputs.last_hidden_state[0, :len(sequence), :]  # 去掉 padding
seq_embedding = embeddings.mean(0)  # 序列级别 embedding (1024,)
print(f"Per-residue embedding shape: {embeddings.shape}")
print(f"Sequence embedding shape: {seq_embedding.shape}")

import torch
import torch.nn as nn
import esm
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report

# ---------- 步骤 1：批量提取 embedding ----------
def extract_esm2_embeddings(sequences, model_name="esm2_t33_650M_UR50D"):
    """批量提取 ESM-2 sequence-level embedding"""
    model, alphabet = esm.pretrained.load_model_and_alphabet(model_name)
    batch_converter = alphabet.get_batch_converter()
    model.eval()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    embeddings = []
    # 分批处理以避免显存溢出
    batch_size = 8
    for i in range(0, len(sequences), batch_size):
        batch = [(f"seq_{j}", seq) for j, seq in enumerate(sequences[i:i+batch_size])]
        _, _, batch_tokens = batch_converter(batch)
        batch_tokens = batch_tokens.to(device)

        with torch.no_grad():
            results = model(batch_tokens, repr_layers=[33])

        for j, (_, seq) in enumerate(batch):
            seq_len = len(seq)
            emb = results["representations"][33][j, 1:seq_len + 1].mean(0)
            embeddings.append(emb.cpu().numpy())

    return np.array(embeddings)

# ---------- 步骤 2：准备数据（以溶解度预测为例）----------
# 假设你有一组蛋白质序列和对应的溶解/不溶标签
sequences = ["MKVLA...", "AKTGR...", ...]   # 替换为真实数据
labels = [1, 0, ...]                         # 1=可溶, 0=不溶

embeddings = extract_esm2_embeddings(sequences)

# ---------- 步骤 3：训练分类器 ----------
X_train, X_test, y_train, y_test = train_test_split(
    embeddings, labels, test_size=0.2, random_state=42
)

clf = LogisticRegression(max_iter=1000)
clf.fit(X_train, y_train)

# ---------- 步骤 4：评估 ----------
y_pred = clf.predict(X_test)
print(classification_report(y_test, y_pred))

from deepfrier import Predictor

# 初始化预测器（加载预训练模型）
predictor = Predictor.from_pretrained(
    gcn_model="trained_models/DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_ext_desc_MF.hdf5"
)

# 方式 1：从 PDB 文件预测（利用结构信息）
predictions = predictor.predict_from_pdb("your_protein.pdb")
for pred in predictions:
    print(f"GO:{pred.go_id}  {pred.go_name}  score={pred.score:.3f}")

# 方式 2：仅从序列预测
predictions = predictor.predict_from_sequence("MKTVRQERLK...")

# UniMol 使用示例（通过 uni-mol 包）
from unimol_tools import MolTrain, MolPredict

# 训练亲和力预测模型（需要提供训练数据）
clf = MolTrain(task='regression', data_type='molecule')
clf.fit(data='training_data.csv')

# 预测
pred = MolPredict(load_model='./exp')
result = pred.predict(data='test_data.csv')

# ThermoMPNN 使用示例
# 安装: pip install thermompnn (或从 GitHub 克隆)

from thermompnn import ThermoMPNN

model = ThermoMPNN.from_pretrained()

# 预测单点突变的 ΔΔG
result = model.predict(
    pdb_path="wild_type.pdb",
    chain_id="A",
    mutations=["A42G", "L108W", "V55I"]  # 格式：原始残基+位置+突变残基
)

for mut, ddg in zip(result.mutations, result.ddg_values):
    stability = "稳定化" if ddg < 0 else "去稳定化"
    print(f"{mut}: ΔΔG = {ddg:.2f} kcal/mol ({stability})")

# RaSP 使用
# GitHub: https://github.com/KULL-Centre/_2022_rasp

from rasp import predict_stability

# 对整个蛋白质的每个位置的所有可能突变进行扫描
results = predict_stability(
    pdb_file="protein.pdb",
    chain="A"
)
# 输出: DataFrame，每行一个突变，列包含位置、野生型、突变型、预测 ΔΔG

import pandas as pd

# 下载预计算数据
# 来源：https://zenodo.org/records/8208688
# 文件: AlphaMissense_aa_substitutions.tsv.gz (~4.2 GB)

# 加载数据
df = pd.read_csv(
    "AlphaMissense_aa_substitutions.tsv.gz", 
    sep='\t', 
    comment='#',
    compression='gzip'
)

# 查找特定蛋白质（以 BRCA1 为例，UniProt: P38398）
brca1 = df[df['uniprot_id'] == 'P38398']

# 查找特定突变
mutation = brca1[
    (brca1['protein_variant'] == 'C61G')  # BRCA1 的已知致病突变
]
print(mutation[['protein_variant', 'am_pathogenicity', 'am_class']])
# 输出类似: C61G  0.9876  likely_pathogenic

# 统计整个蛋白质的致病突变分布
print(brca1['am_class'].value_counts())

import matplotlib.pyplot as plt
import numpy as np

# 构建突变热图
protein_id = "P38398"  # BRCA1
protein_data = df[df['uniprot_id'] == protein_id].copy()

# 解析突变信息
protein_data['position'] = protein_data['protein_variant'].str.extract(r'(\d+)').astype(int)
protein_data['ref_aa'] = protein_data['protein_variant'].str[0]
protein_data['alt_aa'] = protein_data['protein_variant'].str[-1]

# 创建热图矩阵
amino_acids = list("ACDEFGHIKLMNPQRSTVWY")
positions = sorted(protein_data['position'].unique())
heatmap = np.full((len(amino_acids), len(positions)), np.nan)

for _, row in protein_data.iterrows():
    if row['alt_aa'] in amino_acids:
        aa_idx = amino_acids.index(row['alt_aa'])
        pos_idx = positions.index(row['position'])
        heatmap[aa_idx, pos_idx] = row['am_pathogenicity']

# 可视化（截取前 100 个位点）
fig, ax = plt.subplots(figsize=(20, 6))
im = ax.imshow(heatmap[:, :100], aspect='auto', cmap='RdYlGn_r', vmin=0, vmax=1)
ax.set_yticks(range(len(amino_acids)))
ax.set_yticklabels(amino_acids)
ax.set_xlabel('Residue Position')
ax.set_ylabel('Mutant Amino Acid')
ax.set_title(f'AlphaMissense Pathogenicity Landscape ({protein_id})')
plt.colorbar(im, label='Pathogenicity Score')
plt.tight_layout()
plt.savefig("alphamissense_landscape.png", dpi=150)
plt.show()

import torch
import esm
import numpy as np

def zero_shot_mutation_scoring(sequence, mutations, model_name="esm2_t33_650M_UR50D"):
    """
    使用 ESM-2 的遮蔽边际概率对突变进行零样本评分

    参数：
        sequence: 野生型序列 (str)
        mutations: 突变列表，格式如 ["A42G", "L108W"] (原始残基+位置+突变残基)
        model_name: ESM 模型名称

    返回：
        每个突变的 log-likelihood ratio 分数（越负 = 越有害）
    """
    model, alphabet = esm.pretrained.load_model_and_alphabet(model_name)
    batch_converter = alphabet.get_batch_converter()
    model.eval()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    scores = {}
    for mut in mutations:
        wt_aa = mut[0]
        pos = int(mut[1:-1]) - 1   # 转换为 0-indexed
        mut_aa = mut[-1]

        # 验证野生型氨基酸
        assert sequence[pos] == wt_aa, f"位置 {pos+1} 的氨基酸应为 {wt_aa}，实际为 {sequence[pos]}"

        # 构建遮蔽序列（将目标位置替换为 <mask>）
        masked_seq = list(sequence)
        masked_seq[pos] = "<mask>"
        masked_seq_str = "".join(masked_seq).replace("<mask>", "<mask>")

        # 需要用其他方式处理 mask token
        # 实际上，ESM 用 token index 来处理
        data = [("protein", sequence)]
        _, _, batch_tokens = batch_converter(data)
        batch_tokens = batch_tokens.to(device)

        # 遮蔽目标位置
        masked_tokens = batch_tokens.clone()
        masked_tokens[0, pos + 1] = alphabet.mask_idx  # +1 因为有 <cls>

        with torch.no_grad():
            logits = model(masked_tokens)["logits"]

        # 获取该位置的 log 概率
        log_probs = torch.log_softmax(logits[0, pos + 1], dim=-1)

        wt_prob = log_probs[alphabet.get_idx(wt_aa)].item()
        mut_prob = log_probs[alphabet.get_idx(mut_aa)].item()

        score = mut_prob - wt_prob  # log-likelihood ratio
        scores[mut] = score
        effect = "可能有害" if score < -2 else ("可能中性" if score > -1 else "效应不确定")
        print(f"{mut}: Δscore = {score:.3f} ({effect})")

    return scores

# 使用示例
sequence = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"
mutations = ["K2A", "R7A", "V5D"]  # 测试几个突变

scores = zero_shot_mutation_scoring(sequence, mutations)

import torch
import esm
import pandas as pd

def score_mutation_library(wt_sequence, mutation_list):
    """用 ESM-2 对突变库打分，返回排序后的突变列表"""
    model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
    batch_converter = alphabet.get_batch_converter()
    model.eval()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    # 先计算野生型各位置的 log 概率
    data = [("wt", wt_sequence)]
    _, _, tokens = batch_converter(data)
    tokens = tokens.to(device)

    results_list = []

    for mut_str in mutation_list:
        # 解析突变 (如 "A42G" 或多点突变 "A42G/L108W")
        total_score = 0.0
        single_muts = mut_str.split("/")

        for single_mut in single_muts:
            wt_aa = single_mut[0]
            pos = int(single_mut[1:-1])
            mut_aa = single_mut[-1]

            # 遮蔽该位置
            masked = tokens.clone()
            masked[0, pos] = alphabet.mask_idx  # ESM token 索引 = 序列位置 + 1 for <cls>

            with torch.no_grad():
                logits = model(masked)["logits"]

            log_probs = torch.log_softmax(logits[0, pos], dim=-1)
            wt_score = log_probs[alphabet.get_idx(wt_aa)].item()
            mut_score = log_probs[alphabet.get_idx(mut_aa)].item()
            total_score += (mut_score - wt_score)

        results_list.append({
            "mutation": mut_str,
            "esm2_score": total_score,
            "prediction": "有益" if total_score > 0 else "有害"
        })

    df = pd.DataFrame(results_list).sort_values("esm2_score", ascending=False)
    return df

# 使用示例
wt_seq = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"
mutations = ["K2R", "K2A", "K2D", "R7K", "R7A", "R7D", "V5I", "V5D"]

ranking = score_mutation_library(wt_seq, mutations)
print(ranking.to_string(index=False))