优化Python读取含数百个160×160复数方阵的文本文件代码

我有一个数值模拟程序，输出的文本文件包含数百个160×160的复数方阵，格式示例如下：

(#,#)   (#,#)   (#,#)
 (#,#)   (#,#)   (#,#)
 (#,#)   (#,#)   (#,#)
-----

 (#,#)   (#,#)   (#,#)
 (#,#)   (#,#)   (#,#)
 (#,#)   (#,#)   (#,#)
-----

...

原本写的Python读取代码耗时极长，原代码如下：

with open(f'fort.{n}', 'r') as f:
    l = [[num for num in line.split(' ') if num != ''] for line in f]

l = list(filter((l[-2]).__ne__, l))

l = [[num for num in line if num != '\n'] for line in l]
l = [[split(num) for num in line] for line in l]

l = list(filter(([['-----\n']]).__ne__, l))

l = [[float(num[0])+1j*float(num[1]) for num in line] for line in l]

Solutions = []
for i in range(len(l)):
    if (i+1)%160 == 0:
        Solutions.append(l[i-159:i+1])

优化方案 #

核心优化思路 #

• 一次性读完整文件，减少磁盘IO的频繁调用（这是原代码慢的主要原因之一）
• 用正则批量提取所有复数的实部虚部，替代逐行逐元素的零散字符串操作
• 用numpy矢量化操作构造复数数组，比Python循环快得多
• 直接通过reshape拆分得到所有矩阵，避免手动分组的循环

优化后的代码 #

import re
import numpy as np

def load_complex_matrices(file_path, matrix_size=160):
    # 一次性读取全部文件内容，减少IO开销
    with open(file_path, 'r') as f:
        content = f.read()
    
    # 正则匹配所有 (实部,虚部) 格式的数值，自动跳过分隔线和空行
    # 支持正负整数、小数
    pattern = re.compile(r'\((-?\d+\.?\d*),(-?\d+\.?\d*)\)')
    all_num_pairs = pattern.findall(content)
    
    # 转成numpy数组，直接构造复数
    real = np.array([float(p[0]) for p in all_num_pairs], dtype=np.float64)
    imag = np.array([float(p[1]) for p in all_num_pairs], dtype=np.float64)
    complex_flat = real + 1j * imag
    
    # 计算总矩阵数，reshape成三维数组后拆分
    total_matrices = len(complex_flat) // (matrix_size * matrix_size)
    matrices_3d = complex_flat.reshape(total_matrices, matrix_size, matrix_size)
    # 转成二维数组的列表（如果需要保留三维数组也可以直接返回）
    solutions = [mat.squeeze() for mat in np.split(matrices_3d, total_matrices)]
    
    return solutions

# 使用示例
Solutions = load_complex_matrices(f'fort.{n}', matrix_size=160)

效率提升的原因 #

1. IO优化：一次性读取文件比逐行读取减少了大量系统调用，大文件下提升明显。
2. 批量字符串处理：正则一次性抓所有目标数值，替代原代码里多次split、filter的零散操作，字符串处理效率更高。
3. numpy矢量化：用numpy数组处理数值转换和复数构造，比纯Python循环快一个数量级以上。
4. 直接构造矩阵：通过reshape和split直接生成所有矩阵，省去了原代码里手动计数分组的循环，减少Python层面的迭代开销。

额外优化建议 #

如果文件大到内存放不下，可以分块处理：按-----拆分内容，每次处理一个矩阵的块；或者按固定行数读取（每个矩阵160行+1行分隔线）。另外也可以试试numpy.fromregex直接从文件生成数组：

# 用numpy.fromregex的简化写法
dt = np.dtype([('real', np.float64), ('imag', np.float64)])
pattern = r'\((?P<real>-?\d+\.?\d*),(?P<imag>-?\d+\.?\d*)\)'
data = np.fromregex(f'fort.{n}', pattern, dt)
complex_flat = data['real'] + 1j * data['imag']
# 后续reshape和拆分同上

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