Why is numpy.random.choice so slow?

Question

While writing a script I discovered the numpy.random.choice function. I implemented it because it was was much cleaner than the equivalent if statement. However, after running the script I realized it is significantly slower than the if statement.

The following is a MWE. The first method takes 0.0 s, while the second takes 7.2 s. If you scale up the i loop, you will see how fast random.choice slows down.

Can anyone comment on why random.choice is so much slower?

import numpy as np
import numpy.random as rand
import time as tm

#-------------------------------------------------------------------------------

tStart = tm.time()
for i in xrange(100):
    for j in xrange(1000):
        tmp = rand.rand()
        if tmp < 0.25:
            var = 1
        elif tmp < 0.5:
            var = -1
print('Time: %.1f s' %(tm.time() - tStart))

#-------------------------------------------------------------------------------

tStart = tm.time()
for i in xrange(100):
    for j in xrange(1000):
        var = rand.choice([-1, 0, 1], p = [0.25, 0.5, 0.25])
print('Time: %.1f s' %(tm.time() - tStart))

Answer 1

You're using it wrong. Vectorize the operation, or numpy will offer no benefit:

var = numpy.random.choice([-1, 0, 1], size=1000, p=[0.25, 0.5, 0.25])

Timing data:

>>> timeit.timeit('''numpy.random.choice([-1, 0, 1],
...                                      size=1000,
...                                      p=[0.25, 0.5, 0.25])''',
...               'import numpy', number=10000)
2.380380242513752

>>> timeit.timeit('''
... var = []
... for i in xrange(1000):
...     tmp = rand.rand()
...     if tmp < 0.25:
...         var.append(1)
...     elif tmp < 0.5:
...         var.append(-1)
...     else:
...         var.append(0)''',
... setup='import numpy.random as rand', number=10000)
5.673041396894519

Answer 2

This solution with a cumulative score is about 25x faster:

def choice(options,probs):
    x = np.random.rand()
    cum = 0
    for i,p in enumerate(probs):
        cum += p
        if x < cum:
            break
    return options[i]


options = ['a','b','c','d']
probs = [0.2,0.6,0.15,0.05]
runs = 100000


now = time.time()
temp = []
for i in range(runs):
    op = choice(options,probs)
    temp.append(op)
temp = Counter(temp)
for op,x in temp.items():
    print(op,x/runs)
print(time.time()-now)

print("")
now = time.time()
temp = []
for i in range(runs):
    op = np.random.choice(options,p = probs)
    temp.append(op)
temp = Counter(temp)
for op,x in temp.items():
    print(op,x/runs)
print(time.time()-now)

Running it I get:

b 0.59891
a 0.20121
c 0.15007
d 0.04981
0.16232800483703613

b 0.5996
a 0.20138
c 0.14856
d 0.05046
3.8451428413391113

Answer 3

It took me a really long time to figure out that my data generator is very slow due to the random key sampling via np.random.choice.

In case the non-uniform distribution is NOT needed, then here is the workable solution I found.

Replace

def get_random_key(a_huge_key_list):
    return np.random.choice(a_huge_key_list)

with

def get_random_key(a_huge_key_list):
    L = len(a_huge_key_list)
    i = np.random.randint(0, L)
    return a_huge_key_list[i]

which gives a speedup of x60.

Answer 4

I suspect the generality of np.random.choice is slowing it down, more so for small samples than large ones.

A crude vectorization of the if version is:

def foo(n):
    x = np.random.rand(n)
    var = np.zeros(n)
    var[x<.25] = -1
    var[x>.75] = 1
    return var

Running in ipython I get:

timeit np.random.choice([-1,0,1],size=1000,p=[.25,.5,.25])
1000 loops, best of 3: 293 us per loop

timeit foo(1000)
10000 loops, best of 3: 83.4 us per loop

timeit np.random.choice([-1,0,1],size=100000,p=[.25,.5,.25])
100 loops, best of 3: 11 ms per loop

timeit foo(100000)
100 loops, best of 3: 8.12 ms per loop

So for the 1000 size, choice is 3-4x slower, but with larger vectors, the difference starts to disappear.

Answer 5

For others who stumble upon this question and do not draw 1,000 samples 10,000 times but 1 sample 10,000 times, there exists a faster alternative since Python 3.6. The function random.choices is ~20 times faster compared to numpy.random.choice.

timeit("random.choices([-1, 0, 1], k=1, weights=[.25, .5, .25])", 
setup="import random", number=10000)
# >>> 0.018841781999981322

vs

timeit("numpy.random.choice([-1, 0, 1], size=1, p=[.25, .5, .25])", 
setup="import numpy", number=10000)
# >>> 0.40612822200000664

Answer 6

Other answers involve at least one of the following:

1- Using python list as an input for numpy.random.choice and creating overhead.

2- Using prior knowledge that the len(array) will be 3.

3- The distribution is uniform.

For arbitrary length lists, one of the fastest algorithms would divide the list into 2 at each step. For example the following code would work in generic cases.

def my_random_function(collection, p):
    miles = []
    current = 0
    for prob in p:
        miles.append(current)
        current += prob
    if not math.isclose(current,1):
        raise ValueError()
    x = random.random()
    _all = list(zip(collection,miles))
    while(len(_all)!= 1):
        if _all[len(_all)//2][1] < x:
            _all = _all[len(_all)//2:]
        else:
            _all = _all[0: len(_all)//2]
    return _all[0][0]

To compare the differences I prepared two cases:

small_list = list(range(3))
small_array = np.arange(3)
#create a random probability list
small_p = [random.random() for i in range(3)]
small_p = [prob/sum(small_p) for prob in small_p]
small_p_np = np.array(small_p)

large_list = list(range(10000))
large_array = np.arange(10000)
#create a random probability list
large_p = [random.random() for i in range(10000)]
large_p = [prob/sum(large_p) for prob in large_p]
large_p_np = np.array(large_p)

The results are as follows:

%timeit np.random.choice(small_array, p= small_p_np)

68.1 µs ± 196 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit my_random_function(small_list, small_p)

5.13 µs ± 26.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit np.random.choice(large_array, p= large_p_np)

279 µs ± 1.15 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit my_random_function(large_list, large_p)

3.26 ms ± 5.82 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

As the results show numpy.random.choice takes more than x10 time for small collections, but quickly becomes the better choice when there are more elements. It seems that this function has a big overhead, and better be avoided for small lists in performance critical parts of the code.