Python实现PDB文件原子坐标距离筛选的代码开发需求

核心逻辑 #

• 遍历目标整数列表，为每个元素匹配PDB文件中对应行，提取基准坐标
• 逐行扫描PDB文件，计算每行坐标与基准坐标的欧氏距离
• 根据距离阈值筛选符合条件的行，提取指定字段存入新列表
• 去重并排序，补充初始目标元素，保证输出符合预期

完整代码 #

def process_pdb(target_list, pdb_file):
    new_list = []
    # 一次性读取文件所有行，减少IO操作
    with open(pdb_file, 'r') as f:
        lines = f.readlines()
    
    for target in target_list:
        # 查找匹配的基准行，提取坐标
        base_x, base_y, base_z = None, None, None
        for line in lines:
            # 提取第6-11位的序号（Python索引左闭右开，对应5:11）
            res_id = line[5:11].strip()
            if res_id == str(target):
                # 提取30-38、38-46、46-54位的坐标
                base_x = float(line[29:38].strip())
                base_y = float(line[38:46].strip())
                base_z = float(line[46:54].strip())
                break
        if not base_x:
            continue  # 未找到匹配行则跳过
        
        # 遍历所有行计算距离并筛选
        for line in lines:
            try:
                # 提取当前行坐标
                curr_x = float(line[29:38].strip())
                curr_y = float(line[38:46].strip())
                curr_z = float(line[46:54].strip())
            except ValueError:
                continue  # 跳过非坐标行
            
            # 计算欧氏距离
            sd = ((curr_x - base_x)**2 + (curr_y - base_y)**2 + (curr_z - base_z)**2)**0.5
            
            # 筛选距离在2到6之间的行
            if 2 < sd < 6:
                res_num = line[5:11].strip()
                if res_num.isdigit():
                    new_list.append(int(res_num))
    
    # 去重排序，补充初始目标元素（预期输出包含这些值）
    new_list = sorted(list(set(new_list)))
    for num in target_list:
        if num not in new_list:
            new_list.append(num)
    new_list.sort()
    return new_list

# 调用示例
targets = [1, 3, 4]
result = process_pdb(targets, 'test.pdb')
print(f"new_list = {result}")

关键细节说明 #

• 索引对应：注意PDB文件的位数描述对应Python字符串的左闭右开索引，比如第6-11位实际是line[5:11]
• 异常处理：跳过格式错误的行，避免程序崩溃
• 结果修正：初始目标元素的欧氏距离为0（小于2），不会被自动存入，所以手动补充进结果列表，保证与预期输出一致

内容的提问来源于stack exchange，提问作者Mehdi Irani