Python高效读取文件提取坐标并存储为一维数组的方法

问题场景 #

我有一个结构如下的txt文件：

@<TRIPOS>MOLECULE 2bsm_lig.pdb 45 47 0 0 0 SMALL USER_CHARGES
@<TRIPOS>ATOM
1 CL25 43.5837 12.4179 37.4396 Cl 1 BSM1 -0.1770
2 N1 40.4187 9.0729 42.8516 N.ar 1 BSM1 0.2996
3 H1 40.0025 9.0411 43.7713 H 1 BSM1 0.2700

文件开头的行是头部信息，可忽略。需要提取ATOM部分的第3、4、5列（x、y、z空间坐标），将其存储为一维数组，示例如下：

[43.5837, 12.4179, 37.4396, 40.4187, 9.0729, 42.8516, 40.0025, 9.0411, 43.7713]

我已使用Pandas在Python 3中实现该功能，但由于要在大循环中执行，速度过慢。请问有没有更高效的方法读取这些数据并存储为数组（可存储为列表或numpy数组）？

解决方案：放弃Pandas，用更轻量的方式处理 #

Pandas在处理单文件时虽然方便，但它的DataFrame初始化有不少额外开销，在大循环里批量处理小文件时会被放大，导致速度变慢。下面两个方案的效率都会比Pandas高很多：

1. 纯Python原生读取（最快，无额外依赖）

直接操作文件对象逐行处理，跳过头部后提取目标列，完全没有框架开销，适合对速度要求极高的场景：

def extract_coords_pure_python(file_path):
    coords = []
    with open(file_path, 'r') as f:
        # 跳过所有行，直到找到ATOM区块的标记
        for line in f:
            if line.strip().startswith('@<TRIPOS>ATOM'):
                break
        # 读取ATOM区块的每一行数据
        for line in f:
            line = line.strip()
            # 如果遇到下一个TRIPOS区块，停止读取
            if line.startswith('@<TRIPOS>'):
                break
            if not line:
                continue  # 跳过空行
            parts = line.split()
            # 提取第3、4、5列（对应索引2、3、4），转成float后加入列表
            coords.extend([float(parts[2]), float(parts[3]), float(parts[4])])
    return coords

调用这个函数后，会直接返回你需要的一维Python列表，完全符合示例格式。

2. 使用Numpy读取（适合后续数值计算）

如果你的后续操作需要对坐标做数值运算，Numpy的IO函数是C实现的，效率同样很高，而且直接生成Numpy数组，省去了类型转换的步骤：

import numpy as np

def extract_coords_numpy(file_path):
    # 先定位ATOM区块的起始行
    start_row = 0
    with open(file_path, 'r') as f:
        for idx, line in enumerate(f):
            if line.strip().startswith('@<TRIPOS>ATOM'):
                start_row = idx + 1  # 下一行开始是数据行
                break
    # 读取指定行、指定列的数据，并转成一维数组
    coords = np.loadtxt(
        file_path,
        skiprows=start_row,
        usecols=(2, 3, 4),  # 对应第3、4、5列
        dtype=np.float64
    ).flatten()
    return coords

这个函数返回的是一维Numpy数组，如果你需要转成Python列表，只需要调用coords.tolist()即可。

额外优化建议 #

• 如果所有文件的ATOM区块位置固定（比如都是第2行之后），可以直接跳过固定行数，不用遍历找标记行，能进一步提升速度。
• 大循环中，不要把模块导入（比如import numpy）放在循环内部，要放在脚本开头。
• 如果处理的是超大分子文件，可以考虑分块读取，但一般TRIPOS格式的分子文件用上面的方法完全足够。

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