小能豆

如何优化python中的多层循环?

py

我用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

我删除了一些常量以使其看起来更简单。


阅读 6

收藏
2024-11-04

共1个答案

小能豆

为了获得最佳性能,我推荐使用 Numba(易于使用,性能良好)。或者 Cython 可能是一个好主意,但需要对代码进行一些更改。

实际上,您已经把所有事情都做对了,并且实现了一个易于理解(对于人类来说,最重要的是对于编译器来说)的解决方案。

提高绩效的方法基本上有两种

  1. 按照@scnerd 所示对代码进行矢量化。这通常比简单地编译一个仅使用一些 for 循环的相当简单的代码要慢一些,也更复杂一些。不要对代码进行矢量化,然后使用编译器。从简单的循环方法来看,这通常需要做一些工作,并导致更慢、更复杂的结果。这个过程的优点是你只需要 numpy,它是几乎每个处理一些数值计算的 Python 项目的标准依赖项。
  2. 编译代码。如果你已经有了一个只包含几个循环而没有其他循环的解决方案,或者只包含几个非 numpy 函数,那么这通常是最简单、最快的解决方案。

使用 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))。我还并行化了该函数。

@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 倍。

2024-11-04