Speeding up take with numba?

Question

Is it possible to speed up np.take using numba?

Here is my attempt, but it is a lot slower. I can't use nopython mode as it does not like the np.empty_like command.

import numba
import numpy as np
from timer import Timer    

def take( x, indices ):
    result = np.empty_like( indices, dtype=x.dtype )    
    for i in range( len( indices ) ):
        result[i] = x[ indices[ i ] ]
    return result

jtake = numba.jit("f4[:](f4[:],i4[:])" )( take )

if __name__=='__main__':

    N = 100000
    m = 100
    idx = np.random.random_integers( 0, N, m )
    x = np.random.randn( N )

    num_tests=10000

    with Timer( 'take' ):
        for i in range( num_tests ):    
            r0 = take( x, idx )

    with Timer( 'Numba take' ):
        for i in range( num_tests ):    
            r1 = jtake( x, idx )                

    with Timer( 'Numpy.take' ):
        for i in range( num_tests ):
            r2 = x.take( idx )

Which has results:

Beginning take
take took 2.46 seconds
Beginning Numba take
Numba take took 1.11 seconds
Beginning Numpy.take
Numpy.take took 0.04 seconds

Answer 1

The answer is no.

Numba will not act on compiled functions such as np.take() or the array methods that work with fancy indexing, which is the basis for your algorithm. Numba acts on the interpreted part of your code.

Your take() function probably has more overhead than NumPy's and Numba has improved the for loops (interpreted).

With NumPy > 1.9 your code should get closer to NumPy's take since your algorithm is based on fancy indexing and they improved the fancy indexing efficiency to the level of np.take().