Numba @jit 无法优化简单函数

Question

我有一个非常简单的函数，它使用 Numpy 数组和 for 循环，但是添加 Numba @jit 装饰器绝对没有加速：

# @jit(float64[:](int32,float64,float64,float64,int32))
@jit
def Ising_model_1D(N=200,J=1,T=1e-2,H=0,n_iter=1e6):
    beta = 1/T
    s = randn(N,1) > 10  
    s[N-1] = s[0]
    mag = zeros((n_iter,1))
    aux_idx =  randint(low=0,high=N,size=(n_iter,1))

    for i1 in arange(n_iter):
        rnd_idx = aux_idx[i1]
        s_1 = s[rnd_idx]*2 - 1
        s_2 = s[(rnd_idx+1)%(N)]*2 - 1
        s_3 = s[(rnd_idx-1)%(N)]*2 - 1
        delta_E = 2.0*J*(s_2+s_3)*s_1 + 2.0*H*s_1

        if(delta_E < 0):
            s[rnd_idx] = np.logical_not(s[rnd_idx]) 
        elif(np.exp(-1*beta*delta_E) >= rand()):
            s[rnd_idx] = np.logical_not(s[rnd_idx])
        s[N-1] = s[0]
        mag[i1] = (s*2-1).sum()*1.0/N 
    return mag

另一方面，MATLAB 只需不到 0.5 秒即可完成运行！为什么 Numba 缺少如此基本的东西？

Answer 1

这是对您的代码进行的修改，在我的机器上运行大约需要 0.4 秒：

def ising_model_1d(N=200,J=1,T=1e-2,H=0,n_iter=1e6):
    n_iter = int(n_iter)
    beta = 1/T
    s = randn(N) > 10
    s[N-1] = s[0]

    mag = zeros(n_iter)
    aux_idx =  randint(low=0,high=N,size=n_iter)

    pre_rand = rand(n_iter)

    _ising_jitted(n_iter, aux_idx, s, J, N, H, beta, pre_rand, mag)

    return mag


@jit(nopython=True)
def _ising_jitted(n_iter, aux_idx, s, J, N, H, beta, pre_rand, mag):
    for i1 in range(n_iter):
        rnd_idx = aux_idx[i1]
        s_1 = s[rnd_idx*2] - 1
        s_2 = s[(rnd_idx+1)%(N)]*2 - 1
        s_3 = s[(rnd_idx-1)%(N)]*2 - 1
        delta_E = 2.0*J*(s_2+s_3)*s_1 + 2.0*H*s_1
        t = rand()
        if delta_E < 0:
            s[rnd_idx] = not s[rnd_idx]
        elif np.exp(-1*beta*delta_E) >= pre_rand[i1]:
            s[rnd_idx] = not s[rnd_idx]

        s[N-1] = s[0]
        mag[i1] = (s*2-1).sum()*1.0/N

请确保结果符合预期！我改了很多你的，不能保证计算正确！

使用 numba 需要一点小心。 Python 函数，以及大多数 numpy 函数，无法由编译器优化。我发现有用的一件事是使用 nopython 选项到 @jit。这意味着只要你给它一些它不能真正优化的代码，编译器就会报错。然后您可以查看错误消息并找到可能会降低代码速度的行。

我发现，诀窍是在 Python 中编写一个 "gateway" 函数，使用 numpy 及其向量化函数完成尽可能多的工作。它应该创建您需要存储结果的空数组。它应该打包您在计算过程中需要的所有数据。然后它应该将所有这些传递到一个又大又长的参数列表中的 jitted 函数中。

恰当的例子：请注意我如何在 jitted 代码中处理随机数生成。在您的原始代码中，您调用了 rand():

elif(np.exp(-1*beta*delta_E) >= rand()):

但是 rand() 不能被 numba 优化（至少在旧版本的 numba 中是这样。在新版本中它可以，前提是 rand 是不带参数调用）。观察结果是，对于 n_iter 次迭代中的每一次迭代，您都需要一个随机数。因此，我们只需在我们的包装函数中使用 numpy 创建一个随机数组，然后将该随机数组提供给 jitted 函数。获得一个随机数就像索引这个数组一样简单。

最后，有关可由最新版本的编译器优化的 numpy 函数的列表，请参阅 here。在我修改你的代码时，我积极地删除了对 numpy 函数的调用，以便代码可以在 numba.

的更多版本上工作

Numba @jit 无法优化简单函数

Numba @jit fails to optimise simple function

python

numpy

numba