Python逐行读取大文件做数值求和平均的性能优化方案咨询

我编写了Python代码实现逐行读取文件并执行平均、求和运算，目前需要对代码进行提速优化。当前待处理的pressurefile文件共有945,670行（后续行数还会进一步增长）。

原始代码 #

这是我最初提交的版本，后续我将根据大家的建议优化后更新了最新版本，一并提供：

def time_average():
try:
    filename = mem.pressurefile
    navg = mem.NFRAMES
    dz = mem.dz
    zlo = mem.zlo
    NZ = mem.NZ
    mass = mem.mass

    dens_fact = amu_to_kg / (mem.slab_V * ang3_to_m3)
    
    array_pxx = np.zeros([NZ,1])
    array_pyy = np.zeros([NZ,1])
    array_pzz = np.zeros([NZ,1])
    array_ndens = np.zeros([NZ,1])
    
    array_density = np.zeros([NZ,1])
    array_enthalpy = np.zeros([NZ,1])
    array_surf_tens = np.zeros([NZ,1])
    
    counter = 0
    with open(filename) as f:
        for line in f:
            line.strip("\n")
            #content = [_ for _ in line.split()]
            content = line.split()
            if len(content) == 7:
                z = float(content[3]) - zlo
                pxx = float(content[4])
                pyy = float(content[5])
                pzz = float(content[6])
                
                loc = math.floor(z/dz)
                if loc >= NZ:
                    loc = loc - NZ
                elif loc < 0:
                    loc = loc + NZ   
                #print(z, loc, zlo)
                
                array_pxx[loc] += pxx
                array_pyy[loc] += pyy
                array_pzz[loc] += pzz
                array_ndens[loc] += 1
            counter += 1
    for col in range(NZ):
        array_pxx[col] /= navg
        array_pyy[col] /= navg
        array_pzz[col] /= navg
        array_ndens[col] /= navg
        array_density[col] = mass * dens_fact * array_ndens[col]
        
    return (array_density, array_enthalpy, array_surf_tens)
except IndexError as err:
    writelog (err)
    writelog(float(content[3]) , loc, zlo)

性能表现与已尝试的优化方案 #

使用cProfile对主代码进行性能分析，发现该辅助函数处理74.4MB文件耗时约10s，该耗时不符合预期。目前已经尝试了如下优化方案：

• 方案1：cython3
  按照如下方式使用cython编译：

  cython3 --embed -o ptythinfile.c ptythinfile.py
  gcc -Os -I /usr/include/python3.8 -o ptythinfile ptythinfile.c -lpython3.8 -lpthread -lm -lutil -ldl
  

  该方案未带来任何性能提升。
• 方案2：C/C++
  将全部代码转换为C/C编译运行。实际上我最初的版本就是C开发的，但调试难度极高，因此切换到Python，不愿意再采用该方案。
• 方案3：Pypy3
  尝试使用Pypy3运行，遇到兼容性问题，我本地安装的是Python3.8和3.9版本，而pypy3依赖Python3.6，因此放弃该方案。
• 方案4：外部C库
  我查阅了将辅助函数编译为C代码再在Python中调用的教程，这是我下一步计划尝试的方案。

我在搜索时还看到了shedskin等可选方案，想请教大家该代码片段的最优优化路径，以及其他可尝试的提速方案。

2021年10月21日更新 #

我根据专家的建议更新了代码，测试运行正常，平均执行时间从约10s降低到约9.4s。
补充说明：pressurefile是LAMMPS软件的输出文件，文件前几行内容示例如下：

ITEM: TIMESTEP
50100
ITEM: NUMBER OF ATOMS
2744
ITEM: BOX BOUNDS pp pp pp
-2.5000000000000000e+01 2.5000000000000000e+01
-2.5000000000000000e+01 2.5000000000000000e+01
-7.5000000000000000e+01 7.5000000000000000e+01
ITEM: ATOMS id x y z c_1[1] c_1[2] c_1[3]
2354 18.8358 -21.02 -70.5731 -21041.8 -3738.18 -2520.84
1708 5.54312 -8.1526 -62.6984 4362.84 -30610.2 -4065.84

仅需要处理最后两行这类长度为7的内容。

最新优化代码 #

def time_average():
try:
    filename = mem.pressurefile
    navg = mem.NFRAMES
    dz = mem.dz
    zlo = mem.zlo
    NZ = mem.NZ
    mass = mem.mass

    dens_fact = amu_to_kg / (mem.slab_V * ang3_to_m3)
    
    array_pxx = np.zeros([NZ,1])
    array_pyy = np.zeros([NZ,1])
    array_pzz = np.zeros([NZ,1])
    array_ndens = np.zeros([NZ,1])
    
    #array_density = np.zeros([NZ,1])
    array_enthalpy = np.zeros([NZ,1])
    array_surf_tens = np.zeros([NZ,1])
    
    counter = 0
    locList = []
    pxxList = []
    pyyList = []
    pzzList = []
    with open(filename) as f:
        for line in f:
            #line.strip("\n")
            #content = [_ for _ in line.split()]
            content = line.split()
            if len(content) == 7:
                z = float(content[3]) - zlo
                pxx = float(content[4])
                pyy = float(content[5])
                pzz = float(content[6])
                
                #loc = math.floor(z/dz)
                loc = int(z // dz)
                
                if loc >= NZ:
                    loc = loc - NZ
                elif loc < 0:
                    loc = loc + NZ   
                #print(z, loc, zlo)
                
                # Not great but much faster than using Numpy functions
                locList.append(loc)
                pxxList.append(pxx)
                pyyList.append(pyy)
                pzzList.append(pzz)
            counter += 1

    # Very fast list-to-Numpy-array conversion
    locList = np.array(locList, dtype=np.int32)
    pxxList = np.array(pxxList, dtype=np.float64)
    pyyList = np.array(pyyList, dtype=np.float64)
    pzzList = np.array(pzzList, dtype=np.float64)

    # Fast accumulate
    np.add.at(array_pxx[:,0], locList, pxxList)
    np.add.at(array_pyy[:,0], locList, pyyList)
    np.add.at(array_pzz[:,0], locList, pzzList)
    np.add.at(array_ndens[:,0], locList, 1)

    array_pxx /= navg
    array_pyy /= navg
    array_pzz /= navg
    array_ndens /= navg
    array_density = mass * dens_fact * array_ndens

    return (array_density, array_enthalpy, array_surf_tens)
except IndexError as err:
    writelog (err)
    print(loc)
    writelog(float(content[3]) , loc, zlo)

测试设备配置 #

• 处理器：Intel® Xeon(R) W-2255 CPU @ 3.70GHz × 20
• 内存：16 GB
• 显卡：NVIDIA Corporation GP107GL [Quadro P620]
• 系统：64位Ubuntu 20.04.3 LTS

当前最新代码平均执行时间约为2.6s，相比原始版本提速3倍，感谢用户@JeromeRichard的贡献。

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