如何提高np.random.choice（）循环效率

Question

I am trying to apply np.random.choice to a big array with different weights, and wondering any way could avoid looping and improve the performance? Over here len(weights) could be millions.

weights = [[0.1, 0.5, 0.4],
           [0.2, 0.4, 0.4],
           ...
           [0.3, 0.3, 0.4]]

choice = [1, 2, 3]
ret = np.zeros((len(weights), 20))
for i in range(len(weights)):
    ret[i] = np.random.choice(choice, 20, p=weights[i])

Answer 1

Here's a generalization of my answer in Fast random weighted selection across all rows of a stochastic matrix :

def vectorized_choice(p, n, items=None):
    s = p.cumsum(axis=1)
    r = np.random.rand(p.shape[0], n, 1)
    q = np.expand_dims(s, 1) >= r
    k = q.argmax(axis=-1)
    if items is not None:
        k = np.asarray(items)[k]
    return k

p is expected to be a two-dimensional array whose rows are probability vectors. n is the number of samples to draw from the distribution defined by each row. If items is None, the samples are integers in range(0, p.shape[1]) . If items is not None, it is expected to be a sequence with length p.shape[1] .

Example:

In [258]: p = np.array([[0.1, 0.5, 0.4], [0.75, 0, 0.25], [0, 0, 1], [1/3, 1/3, 1/3]])                                                   

In [259]: p                                                                                                                              
Out[259]: 
array([[0.1       , 0.5       , 0.4       ],
       [0.75      , 0.        , 0.25      ],
       [0.        , 0.        , 1.        ],
       [0.33333333, 0.33333333, 0.33333333]])

In [260]: vectorized_choice(p, 20)                                                                                                       
Out[260]: 
array([[1, 1, 2, 1, 1, 2, 2, 2, 1, 2, 1, 1, 1, 2, 2, 0, 1, 2, 2, 2],
       [0, 2, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0],
       [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
       [1, 0, 2, 2, 0, 1, 2, 1, 0, 0, 0, 0, 2, 2, 0, 0, 2, 1, 1, 2]])

In [261]: vectorized_choice(p, 20, items=[1, 2, 3])                                                                                      
Out[261]: 
array([[2, 1, 2, 2, 2, 3, 2, 2, 2, 2, 3, 3, 2, 2, 3, 3, 2, 3, 2, 2],
       [1, 1, 1, 1, 1, 1, 1, 1, 3, 3, 3, 1, 1, 3, 3, 1, 3, 1, 1, 1],
       [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3],
       [3, 3, 3, 1, 3, 2, 1, 2, 3, 1, 2, 2, 3, 2, 1, 2, 1, 2, 2, 2]])

Timing for p with shape (1000000, 3) :

In [317]: p = np.random.rand(1000000, 3)

In [318]: p /= p.sum(axis=1, keepdims=True)

In [319]: %timeit vectorized_choice(p, 20, items=np.arange(1, p.shape[1]+1))
1.89 s ± 28.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Here's the timing for Divakar's function:

In [320]: %timeit random_choice_prob_vectorized(p, 20, choice=np.arange(1, p.shape[1]+1))
7.33 s ± 43.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

The difference will be less pronounced if you increase the number of columns in p , and if you make the number of columns big enough, Divakar's function will be faster. Eg

In [321]: p = np.random.rand(1000, 120)

In [322]: p /= p.sum(axis=1, keepdims=True)

In [323]: %timeit vectorized_choice(p, 20, items=np.arange(1, p.shape[1]+1))
6.41 ms ± 20.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [324]: %timeit random_choice_prob_vectorized(p, 20, choice=np.arange(1, p.shape[1]+1))
6.29 ms ± 342 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Answer 2

Borrowing idea from Vectorizing numpy.random.choice for given 2D array of probabilities along an axis alongwith idea from vectorized searchsorted , here's one vectorized way -

def random_choice_prob_vectorized(weights, num_items, choice=None):
    weights = np.asarray(weights)

    w = weights.cumsum(1)
    r = np.random.rand(len(weights),num_items)

    m,n = w.shape
    o = np.arange(m)[:,None]
    w_o = (w+o).ravel()
    r_o = (r+o).ravel()
    idx = np.searchsorted(w_o,r_o).reshape(m,-1)%n
    if choice is not None:
        return np.asarray(choice)[idx]
    else:
        return idx

Sample run to verify using 2D bincount -

In [28]: weights = [[0.1, 0.5, 0.4],
    ...:            [0.2, 0.4, 0.4],
    ...:            [0.3, 0.3, 0.4]]
    ...: 
    ...: choice = [1, 2, 3]
    ...: num_items = 20000

In [29]: out = random_choice_prob_vectorized(weights, num_items, choice)

# Use 2D bincount to get per average occurences and verify against weights
In [75]: bincount2D_vectorized(out)/num_items
Out[75]: 
array([[0.     , 0.09715, 0.4988 , 0.40405],
       [0.     , 0.1983 , 0.40235, 0.39935],
       [0.     , 0.30025, 0.29485, 0.4049 ]])

Answer 3

Looks like each row of the resulting array is independent of other rows. I am not sure how bad is the performance now. If it really is a concern, I would try to use python's multiprocessing module to run random number generations with several processes in parallel. It should help.