如何加速Python中分子动力学模拟相关的for循环？

问题背景 #

以下是用于分子动力学模拟的Python代码，其中resampling函数内的1000万次迭代循环运行极慢（100万步耗时约50分钟），尝试GPU加速后效果更差，需要优化循环效率：

import MDAnalysis as mda
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from tqdm import tqdm as tq
import MDAnalysis.analysis.pca as pca
import random
import math  # 修复原代码未导入math的问题

def PCA_projection(pdb,dcd,atomgroup):
    u = mda.Universe(pdb,dcd)
    PSF_pca = pca.PCA(u, select=atomgroup)
    PSF_pca.run(verbose=True)
    n_pcs = np.where(PSF_pca.results.cumulated_variance > 0.95)[0][0]
    atomgroup = u.select_atoms(atomgroup)
    pca_space = PSF_pca.transform(atomgroup, n_components=n_pcs)
    PC1_proj = [pca_space[i][0] for i in range(len(pca_space))]
    PC2_proj = [pca_space[i][1] for i in range(len(pca_space))]
    return PC1_proj, PC2_proj


def Read_bias_potential(bias_potential):
    Bias_potential = pd.read_csv(bias_potential)
    Bias_potential = Bias_potential['En-User']
    Bias_potential = Bias_potential.values.tolist()
    W = [math.exp((-1 * i) / (0.001987*300)) for i in Bias_potential]
    return W

def Bin(PC1_prj, PC2_prj, frame_num, min_br1, max_br1, min_br2, max_br2, bin_num, W):
    data1 = PC1_prj[0:frame_num]
    bins1 = np.linspace(min_br1, max_br1, bin_num)
    bins1 = np.round(bins1,2)
    digitized1 = np.digitize(data1, bins1)
    binc1 = np.arange(min_br1 + (max_br1 - min_br1)/2*bin_num, 
max_br1 + (max_br1 - min_br1)/2*bin_num, (max_br1 - min_br1)/bin_num, dtype = float)
    binc1 = np.around(binc1,3)

    data2 = PC2_prj[0:frame_num]
    bins2 = np.linspace(min_br2, max_br2, bin_num)
    bins2 = np.round(bins2,2)
    digitized2 = np.digitize(data2, bins2)
    binc2 = np.arange(min_br2 + (max_br2 - min_br2)/2*bin_num, max_br2 + (max_br2 - min_br2)/2*bin_num, (max_br2 - min_br2)/bin_num, dtype = float)
    binc2 = np.around(binc2,3)

    w_array = np.zeros((bin_num,bin_num))

    for j in range(frame_num):
        w_array[digitized1[j]][digitized2[j]] += (W[digitized1[j]] + W[digitized2[j]])

    for m in range(bin_num):
        for n in range(bin_num):
            if w_array[m][n] == 0:
                w_array[m][n] = 1e-100

    return w_array, binc1, binc2

def gaussian(Sj1,Slj1,Sj2,Slj2,count):
    sigma1 = 0.5
    sigma2 = 0.5
    Kb = 0.001987204
    T = 300
    h0 = 0.0001
    g = 0
    C1 = 0
    C2 = 0
    idx_j1 = np.where(Slj1 == Sj1)[0][0]
    idx_j2 = np.where(Slj2 == Sj2)[0][0]
    for i in range(idx_j2 -5, idx_j2 +6):
        C2 = i if 0<=i<1000 else (i+1000 if i<0 else i-1000)
        for j in range(idx_j1 -5, idx_j1 +6):
            C1 = j if 0<=j<1000 else (j+1000 if j<0 else j-1000)
            g += count[C2,C1] * h0 * np.exp( (-(Sj1 - Slj1[C1])**2/(2*sigma1**2)) + (-(Sj2 - Slj2[C2])**2/(2*sigma2**2)) )
    return np.exp(-g/(Kb*T))

