为什么numpy.where比替代品快得多

Question

im trying to speedup the following code:

import time
import numpy as np
np.random.seed(10)
b=np.random.rand(10000,1000)
def f(a=1):
    tott=0
    for _ in range(a):
        q=np.array(b)
        t1 = time.time()
        for i in range(len(q)):
            for j in range(len(q[0])):
                if q[i][j]>0.5:
                    q[i][j]=1
                else:
                    q[i][j]=-1
        t2=time.time()
        tott+=t2-t1
    print(tott/a)

As you can see, mainly func is about iterating in double cycle. So, i've tried to use np.nditer , np.vectorize and map instead of it. If gives some speedup (like 4-5 times except np.nditer ), but! with np.where(q>0.5,1,-1) speedup is almost 100x. How can i iterate over numpy arrays as fast as np.where does it? And why is it so much faster?

Answer 1

It's because the core of numpy is implemented in C. You're basically comparing the speed of C with Python.

If you want to use the speed advantage of numpy, you should make as few calls as possible in your Python code. If you use a Python-loop, you have already lost, even if you use numpy functions in that loop only. Use higher-level functions provided by numpy (that's why they ship so many special functions). Internally, it will use a much more efficient (C-)loop

You can implement a function in C (with loops) yourself and call that from Python. That should give comparable speeds.

Answer 2

To answer this question, you can gain the same speed (100x acceleration) by using the numba library:

from numba import njit

def f(b):
    q = np.zeros_like(b)

    for i in range(b.shape[0]):
        for j in range(b.shape[1]):
            if q[i][j] > 0.5:
                q[i][j] = 1
            else:
                q[i][j] = -1

    return q

@njit
def f_jit(b):
    q = np.zeros_like(b)

    for i in range(b.shape[0]):
        for j in range(b.shape[1]):
            if q[i][j] > 0.5:
                q[i][j] = 1
            else:
                q[i][j] = -1

    return q

Compare the speed:

Plain Python

%timeit f(b)
592 ms ± 5.72 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Numba (just-in-time compiled using LLVM ~ C speed)

%timeit f_jit(b)
5.97 ms ± 105 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)