如何用Python处理含关联数据的列：原子与键数据清理需求

需求与挑战 #

• 原始数据包含Atoms和Bonds两个分段，核心需求为：删除Atoms段中第5、6、7列（x/y/z坐标）数值大于50的行
• 附带关联处理要求：
  1. Atoms段中被删除的Atom_id，需要同步清理Bonds段中包含该Atom_id的所有行
  2. 清理完成后，需将Atom_id、Bond_id更新为连续序列，同时同步更新Bonds段内关联的Atom_id

示例原始数据 #

#Atoms
#Atom_id    molecules_id    atom_type   charge  x_coor  y_coor  z_coor
       1       1     2    -0.834           60.243           56.013           55.451  
       2       1     1     0.417           1.061            6.406            5.263  
       3       1     1     0.417           3.513            2.071             5.526  
       4       2     2    -0.834             4.14           6.861            5.328  
       5       2     1     0.417            4.322            6.96            5.317  
       6       2     1     0.417            3.303           1.922            4.912  
       7       3     2    -0.834           12.756           53.344            3.856  
       8       3     1     0.417           12.527           53.366            4.833  
       9       3     1     0.417           12.039           53.747            3.454  
      10       4     2    -0.834           1.402            7.122           13.392  
       
 #Bonds
 #Bond_id   bond_type   atom_id atom_id
       1       2       1       2    
       2       2       1       3    
       3       2       4       5  
       4       2       4       6   
       5       2       7       8   
       6       2       7       9   
       7       2      10      11   
       8       2      10      12   
       9       2      13      14   
      10       2      13      15    
      11       2      16      17 

拆分文件的优势 #

将Atoms和Bonds拆分到两个独立文件确实能显著降低脚本编写的复杂度，核心原因有两点：

1. 逻辑拆分更清晰：可以先单独处理Atoms段，筛选有效原子并生成旧ID到新ID的映射表；再基于该映射表处理Bonds段，过滤无效行并更新关联ID，两段处理逻辑完全解耦
2. 减少状态管理难度：单文件处理需要同时维护两段的读取指针、筛选状态，拆分后可分别完成读取、处理、写入操作，避免跨段的状态冲突

基于拆分文件的处理步骤 #

1. 处理Atoms文件 #

• 读取Atoms文件，保留段标识和表头行，遍历数据行：
  • 检查第5、6、7列的数值，只要任意一个大于50则跳过该行
  • 为保留的行分配连续的新Atom_id，同时记录旧Atom_id -> 新Atom_id的映射字典
• 将处理后的Atoms数据（含更新后的ID）写入新文件

2. 处理Bonds文件 #

• 读取Bonds文件，保留段标识和表头行，遍历数据行：
  • 提取第3、4列的Atom_id，若任意一个不在映射字典中（对应原子已被删除），则跳过该行
  • 对有效行，用映射字典替换为新的Atom_id，同时分配连续的新Bond_id
• 将处理后的Bonds数据写入新文件

示例Python脚本 #

# 处理Atoms段，生成ID映射
atom_id_map = {}
new_atom_id = 1
processed_atoms_content = []

# 读取原始Atoms文件
with open("atoms_raw.txt", "r") as f:
    lines = f.readlines()
    # 保留段标识和表头行
    processed_atoms_content.append(lines[0])
    processed_atoms_content.append(lines[1])
    
    for line in lines[2:]:
        line = line.strip()
        if not line:
            continue
        cols = line.split()
        # 提取x/y/z坐标并判断
        x = float(cols[4])
        y = float(cols[5])
        z = float(cols[6])
        if x > 50 or y > 50 or z > 50:
            continue
        # 记录旧ID到新ID的映射
        old_atom_id = int(cols[0])
        atom_id_map[old_atom_id] = new_atom_id
        # 更新当前行的Atom_id
        cols[0] = str(new_atom_id)
        processed_atoms_content.append(" ".join(cols) + "\n")
        new_atom_id += 1

# 写入处理后的Atoms文件
with open("atoms_processed.txt", "w") as f:
    f.writelines(processed_atoms_content)

# 处理Bonds段
new_bond_id = 1
processed_bonds_content = []

# 读取原始Bonds文件
with open("bonds_raw.txt", "r") as f:
    lines = f.readlines()
    # 保留段标识和表头行
    processed_bonds_content.append(lines[0])
    processed_bonds_content.append(lines[1])
    
    for line in lines[2:]:
        line = line.strip()
        if not line:
            continue
        cols = line.split()
        atom1 = int(cols[2])
        atom2 = int(cols[3])
        # 检查两个原子是否都保留
        if atom1 not in atom_id_map or atom2 not in atom_id_map:
            continue
        # 更新Bond_id和关联的Atom_id
        cols[0] = str(new_bond_id)
        cols[2] = str(atom_id_map[atom1])
        cols[3] = str(atom_id_map[atom2])
        processed_bonds_content.append(" ".join(cols) + "\n")
        new_bond_id += 1

# 写入处理后的Bonds文件
with open("bonds_processed.txt", "w") as f:
    f.writelines(processed_bonds_content)

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