Numba : Cell Vars Are Not Supported

I'd like to use numba to speed up this function: 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(wy

Solution 1:

Numba does not work particularly well currently on nested list of lists (as of v0.21 at least). I believe that this is what the 'cell vars' error is referring to, but I'm not 100% sure. Below, I convert everything to numpy arrays to enable the code to be optimized by numba:

import numpy as np
import numba as nb
import math

defrownowaga(u, v):
    wymiar_x = len(u)
    wymiar_y = len(u[1])
    f = [[[0for j inrange(wymiar_y)] for i inrange(wymiar_x)] for k inrange(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 inrange( wymiar_x):
        for j inrange (wymiar_y):
            for k inrange(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)defrownowaga_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

Now let's set up some test arrays:

u = [[math.sin(i) + math.cos(j) for j in range(40)] for i in range(1000)]
y = [[math.sin(i) + math.cos(j) for j in range(40)] for i in range(1000)]

u_np = np.array(u)
y_np = np.array(y)

First let's verify that my numba code is giving the same answer as the OP's code:

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

From an ipython notebook:

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

And now let's time things on my laptop:

In [15] %timeitf1=rownowaga(u,y)1loops,best of 3:288msperloopIn [16] %timeitf2=rownowaga_numba(u_np,y_np)1000 loops,best of 3:973µsperloop

So we get a nice 300x speed-up with minimal code changes. Just to note, I'm using a nightly build of Numba from a little before 0.22:

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