def resampling(binc1, binc2, w_array):
    l =1000
    F = np.zeros((l,l))
    count = np.zeros((l,l))
    Wn = w_array

    for i in tq(range(10000000)):
        SK1 = random.choice(binc1)
        SK2 = random.choice(binc2)
        SL1 = random.choice(binc1)
        SL2 = random.choice(binc2)
        while SK1 == SL1:
            SL1 = random.choice(binc1)
        while SK2 == SL2:
            SL2 = random.choice(binc2)
        idx_sk1 = np.where(binc1 == SK1)[0][0]
        idx_sk2 = np.where(binc2 == SK2)[0][0]
        idx_sl1 = np.where(binc1 == SL1)[0][0]
        idx_sl2 = np.where(binc2 == SL2)[0][0]
        
        F[idx_sk2][idx_sk1] = gaussian(SK1,binc1,SK2,binc2,count)
        F[idx_sl2][idx_sl1] = gaussian(SL1,binc1,SL2,binc2,count)
        
        W_SK = Wn[idx_sk2][idx_sk1] * F[idx_sk2][idx_sk1]
        W_SL = Wn[idx_sl2][idx_sl1] * F[idx_sl2][idx_sl1]
        
        if W_SK <= W_SL:
            selected_idx1, selected_idx2 = idx_sl1, idx_sl2
        else:
            a = random.random()
            if W_SL/W_SK >= a:
                selected_idx1, selected_idx2 = idx_sl1, idx_sl2
            else:
                selected_idx1, selected_idx2 = idx_sk1, idx_sk2
        count[selected_idx2][selected_idx1] += 1
    return F

注：已修复原代码两处问题：补充math模块导入；修正resampling中重复赋值F的错误。

核心优化方案 #

1. 预构建值-索引映射，消除重复查找 #

原代码每次循环多次调用np.where查找索引，时间复杂度为O(n)，提前构建字典映射可将查找降为O(1)：

def resampling(binc1, binc2, w_array):
    # 预构建值到索引的映射
    binc1_map = {val: idx for idx, val in enumerate(binc1)}
    binc2_map = {val: idx for idx, val in enumerate(binc2)}
    l =1000
    F = np.zeros((l,l))
    count = np.zeros((l,l))
    Wn = w_array

    for i in tq(range(10000000)):
        SK1 = random.choice(binc1)
        SK2 = random.choice(binc2)
        SL1 = random.choice(binc1)
        SL2 = random.choice(binc2)
        while SK1 == SL1:
            SL1 = random.choice(binc1)
        while SK2 == SL2:
            SL2 = random.choice(binc2)
        # 直接通过字典获取索引
        idx_sk1 = binc1_map[SK1]
        idx_sk2 = binc2_map[SK2]
        idx_sl1 = binc1_map[SL1]
        idx_sl2 = binc2_map[SL2]
        # ... 后续代码不变

2. 向量化gaussian函数，消除内部嵌套循环 #

用numpy切片和广播替代gaussian中的嵌套循环，利用numpy的C级运算加速：

def gaussian(Sj1, Slj1, Sj2, Slj2, count, sigma1=0.5, sigma2=0.5, Kb=0.001987204, T=300, h0=0.0001):
    idx_j1 = np.where(Slj1 == Sj1)[0][0]
    idx_j2 = np.where(Slj2 == Sj2)[0][0]
    
    # 生成索引范围并处理边界循环
    idx_range1 = np.arange(idx_j1 -5, idx_j1 +6)
    idx_range1 = np.where(idx_range1 <0, idx_range1+1000, idx_range1)
    idx_range1 = np.where(idx_range1 >=1000, idx_range1-1000, idx_range1)
    
    idx_range2 = np.arange(idx_j2 -5, idx_j2 +6)
    idx_range2 = np.where(idx_range2 <0, idx_range2+1000, idx_range2)
    idx_range2 = np.where(idx_range2 >=1000, idx_range2-1000, idx_range2)
    
    # 广播生成二维矩阵，向量化计算
    slj1_vals = Slj1[idx_range1]
    slj2_vals = Slj2[idx_range2]
    count_vals = count[idx_range2[:, None], idx_range1]
    
    gauss_term1 = np.exp(-(Sj1 - slj1_vals)**2/(2*sigma1**2))
    gauss_term2 = np.exp(-(Sj2 - slj2_vals)**2/(2*sigma2**2))
    total_gauss = count_vals * h0 * gauss_term1 * gauss_term2
    
    g = total_gauss.sum()
    return np.exp(-g/(Kb*T))

