Numba:不支持单元变量

Question

我想用numba来加速这个功能：

from numba import jit
@jit
def rownowaga_numba(u, v):
    wymiar_x = len(u)
    wymiar_y = len(u[1])
    f = [[[0 for j in range(wymiar_y)] for i in range(wymiar_x)] for k in range(9)]
    cx = [0., 1., 0., -1., 0., 1., -1., -1., 1.]
    cy = [0., 0., 1., 0., -1., 1., 1., -1., -1.]
    w = [4./9, 1./9, 1./9, 1./9, 1./9, 1./36, 1./36, 1./36, 1./36] 
    for i in range( wymiar_x):
        for j in range (wymiar_y):
            for k in range(9):
                up = u[i][j]
                vp = v[i][j]
                udot = (up**2 + vp**2)
                cu = up*cx[k] + vp*cy[k]
                f[k][i][j] =  w[k] + w[k]*(3.0*cu + 4.5*cu**2 - 1.5*udot)
     return f

我用这样的数据进行测试：

^{pr2}$

我得到了这个错误：

    Traceback (most recent call last):
  File "C:\Users\Ricevind\Desktop\PyPy\Skrypty\Rownowaga.py", line 29, in <module>
    f = rownowaga_pypy(u,y)
  File "C:\pyzo2014a\lib\site-packages\numba\dispatcher.py", line 171, in _compile_for_args
    return self.compile(sig)
  File "C:\pyzo2014a\lib\site-packages\numba\dispatcher.py", line 348, in compile
    flags=flags, locals=self.locals)
  File "C:\pyzo2014a\lib\site-packages\numba\compiler.py", line 637, in compile_extra
    return pipeline.compile_extra(func)
  File "C:\pyzo2014a\lib\site-packages\numba\compiler.py", line 356, in compile_extra
    raise e
  File "C:\pyzo2014a\lib\site-packages\numba\compiler.py", line 351, in compile_extra
    bc = self.extract_bytecode(func)
  File "C:\pyzo2014a\lib\site-packages\numba\compiler.py", line 343, in extract_bytecode
    bc = bytecode.ByteCode(func=self.func)
  File "C:\pyzo2014a\lib\site-packages\numba\bytecode.py", line 343, in __init__
    raise NotImplementedError("cell vars are not supported")
NotImplementedError: cell vars are not supported

我最感兴趣的是“不支持cell vars”的含义，因为Google返回的结果没有意义。在

Answer 1

Numba目前在嵌套列表列表（至少从v0.21开始）上工作得不太好。我相信这就是'cell vars'错误所指的，但我不是100%确定。下面，我将所有内容转换为numpy数组，以便numba对代码进行优化：

import numpy as np
import numba as nb
import math

def rownowaga(u, v):
    wymiar_x = len(u)
    wymiar_y = len(u[1])
    f = [[[0 for j in range(wymiar_y)] for i in range(wymiar_x)] for k in range(9)]
    cx = [0., 1., 0., -1., 0., 1., -1., -1., 1.]
    cy = [0., 0., 1., 0., -1., 1., 1., -1., -1.]
    w = [4./9, 1./9, 1./9, 1./9, 1./9, 1./36, 1./36, 1./36, 1./36] 
    for i in range( wymiar_x):
        for j in range (wymiar_y):
            for k in range(9):
                up = u[i][j]
                vp = v[i][j]
                udot = (up**2 + vp**2)
                cu = up*cx[k] + vp*cy[k]
                f[k][i][j] =  w[k] + w[k]*(3.0*cu + 4.5*cu**2 - 1.5*udot)
    return f

# Pull these out so that numba treats them as constant arrays
cx = np.array([0., 1., 0., -1., 0., 1., -1., -1., 1.])
cy = np.array([0., 0., 1., 0., -1., 1., 1., -1., -1.])
w = np.array([4./9, 1./9, 1./9, 1./9, 1./9, 1./36, 1./36, 1./36, 1./36]) 

@nb.jit(nopython=True)
def rownowaga_numba(u, v):
    wymiar_x = u.shape[0]
    wymiar_y = u[1].shape[0]
    f = np.zeros((9, wymiar_x, wymiar_y))

    for i in xrange( wymiar_x):
        for j in xrange (wymiar_y):
            for k in xrange(9):
                up = u[i,j]
                vp = v[i,j]
                udot = (up*up + vp*vp)
                cu = up*cx[k] + vp*cy[k]
                f[k,i,j] =  w[k] + w[k]*(3.0*cu + 4.5*cu**2 - 1.5*udot)
    return f

现在让我们设置一些测试数组：

^{pr2}$

首先，让我们验证我的numba代码是否给出了与OP代码相同的答案：

f1 = rownowaga(u, y)
f2 = rownowaga_numba(u_np, y_np)

从ipython笔记本上：

In [13]: np.allclose(f2, np.array(f1))
Out[13]:
True

现在让我们在我的笔记本电脑上计时：

In [15] %timeit f1 = rownowaga(u, y)
1 loops, best of 3: 288 ms per loop


In [16] %timeit f2 = rownowaga_numba(u_np, y_np)
1000 loops, best of 3: 973 µs per loop

因此，我们得到了一个不错的300倍的加速与最小的代码变化。我要注意的是，我使用的是从0.22之前的一个夜间构建的Numba：

In [16]: nb.__version__
Out[16]:
'0.21.0+137.gac9929d'