Numba code slower than pure python

cython, numba, numpy, performance, python

Solution

The problem is that numba can't intuit the type of `lookup`. If you put a `print nb.typeof(lookup)` in your method, you'll see that numba is treating it as an object, which is slow. Normally I would just define the type of `lookup` in a locals dict, but I was getting a strange error. Instead I just created a little wrapper, so that I could explicitly define the input and output types.

@nb.jit(nb.f8[:](nb.f8[:]))
def numba_cumsum(x):
    return np.cumsum(x)

@nb.autojit
def numba_resample2(qs, xs, rands):
    n = qs.shape[0]
    #lookup = np.cumsum(qs)
    lookup = numba_cumsum(qs)
    results = np.empty(n)

    for j in range(n):
        for i in range(n):
            if rands[j] < lookup[i]:
                results[j] = xs[i]
                break
    return results

Then my timings are:

print "Timing Numba Function:"
%timeit numba_resample(qs, xs, rands)

print "Timing Revised Numba Function:"
%timeit numba_resample2(qs, xs, rands)

Timing Numba Function:
100 loops, best of 3: 8.1 ms per loop
Timing Revised Numba Function:
100000 loops, best of 3: 15.3 µs per loop

You can go even a little faster still if you use `jit` instead of `autojit`:

@nb.jit(nb.f8[:](nb.f8[:], nb.f8[:], nb.f8[:]))

For me that lowers it from 15.3 microseconds to 12.5 microseconds, but it's still impressive how well autojit does.

Problem

I've been working on speeding up a resampling calculation for a particle filter. As python has many ways to speed it up, I though I'd try them all. Unfortunately, the numba version is incredibly slow. As Numba should result in a speed up, I assume this is an error on my part. I tried 4 different versions: - Numba - Python - Numpy - Cython The code for each is below: ``` import numpy as np import scipy as sp import numba as nb from cython_resample import cython_resample @nb.autojit def numba_resample(qs, xs, rands): n = qs.shape[0] lookup = np.cumsum(qs) results = np.empty(n) for j in range(n): for i in range(n): if rands[j] < lookup[i]: results[j] = xs[i] break return results def python_resample(qs, xs, rands): n = qs.shape[0] lookup = np.cumsum(qs) results = np.empty(n) for j in range(n): for i in range(n): if rands[j] < lookup[i]: results[j] = xs[i] break return results def numpy_resample(qs, xs, rands): results = np.empty_like(qs) lookup = sp.cumsum(qs) for j, key in enumerate(rands): i = sp.argmax(lookup>key) results[j] = xs[i] return results #The following is the code for the cython module. It was compiled in a #separate file, but is included here to aid in the question. """ import numpy as np cimport numpy as np cimport cython DTYPE = np.float64 ctypedef np.float64_t DTYPE_t @cython.boundscheck(False) def cython_resample(np.ndarray[DTYPE_t, ndim=1] qs, np.ndarray[DTYPE_t, ndim=1] xs, np.ndarray[DTYPE_t, ndim=1] rands): if qs.shape[0] != xs.shape[0] or qs.shape[0] != rands.shape[0]: raise ValueError("Arrays must have same shape") assert qs.dtype == xs.dtype == rands.dtype == DTYPE cdef unsigned int n = qs.shape[0] cdef unsigned int i, j cdef np.ndarray[DTYPE_t, ndim=1] lookup = np.cumsum(qs) cdef np.ndarray[DTYPE_t, ndim=1] results = np.zeros(n, dtype=DTYPE) for j in range(n): for i in range(n): if rands[j] < lookup[i]: results[j] = xs[i] break return results """ if __name__ == '__main__': n = 100 xs = np.arange(n, dtype=np.float64) qs = np.array([1.0/n,]*n) rands = np.random.rand(n) print "Timing Numba Function:" %timeit numba_resample(qs, xs, rands) print "Timing Python Function:" %timeit python_resample(qs, xs, rands) print "Timing Numpy Function:" %timeit numpy_resample(qs, xs, rands) print "Timing Cython Function:" %timeit cython_resample(qs, xs, rands) ``` This results in the following output: ``` Timing Numba Function: 1 loops, best of 3: 8.23 ms per loop Timing Python Function: 100 loops, best of 3: 2.48 ms per loop Timing Numpy Function: 1000 loops, best of 3: 793 µs per loop Timing Cython Function: 10000 loops, best of 3: 25 µs per loop ``` Any idea why the numba code is so slow? I assumed it would be at least comparable to Numpy. Note: if anyone has any ideas on how to speed up either the Numpy or Cython code samples, that would be nice too:) My main question is about Numba though.

Original source