3. 用Numba编译核心函数，将Python循环转为机器码 #

Numba对循环密集型代码加速效果显著，给gaussian和resampling添加编译装饰器：

from numba import njit, prange

@njit(fastmath=True)
def gaussian(Sj1, Slj1, Sj2, Slj2, count, sigma1=0.5, sigma2=0.5, Kb=0.001987204, T=300, h0=0.0001):
    idx_j1 = np.where(Slj1 == Sj1)[0][0]
    idx_j2 = np.where(Slj2 == Sj2)[0][0]
    g = 0.0
    for i in range(idx_j2 -5, idx_j2 +6):
        C2 = i
        if C2 <0:
            C2 +=1000
        elif C2 >=1000:
            C2 -=1000
        for j in range(idx_j1 -5, idx_j1 +6):
            C1 = j
            if C1 <0:
                C1 +=1000
            elif C1 >=1000:
                C1 -=1000
            term1 = -(Sj1 - Slj1[C1])**2/(2*sigma1**2)
            term2 = -(Sj2 - Slj2[C2])**2/(2*sigma2**2)
            g += count[C2, C1] * h0 * np.exp(term1 + term2)
    return np.exp(-g/(Kb*T))

@njit(fastmath=True)
def resampling(binc1, binc2, w_array):
    l = 1000
    F = np.zeros((l,l), dtype=np.float64)
    count = np.zeros((l,l), dtype=np.float64)
    Wn = w_array
    len_binc1 = len(binc1)
    len_binc2 = len(binc2)
    
    for i in range(10000000):
        # 直接随机选择索引，避免后续查找
        idx_sk1 = np.random.randint(0, len_binc1)
        idx_sl1 = np.random.randint(0, len_binc1-1)
        if idx_sl1 >= idx_sk1:
            idx_sl1 +=1
        
        idx_sk2 = np.random.randint(0, len_binc2)
        idx_sl2 = np.random.randint(0, len_binc2-1)
        if idx_sl2 >= idx_sk2:
            idx_sl2 +=1
        
        SK1 = binc1[idx_sk1]
        SK2 = binc2[idx_sk2]
        SL1 = binc1[idx_sl1]
        SL2 = binc2[idx_sl2]
        
        F[idx_sk2, idx_sk1] = gaussian(SK1, binc1, SK2, binc2, count)
        F[idx_sl2, idx_sl1] = gaussian(SL1, binc1, SL2, binc2, count)
        
        W_SK = Wn[idx_sk2, idx_sk1] * F[idx_sk2, idx_sk1]
        W_SL = Wn[idx_sl2, idx_sl1] * F[idx_sl2, idx_sl1]
        
        if W_SK <= W_SL:
            selected_idx1, selected_idx2 = idx_sl1, idx_sl2
        else:
            a = np.random.rand()
            if W_SL/W_SK >= a:
                selected_idx1, selected_idx2 = idx_sl1, idx_sl2
            else:
                selected_idx1, selected_idx2 = idx_sk1, idx_sk2
        count[selected_idx2, selected_idx1] +=1
    return F

4. 优化随机选择逻辑，消除while循环 #

原代码用while确保SK和SL不同，改用索引偏移法直接生成符合条件的随机数：

# 替换resampling中的随机选择部分
idx_sk1 = np.random.randint(0, len_binc1)
# 生成不等于idx_sk1的随机索引
idx_sl1 = np.random.randint(0, len_binc1-1)
if idx_sl1 >= idx_sk1:
    idx_sl1 +=1

idx_sk2 = np.random.randint(0, len_binc2)
idx_sl2 = np.random.randint(0, len_binc2-1)
if idx_sl2 >= idx_sk2:
    idx_sl2 +=1

效果预期 #

通过上述优化，尤其是Numba编译和向量化改造，循环速度可提升10-100倍，100万步耗时可压缩至几分钟以内。GPU效果差的原因是单步计算量过小，内存拷贝开销抵消了计算优势，而Numba的CPU编译更适配这种循环次数多、单步计算量小的场景。

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