Random sample in numba

Viewed 2552

For performance reasons I often use numba and for my code I need to take a random sample without replacement. I found, that I could use the numpy.random function for that, but I noticed that it is extremely slow compared to the random.sample function. Am I doing something wrong? How could I improve the performance for the numba function? I boiled down my code to this minimal example:

import numpy as np
import numba as nb

def func2():
    List = range(100000)
    for x in range(20000):
        random.sample(List, 10)

@nb.njit()
def func3():
    Array = np.arange(100000)
    for x in range(20000):
        np.random.choice(Array, 10, False)

print(timeit(lambda: func2(), number=1))
print(timeit(lambda: func3(), number=1))
>>>0.1196
>>>20.1245

Edit: I'm now using my own sample function, which is much faster than np.random.choice.

@nb.njit()
def func4():
    for x in range(20000):
        rangeList = list(range(100000))
        result = []
        for x in range(10):
            randint = random.randint(0, len(rangeList) - 1)
            result.append(rangeList.pop(randint))
        return result
print(timeit(lambda: func4(), number=count))
>>>0.1767
1 Answers

Because I did some time measurements, I want to show you the results (regarding my comments on your question)

import numpy as np
from timeit import timeit
import numba as nb
import random

def func2():
    List = range(100000)
    for x in range(1000):
        random.sample(List, 10)

@nb.njit()
def func3():
    Array = np.arange(100000)
    for x in range(1000):
        np.random.choice(Array, 10, replace=False)

def func4():
    Array = np.arange(100000)
    for x in range(1000):
        np.random.choice(Array, 10, replace=False)

def func5():
    Array = np.arange(100000)
    for x in range(1000):
        np.random.default_rng().choice(Array, 10, replace=False)

print(f"random.sample {timeit(lambda: func2(), number=1)}")
print(f"np.random.choice JIT incl. compiling {timeit(lambda: func3(), number=1)}")
print(f"np.random.choice JIT excl. compiling {timeit(lambda: func3(), number=1)}")
print(f"np.random.choice {timeit(lambda: func4(), number=1)}")
print(f"np.random.default_rng.choice {timeit(lambda: func5(), number=1)}")

Giving you:

random.sample 0.0090606
np.random.choice JIT incl. compiling 1.9129443
np.random.choice JIT excl. compiling 0.8365084999999999
np.random.choice 1.8339632999999997
np.random.default_rng.choice 0.049018499999999854
Related