Python大数组列表创建加速：万级原子势能计算性能优化

• 研究体系规模：12134个原子（三维点），计算目标为原子相关物理量
• 已有输入数据：
  • 字典对象：存储每个原子的类型、电荷及各类计算参数
  • 近邻列表：存储每个原子距离小于阈值的所有近邻原子索引
  • 平方距离矩阵：存储所有原子对之间的平方距离
• 计算目标：给定原子与其所有近邻的势能，计算过程需用到原子间距离、原子属性参数、近邻列表信息
• 现有纯Python实现代码：

def LJ(sii,sjj,Eii,Ejj,rij):
    """
    Fonction that compute the LJ potential. one of the quantities

    sii : sigma for atom i ( a parameter )
    sjj : sigma for atom j ( a parameter )
    Eii : epsilon for atom i ( a parameter )
    Ejj : epsilon for atom j ( a parameter )
    rij : squared distance between atom i and j 
    """
    sij = (sii+sjj)*0.5
    Eij = (Eii*Ejj)**0.5
    return ((4*Eij*(sij**12))/(rij**6)) - ((4*Eij*(sij**6))/(rij**3))

def Coul(qi,qj,rij):
    """
    Fonction that compute the Coul potential. another quantities
    qi : charge of atom i ( a parameter )
    qj : charge of atom j ( a parameter )
    rij : squared distance between atom i and j 
    """
    V = 138.935458*((qi*qj)/(rij**0.5)) # because rij is squared i take the square root
    return V

def energies(NL,sP,D):
    """
    NL : a vector of neighbors with as index the position and as value a list of indexes of neighbors
        for example [[1,2,3],[0,2]...] the first atom at position 0 as 3 neighbors 1,2 and 3. The second atom at position 1 as 2 neighbors 0 and 2 etc...
    sP : a distionnary giving the parameters for a given atoms, and also neighbors distant of 4 atoms in the mist 'dihedrals'. The key correspond to the index of the distance matric and the neighboring list, but it's a string
        sample : {'key': '0_NH3', 'num': 0, 'type': 'NH3', 'resnr': 0, 'residu': 'ASP', 'atom': 'N', 'charge': -0.3, 'sigma': 0.329632525712, 'epsilon': 0.8368, 'voisins': [1, 2, 3, 4, 5, 6, 12], 'dihedrals': [7, 8, 9, 13, 14]}
    D : the squared distance matrix between all atoms
    """
    
    coul,ljk,ljk14,coul14 = [],[],[],[] # I want the potential for each atom and I want 4 types of potentials
    # coul and ljk are potentials calculated on neighboring atoms in the NL neighbors list
    # coul14 and ljk14 are potentials calculated on neighboring atoms in the sP[key]['dihedrals'] list

    for k in sP.keys(): # For each atom 
        # k is a string that represent an int sP['0'] as the informations for D[0][:] or D[:][0] and NL[0] 
        ljk.append(sum([LJ(sP[k]['sigma'],sP[str(n)]['sigma'],sP[k]['epsilon'],sP[str(n)]['epsilon'],D[int(k)][n]) for n in NL[int(k)]]))
        coul.append(sum([Coul(sP[k]['charge'],sP[str(n)]['charge'],D[int(k)][n]) for n in NL[int(k)]]))
        ljk14.append(sum([LJ(sP[k]['sigma'],sP[str(n)]['sigma'],sP[k]['epsilon'],sP[str(n)]['epsilon'],D[int(k)][n]) for n in sP[k]['dihedrals']]))
        coul14.append(sum([Coul(sP[k]['charge'],sP[str(n)]['charge'],D[int(k)][n]) for n in sP[k]['dihedrals']]))
        # because I want the sum of all potentials for each atom I use a sum on a comprehensive list of potential between a given atom and it's neighbors

• 现有性能问题：12134原子体系下运行耗时约40秒，需通过numpy向量化等手段进一步压缩计算耗时。

1. 预处理消除冗余Python层开销 #

