Search code examples
pythonnumpyvectorization

Numpy: how to vectorize a for loop generating random numbers?

Say I have the following simple function that I use to generate random numbers:

def my_func():

  rvs = np.random.random(size=3)
  
  return rvs[2] - rvs[1]

I want to call this function a number of times, lets say 1000 and I want to store the results in an array, for example:

result = []
for _ in range(1000):
   result += [my_func()]

Is there a way to use numpy to vectorize this operation and make everything faster? I don't mind if the workflow changes.

Solution

  • If I understand your question correctly, you just need to use the np.random.rand function:

    np.random.rand(1000)

    This function create an array of the given shape and populate it with random samples from a uniform distribution over [0, 1).

    You can vectorize as follows:

    rvs_vect = np.random.rand(10000, 3)
result = rvs_vect[:,2] - rvs_vect[:,1]

    rvs_vect[:,1] selects all rows in column 1.
    rvs_vect[:,2] selects all rows in column 2.

    Execution times for instances of 10000 elements on my machine are about 100 times faster than your solution and the other proposed ones (np.vectorize and list comprehension).

    Extras

    I have prepared an example for you with Numba. Numba is an open source JIT compiler that translates a subset of Python and NumPy code into fast machine code. Although you will not gain substantial advantages over numpy on this type of operation.

    import numba as nb
nb.njit 
def my_rand(n):
    rvs_vect = np.random.rand(n, 3)
    return rvs_vect[:,2] - rvs_vect[:,1]