evaluation.py

import os
import sys
import numpy as np
import pandas as pd
from Bio.PDB import PDBParser, PPBuilder, Superimposer
from Bio.PDB.Polypeptide import is_aa
from Bio import BiopythonWarning
import warnings
from tmtools import tm_align  # 导入tm_align函数

# 忽略 Biopython 的警告
warnings.simplefilter('ignore', BiopythonWarning)


def extract_sequence(pdb_path):
    """
    从 PDB 文件中提取第一个链的氨基酸序列。

    参数:
        pdb_path (str): PDB 文件的路径。

    返回:
        str: 氨基酸序列（单字母代码），如果无法提取则返回 None。
    """
    parser = PDBParser(QUIET=True)
    try:
        structure = parser.get_structure('structure', pdb_path)
    except Exception as e:
        print(f"无法解析 PDB 文件 {pdb_path}: {e}")
        return None

    ppb = PPBuilder()
    seq = ""
    for pp in ppb.build_peptides(structure):
        seq += str(pp.get_sequence())
        # 只取第一个多肽链
        break

    return seq if seq else None


def get_sequences(directory):
    """
    获取指定目录下所有 PDB 文件的序列。

    参数:
        directory (str): PDB 文件夹路径。

    返回:
        dict: 键为 PDB 文件名，值为其氨基酸序列。
    """
    seq_dict = {}
    for filename in os.listdir(directory):
        if filename.endswith('.pdb'):
            filepath = os.path.join(directory, filename)
            seq = extract_sequence(filepath)
            if seq:
                seq_dict[filename] = seq
            else:
                print(f"警告: 无法提取序列 {filename}")
    return seq_dict


def match_pdbs(wetlab_seqs, pdb_files_seqs):
    """
    根据序列内容匹配 PDB 文件。

    参数:
        wetlab_seqs (dict): wetlab_pdb 文件夹中的 PDB 文件及其序列。
        pdb_files_seqs (dict): pdb_files 文件夹中的 PDB 文件及其序列。

    返回:
        list of tuples: 每个元组包含 (wetlab_pdb, pdb_files_pdb)。
    """
    matches = []
    pdb_files_seqs_reverse = {}
    for pdb, seq in pdb_files_seqs.items():
        pdb_files_seqs_reverse.setdefault(seq, []).append(pdb)

    for wetlab_pdb, wetlab_seq in wetlab_seqs.items():
        matched_pdbs = pdb_files_seqs_reverse.get(wetlab_seq, [])
        if matched_pdbs:
            for matched_pdb in matched_pdbs:
                matches.append((wetlab_pdb, matched_pdb))
        else:
            print(f"警告: 在 pdb_files 中未找到匹配的序列 for {wetlab_pdb} {wetlab_seq}")
    return matches


def get_ca_atoms(pdb_path):
    """
    提取 PDB 文件中第一个链的所有 C-alpha 原子。

    参数:
        pdb_path (str): PDB 文件的路径。

    返回:
        list of Bio.PDB.Atom.Atom: C-alpha 原子列表，如果失败则返回 None。
    """
    parser = PDBParser(QUIET=True)
    try:
        structure = parser.get_structure('structure', pdb_path)
    except Exception as e:
        print(f"无法解析 PDB 文件 {pdb_path}: {e}")
        return None

    ca_atoms = []
    models = list(structure.get_models())
    if not models:
        print(f"警告: PDB 文件 {pdb_path} 没有模型。")
        return None
    first_model = models[0]

    chains = list(first_model.get_chains())
    if not chains:
        print(f"警告: PDB 文件 {pdb_path} 没有链。")
        return None
    first_chain = chains[0]

    for residue in first_chain:
        if is_aa(residue, standard=True):
            if 'CA' in residue:
                ca = residue['CA']
                ca_atoms.append(ca)

    return ca_atoms if ca_atoms else None


