Numba Invalid use of Function with argument(s) of type(s)

user_1977628 0
Viewed 2574

I'm using Numba nonpython mode and some NumPy functions.

@njit
def invert(W, copy=True):
    '''
    Inverts elementwise the weights in an input connection matrix.
    In other words, change the from the matrix of internode strengths to the
    matrix of internode distances.

    If copy is not set, this function will *modify W in place.*

    Parameters
    ----------
    W : np.ndarray
        weighted connectivity matrix
    copy : bool

    Returns
    -------
    W : np.ndarray
        inverted connectivity matrix
    '''

    if copy:
        W = W.copy()
    E = np.where(W)
    W[E] = 1. / W[E]
    return W

In this function, W is a matrix. But I got the following error. It maybe related with W[E] = 1. / W[E] line.

File "/Users/xxx/anaconda3/lib/python3.7/site-packages/numba/dispatcher.py", line 317, in error_rewrite
    reraise(type(e), e, None)
  File "/Users/xxx/anaconda3/lib/python3.7/site-packages/numba/six.py", line 658, in reraise
    raise value.with_traceback(tb)
numba.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of Function(<built-in function getitem>) with argument(s) of type(s): (array(float64, 2d, A), tuple(array(int64, 1d, C) x 2))

So what is the right way to use NumPy and Numba? I know NumPy does a great job on matrix computing. In this case, is NumPy fast enough that Numba provides no more speed up?

numba numpy python
1 Answers

As FBruzzesi mentioned in comments, the reason the code isn't compiling is your use of "fancy indexing", since the E in W[E] is the output of np.where and is a tuple of arrays. (That explains the slightly cryptic error message: Numba doesn't know how to use getitem, i.e. it doesn't know how to find something in brackets, when one of the inputs is a tuple.)

Numba actually supports fancy indexing (also called "advanced indexing") on a single dimension, just not multiple dimensions. In your case, this allows for a simple modification: first using ravel to almost-costlessly make your array one-dimensional, then applying the transformation, then a cheap reshape back. 

@njit
def invert2(W, copy=True):
    if copy:
        W = W.copy()
    Z = W.ravel()
    E = np.where(Z)
    Z[E] = 1. / Z[E]
    return Z.reshape(W.shape)

But this is still slower than it needs to be, because passes the computation through unnecessary intermediate arrays rather than just immediately modifying the array when it encounters a nonzero value. It's faster simply to do a loop:

@njit 
def invert3(W, copy=True): 
    if copy: 
        W = W.copy() 
    Z = W.ravel() 
    for i in range(len(Z)): 
        if Z[i] != 0: 
            Z[i] = 1/Z[i] 
    return Z.reshape(W.shape)

This code works regardless of the dimensions of W. If we know that W is two-dimensional, then we could directly iterate over the two dimensions, but since the two have similar performance I'm going for the more general route.

On my computer, the timings, assuming a 300-by-300 array W where about half of the entries are 0s, and where invert is your original function without Numba compilation, are:

In [80]: %timeit invert(W)                                                                   
2.67 ms ± 49.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [81]: %timeit invert2(W)                                                                  
519 µs ± 24.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [82]: %timeit invert3(W)                                                                  
186 µs ± 11.1 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

so Numba gives us a pretty substantial speedup (after it's already been run once to eliminate compilation time), especially after the code is rewritten in a highly efficient loopy style that Numba can take advantage of.

user_11030440
Related