如何以最快方式将Python列表追加到Numpy矩阵以处理超大规模分块科研数据

问题背景 #

我正在编写代码读取最高可达十亿行的科研数据，数据由多个分块组成，每个分块的头信息、数据集均不相同，因此需要逐行读取数据。我希望将有效数据存入Numpy矩阵中用于后续矩阵运算，当前核心代码如下：

with open(datafile, "r") as dump:
    i = 0           # block line number
    line_no = 0     # total line number
    block_size = 0
    block_count = 0
    for line in dump:
        values = line.rstrip().rsplit()
        i += 1
        line_no += 1
        if i <= self.head_line_no:
            print(line)  # for test
            if self.tag_block in line or i == 1:      # 1st line of a block
                # save block size after reading 1st block
                if block_size == 0 and block_count == 0:
                    block_size = line_no - 1
                    i = 1               # reset block line number
                    self.box = []       # reset box constant
                    print(self.matrix)
                    self.matrix = np.zeros((0, 0), dtype="float")   # reset matrix

                block_count += 1
            elif i == 2:
                self.timestamp.append(values[0])
            elif i == 3 or i == 5:
                continue
            elif i == 4:
                if self.atom_no != 0 and self.atom_no != values[0]:
                    self.warning_message = "atom number in timestep " + self.timestamp[-1] + "is inconsistent with" + self.timestamp[-2]
                    config.ConfigureUserEnv.log(self.warning_message)
                else:
                    pass
                self.atom_no = values[0]
            elif i == 6 or i == 7 or i == 8:
                self.box.append(values[0])
                self.box.append(values[1])
            elif i == self.head_line_no:
                values = line.rstrip().rsplit(":")
                for j in range(1,len(values)):
                    self.column_name.append(values[j])
        else:
            if self.matrix.size != 0:
                np_array = np.array(values)
                self.matrix = np.append(self.matrix, np.array(np.asarray(values)), 0)     
            else:
                np_array = np.array(values)
                self.matrix = np.zeros((1,len(values)), dtype="float")
                self.matrix = np.asarray(values)
    dump.close()
    print(self.matrix)       # for test
    print(self.matrix.size)  # for test

原始数据格式示例 #

ITEM: TIMESTEP
100
ITEM: NUMBER OF ATOMS
17587
ITEM: BOX BOUNDS pp pp pp
0.0000000000000000e+00 4.3491000000000000e+01
0.0000000000000000e+00 4.3491000000000000e+01
0.0000000000000000e+00 1.2994000000000000e+02
ITEM: ATOMS id type q xs ys zs 
59 1 1.80278 0.110598 0.129682 0.0359397 
297 1 1.14132 0.139569 0.0496654 0.00692627 
315 1 1.17041 0.0832356 0.00620818 0.00507927 
509 1 1.67165 0.0420777 0.113817 0.0313991 
590 1 1.65209 0.114966 0.0630015 0.0447129 
731 1 1.65143 0.0501253 0.13658 0.0108512 
1333 2 1.049 0.00850751 0.0526546 0.0406341 
...... 

期望输出格式 #

matrix = [[59 1 1.80278 0.110598 0.129682 0.0359397],
[297 1 1.14132 0.139569 0.0496654 0.00692627],
[315 1 1.17041 0.0832356 0.00620818 0.00507927],
...]

优化需求 #

当前处理的数据集规模极大，逐行追加Numpy数组的速度过慢，需要找到最高效的矩阵数据写入方案。 #

解决方案 #

现有代码性能瓶颈 #

现有代码最大的问题是逐行调用np.append，每次追加操作都会重新申请整块内存、拷贝原有数据，十亿行场景下时间复杂度为O(n²)，完全无法满足性能要求。此外逐行做字符串转numpy数组的操作也会产生大量冗余开销。

优化方案 #

1. 预分配内存
   提前统计数据总有效行数：你可以在第一次读取分块时拿到单个分块的原子数，再统计总分块数，直接创建大小为(总原子数, 列数)的空Numpy矩阵，后续只做索引赋值，完全避免追加操作。如果无法提前统计总行数，可采用动态扩容策略：初始预留10万行空间，空间用尽时按1.5~2倍大小扩容，大幅减少内存分配次数。
2. 批量转换写入
   不要逐行处理数据，每读完一个完整分块的所有原子行后，一次性将整组字符串列表转换成float类型的numpy数组，再整块赋值到预分配的矩阵对应位置，性能可以提升1~2个数量级。
3. 冗余操作清理
   原代码中rstrip().rsplit()可以直接替换为split()，默认会自动处理首尾空白和任意长度的分隔符，速度更快。同时删除首次矩阵赋值时冗余的零矩阵创建操作。
4. IO性能优化
   超大规模文件可以采用内存映射（np.memmap）方式读取，或者用pandas.read_csv指定skiprows和nrows批量读取分块，比Python层逐行迭代效率高很多。

优化后核心逻辑示例 #

# 提前获取列数、单个分块原子数后预分配矩阵
total_rows = 预估总行数
col_num = len(self.column_name)
self.matrix = np.empty((total_rows, col_num), dtype=np.float32) # 按需选精度，float32比float64省一半空间
current_row = 0
block_atom_num = int(self.atom_no)

# 处理数据行部分替换为批量写入
block_data = []
for line in dump:
    # 头信息处理逻辑保持不变，直到进入数据行部分
    if i > self.head_line_no:
        block_data.append(line.split())
        # 读完当前分块所有原子行
        if len(block_data) == block_atom_num:
            # 整块转换写入
            self.matrix[current_row:current_row+block_atom_num] = np.array(block_data, dtype=np.float32)
            current_row += block_atom_num
            block_data = []
            i = 0 # 重置分块行计数，进入下一个分块处理
# 最后截断预分配多的空间
self.matrix = self.matrix[:current_row]

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