def calculate_rmsd(pdb_path1, pdb_path2):
    """
    计算两个 PDB 文件之间的 RMSD。

    参数:
        pdb_path1 (str): 第一个 PDB 文件路径。
        pdb_path2 (str): 第二个 PDB 文件路径。

    返回:
        float: RMSD 值，如果计算失败则返回 None。
    """
    atoms1 = get_ca_atoms(pdb_path1)
    atoms2 = get_ca_atoms(pdb_path2)

    if atoms1 is None or atoms2 is None:
        print(f"警告: 无法获取 C-alpha 原子 for {pdb_path1} and/or {pdb_path2}")
        return None

    if len(atoms1) != len(atoms2):
        print(f"警告: {os.path.basename(pdb_path1)} 和 {os.path.basename(pdb_path2)} 的 C-alpha 原子数量不同 ({len(atoms1)} vs {len(atoms2)}).")
        return None

    # 创建 Superimposer 对象
    sup = Superimposer()
    sup.set_atoms(atoms1, atoms2)
    sup.apply(atoms2)  # 应用旋转和平移到 atoms2，如果需要
    rmsd = sup.rms

    return rmsd


def calculate_tm_score_tmtools(coords1, coords2, seq1, seq2):
    """
    使用 tmtools 计算两个坐标阵列之间的 TM-score。

    参数:
        coords1 (np.ndarray): 第一个坐标数组，形状为 N1 x 3。
        coords2 (np.ndarray): 第二个坐标数组，形状为 N2 x 3。
        seq1 (str): 第一个序列。
        seq2 (str): 第二个序列。

    返回:
        float: TM-score 值，如果计算失败则返回 None。
    """
    try:
        res = tm_align(coords1, coords2, seq1, seq2)
        tm_score = res.tm_norm_chain2  # 使用 tm_norm_chain2 作为 TM-score
        return tm_score
    except Exception as e:
        print(f"Error calculating TM-score with tmtools: {e}")
        return None


def process_pair_rmsd_tm(pdb1_path, pdb2_path, seq1, seq2):
    """
    处理一对 PDB 文件，计算 TM-score 和 RMSD。

    参数:
        pdb1_path (str): 第一个 PDB 文件路径。
        pdb2_path (str): 第二个 PDB 文件路径。
        seq1 (str): 第一个序列。
        seq2 (str): 第二个序列。

    返回:
        tuple: (PDB_ID, TM_score, RMSD)
    """
    pdb1 = os.path.basename(pdb1_path)
    pdb2 = os.path.basename(pdb2_path)

    # 获取 C-alpha 坐标
    atoms1 = get_ca_atoms(pdb1_path)
    atoms2 = get_ca_atoms(pdb2_path)

    if atoms1 is None or atoms2 is None:
        print(f"警告: 无法获取 C-alpha 原子 for {pdb1} and/or {pdb2}")
        return (f"{pdb1} vs {pdb2}", None, None)

    if len(atoms1) != len(atoms2):
        print(f"警告: {pdb1} 和 {pdb2} 的 C-alpha 原子数量不同 ({len(atoms1)} vs {len(atoms2)}).")
        return (f"{pdb1} vs {pdb2}", None, None)

    # 提取坐标为 NumPy 数组
    coords1 = np.array([atom.get_coord() for atom in atoms1])
    coords2 = np.array([atom.get_coord() for atom in atoms2])

    # 计算 RMSD
    rmsd = calculate_rmsd(pdb1_path, pdb2_path)

    # 计算 TM-score
    tm_score = calculate_tm_score_tmtools(coords1, coords2, seq1, seq2)

    return (f"{pdb1} vs {pdb2}", tm_score, rmsd)