现有代码的核心性能损耗首先来自循环内的重复操作：反复执行int(k)/str(n)类型转换、字符串键字典查询、原生列表遍历，这类操作没有利用硬件加速，单步耗时远高于数值计算。
在正式计算前先做一次预处理，把所有原子参数按索引顺序抽为numpy一维数组，把近邻、二面角列表提前规整为整数索引格式，预处理耗时可忽略，能直接砍掉大半纯Python开销：

import numpy as np
n_atoms = len(NL)
# 按原子索引存储参数的numpy数组
sigma = np.zeros(n_atoms, dtype=np.float64)
epsilon = np.zeros(n_atoms, dtype=np.float64)
charge = np.zeros(n_atoms, dtype=np.float64)
dihedrals = []
for idx in range(n_atoms):
    atom_info = sP[str(idx)]
    sigma[idx] = atom_info['sigma']
    epsilon[idx] = atom_info['epsilon']
    charge[idx] = atom_info['charge']
    dihedrals.append(atom_info['dihedrals'])
# 把平方距离矩阵转为numpy数组，索引取值速度比原生列表快数十倍
D = np.asarray(D, dtype=np.float64)

2. 全numpy向量化计算替换逐元素循环 #

不要在Python层逐对原子调用势能函数，先把所有需要计算的原子对拼成连续索引数组，一次性批量取数、计算，最后用np.bincount按中心原子索引做分组求和，比纯Python列表推导快1~2个数量级。
核心实现逻辑：

def calc_energies_vec(NL, dihedrals, sigma, epsilon, charge, D):
    # 子函数：输入近邻列表，返回对应的LJ和库仑势能总和数组
    def _calc_pairs(neighbor_list):
        # 拼接所有(中心原子i, 近邻j)的索引对
        ij_pairs = []
        for i in range(n_atoms):
            js = np.array(neighbor_list[i], dtype=np.int64)
            ij_pairs.append(np.stack([np.full(len(js), i, dtype=np.int64), js], axis=1))
        ij_pairs = np.concatenate(ij_pairs, axis=0)
        i_idx, j_idx = ij_pairs[:,0], ij_pairs[:,1]
        
        # 批量取所有原子对的参数和距离
        sii, sjj = sigma[i_idx], sigma[j_idx]
        eii, ejj = epsilon[i_idx], epsilon[j_idx]
        qi, qj = charge[i_idx], charge[j_idx]
        rij = D[i_idx, j_idx]
        
        # 批量计算LJ势能
        sij = (sii + sjj) * 0.5
        eij = np.sqrt(eii * ejj)
        inv_r3 = 1 / (rij**3)
        inv_r6 = inv_r3 ** 2
        lj_total = 4 * eij * (sij**12 * inv_r6 - sij**6 * inv_r3)
        # 按中心原子i分组求和
        lj_sum = np.bincount(i_idx, weights=lj_total, minlength=n_atoms)
        
        # 批量计算库仑势能
        coul_total = 138.935458 * qi * qj / np.sqrt(rij)
        coul_sum = np.bincount(i_idx, weights=coul_total, minlength=n_atoms)
        
        return lj_sum, coul_sum
    
    # 分别计算普通近邻、1-4二面角近邻的势能
    ljk, coul = _calc_pairs(NL)
    ljk14, coul14 = _calc_pairs(dihedrals)
    return coul, ljk, ljk14, coul14

该实现完全避免了Python层的逐原子、逐近邻循环，所有数值计算都在numpy的C后端执行，单轮计算耗时通常可以压到1秒以内。

3. 可选进阶加速 #

• 若需要进一步提速，可给_calc_pairs函数加numba的@njit装饰器，不需要修改业务逻辑，依靠JIT编译能再提速2~5倍，总耗时可压缩到百毫秒级。
• 当前实现会分别计算i→j和j→i的双向相互作用，如果你的势能计算不需要重复统计双向作用，可在拼接ij_pairs时只保留i < j的原子对，计算完成后把势能同时累加到i和j的总和中，直接减少一半计算量。
• 计算LJ势能时，可提前预计算所有原子的s_pow6 = sigma**6、s_pow12 = s_pow6**2，避免在原子对层面重复计算sigma的幂次，进一步减少计算量。

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