我用python写了一个函数来计算高斯加宽中的Delta函数,其中涉及4级循环。但是效率很低,比用Fortran做类似的事情慢10倍左右。
def Delta_Gaussf(Nw, N_bd, N_kp, hw, eigv): Delta_Gauss = np.zeros((Nw,N_kp,N_bd,N_bd),dtype=float) for w1 in range(Nw): for k1 in range(N_kp): for i1 in range(N_bd): for j1 in range(N_bd): if ( j1 >= i1 ): Delta_Gauss[w1][k1][i1][j1] = np.exp(pow((eigv[k1][j1]-eigv[k1][i1]-hw[w1])/width,2)) return Delta_Gauss
我删除了一些常量以使其看起来更简单。
为了获得最佳性能,我推荐使用 Numba(易于使用,性能良好)。或者 Cython 可能是一个好主意,但需要对代码进行一些更改。
实际上,您已经把所有事情都做对了,并且实现了一个易于理解(对于人类来说,最重要的是对于编译器来说)的解决方案。
提高绩效的方法基本上有两种
使用 Numba 的解决方案
您不需要做太多改变,我将 pow 函数改为 np.power,并对 numpy 中访问数组的方式做了一些细微的改变(这不是真正必要的)。
import numba as nb import numpy as np #performance-debug info import llvmlite.binding as llvm llvm.set_option('', '--debug-only=loop-vectorize') @nb.njit(fastmath=True) def Delta_Gaussf_nb(Nw, N_bd, N_kp, hw, width,eigv): Delta_Gauss = np.zeros((Nw,N_kp,N_bd,N_bd),dtype=float) for w1 in range(Nw): for k1 in range(N_kp): for i1 in range(N_bd): for j1 in range(N_bd): if ( j1 >= i1 ): Delta_Gauss[w1,k1,i1,j1] = np.exp(np.power((eigv[k1,j1]-eigv[k1,i1]-hw[w1])/width,2)) return Delta_Gauss
由于“if”,SIMD 矢量化失败。在下一步中,我们可以将其删除(可能需要在 njited 函数之外进行调用np.triu(Delta_Gauss))。我还并行化了该函数。
np.triu(Delta_Gauss)
@nb.njit(fastmath=True,parallel=True) def Delta_Gaussf_1(Nw, N_bd, N_kp, hw, width,eigv): Delta_Gauss = np.zeros((Nw,N_kp,N_bd,N_bd),dtype=np.float64) for w1 in nb.prange(Nw): for k1 in range(N_kp): for i1 in range(N_bd): for j1 in range(N_bd): Delta_Gauss[w1,k1,i1,j1] = np.exp(np.power((eigv[k1,j1]-eigv[k1,i1]-hw[w1])/width,2)) return Delta_Gauss
表现
Nw = 20 N_bd = 20 N_kp = 20 width=20 hw = np.linspace(0., 1.0, Nw) eigv = np.zeros((N_kp, N_bd),dtype=np.float) Your version: 0.5s first_compiled version: 1.37ms parallel version: 0.55ms
这些简单的优化可使速度提高约 1000 倍。