def main():
    # 配置文件夹路径
    wetlab_dir = "./wetlab_pdb"  # 替换为您的 wetlab_pdb 文件夹路径
    pdb_files_dir = "/home/ubuntu/alphafold3/output_pdb/"  # 替换为您的 pdb_files 文件夹路径
    output_csv = "alphafold3_tm_rmsd_results.csv"  # 输出的 CSV 文件名

    # 检查文件夹是否存在
    if not os.path.isdir(wetlab_dir):
        print(f"错误: 文件夹 '{wetlab_dir}' 不存在。")
        sys.exit(1)
    if not os.path.isdir(pdb_files_dir):
        print(f"错误: 文件夹 '{pdb_files_dir}' 不存在。")
        sys.exit(1)

    print("提取 wetlab_pdb 文件夹中的序列...")
    wetlab_seqs = get_sequences(wetlab_dir)
    print(f"提取到 {len(wetlab_seqs)} 个序列。\n")

    print("提取 pdb_files 文件夹中的序列...")
    pdb_files_seqs = get_sequences(pdb_files_dir)
    print(f"提取到 {len(pdb_files_seqs)} 个序列。\n")

    print("匹配 PDB 文件 based on sequences...")
    matches = match_pdbs(wetlab_seqs, pdb_files_seqs)
    print(f"找到 {len(matches)} 对匹配的 PDB 文件。\n")

    if not matches:
        print("没有找到任何匹配的 PDB 文件。")
        sys.exit(0)

    results = []

    print("开始计算 TM-score 和 RMSD...\n")
    for wetlab_pdb, ref_pdb in matches:
        print(f"处理: {wetlab_pdb} vs {ref_pdb}")
        wetlab_path = os.path.join(wetlab_dir, wetlab_pdb)
        ref_path = os.path.join(pdb_files_dir, ref_pdb)
        try:
            # 获取序列
            seq1 = wetlab_seqs[wetlab_pdb]
            seq2 = pdb_files_seqs[ref_pdb]

            pdb_id, tm_score, rmsd = process_pair_rmsd_tm(wetlab_path, ref_path, seq1, seq2)
            if tm_score is not None and rmsd is not None:
                print(f"TM-score: {tm_score:.4f}, RMSD: {rmsd:.4f} Å\n")
            else:
                print("计算失败。\n")
        except Exception as e:
            print(f"错误: 处理 {wetlab_pdb} vs {ref_pdb} 时出错: {e}\n")
            pdb_id, tm_score, rmsd = (f"{wetlab_pdb} vs {ref_pdb}", None, None)
        results.append((pdb_id, tm_score, rmsd))

    # 创建 DataFrame
    df = pd.DataFrame(results, columns=['PDB_ID', 'TM_score', 'RMSD'])

    # 过滤掉无法计算的结果
    valid_df = df.dropna(subset=['TM_score', 'RMSD'])

    if not valid_df.empty:
        # 计算平均值和标准差
        avg_tm = valid_df['TM_score'].mean()
        std_tm = valid_df['TM_score'].std()
        avg_rmsd = valid_df['RMSD'].mean()
        std_rmsd = valid_df['RMSD'].std()

        # 创建摘要 DataFrame
        summary = pd.DataFrame({
            'PDB_ID': ['Average', 'Std Dev'],
            'TM_score': [avg_tm, std_tm],
            'RMSD': [avg_rmsd, std_rmsd]
        })

        # 合并摘要和详细结果
        final_df = pd.concat([summary, df], ignore_index=True)
    else:
        final_df = df.copy()

    # 保存到 CSV
    final_df.to_csv(output_csv, index=False)
    print(f"结果已保存到 '{output_csv}'。")

    # 打印摘要
    if not valid_df.empty:
        print("\n### 计算摘要 ###")
        print(f"平均 TM-score: {avg_tm:.4f}")
        print(f"TM-score 标准差: {std_tm:.4f}")
        print(f"平均 RMSD: {avg_rmsd:.4f} Å")
        print(f"RMSD 标准差: {std_rmsd:.4f} Å")
    else:
        print("没有成功计算任何 PDB 对的 TM-score 和 RMSD。")


if __name__ == "__main__":
    main()