Быстрый случайный взвешенный выбор по всем строкам стохастической матрицы

numpy.random.choice позволяет взвешенный выбор из вектора, т. е.

 arr = numpy.array([1, 2, 3]) weights = numpy.array([0.2, 0.5, 0.3]) choice = numpy.random.choice(arr, p=weights) 

выбирает 1 с вероятностью 0,2, 2 с вероятностью 0,5 и 3 с вероятностью 0,3.

Что делать, если бы мы хотели сделать это быстро в векторном виде для двумерного массива (матрицы), для которого каждая из строк является вектором вероятностей? То есть, мы хотим вектор выбора из стохастической матрицы? Это супер медленный путь:

 import numpy as np m = 10 n = 100 # Or some very large number items = np.arange(m) prob_weights = np.random.rand(m, n) prob_matrix = prob_weights / prob_weights.sum(axis=0, keepdims=True) choices = np.zeros((n,)) # This is slow, because of the loop in Python for i in range(n): choices[i] = np.random.choice(items, p=prob_matrix[:,i]) 

print(choices) :

 array([ 4., 7., 8., 1., 0., 4., 3., 7., 1., 5., 7., 5., 3., 1., 9., 1., 1., 5., 9., 8., 2., 3., 2., 6., 4., 3., 8., 4., 1., 1., 4., 0., 1., 8., 5., 3., 9., 9., 6., 5., 4., 8., 4., 2., 4., 0., 3., 1., 2., 5., 9., 3., 9., 9., 7., 9., 3., 9., 4., 8., 8., 7., 6., 4., 6., 7., 9., 5., 0., 6., 1., 3., 3., 2., 4., 7., 0., 6., 3., 5., 8., 0., 8., 3., 4., 5., 2., 2., 1., 1., 9., 9., 4., 3., 3., 2., 8., 0., 6., 1.]) 

Это сообщение предполагает, что cumsum и bisect могут быть потенциальным подходом и быстро. Но в то время как numpy.cumsum(arr, axis=1) может делать это вдоль одной оси массива numpy, функция bisect.bisect работает только на одном массиве за раз. Аналогично, numpy.searchsorted работает только с 1D-массивами.

Есть ли быстрый способ сделать это, используя только векторизованные операции?

2 Solutions collect form web for “Быстрый случайный взвешенный выбор по всем строкам стохастической матрицы”

Вот полностью векторизованная версия, которая довольно быстро:

 def vectorized(prob_matrix, items): s = prob_matrix.cumsum(axis=0) r = np.random.rand(prob_matrix.shape[1]) k = (s < r).sum(axis=0) return items[k] 

Теоретически , searchsorted – это правильная функция, используемая для поиска случайного значения в кумулятивно суммируемых вероятностях, но с m относительно небольшим, k = (s < r).sum(axis=0) заканчивается намного быстрее. Его временная сложность – O (m), а метод searchsorted – O (log (m)), но это будет иметь значение только для гораздо большего m . Кроме того , cumsum – O (m), поэтому improved как vectorized , так и @ перимосокардии O (m). (Если ваш m , на самом деле, намного больше, вам нужно будет запустить некоторые тесты, чтобы увидеть, насколько большой m может быть до того, как этот метод будет медленнее.)

Вот время, которое я получаю с m = 10 и n = 10000 (используя original функции и improved из ответа @ perimosocordiae):

 In [115]: %timeit original(prob_matrix, items) 1 loops, best of 3: 270 ms per loop In [116]: %timeit improved(prob_matrix, items) 10 loops, best of 3: 24.9 ms per loop In [117]: %timeit vectorized(prob_matrix, items) 1000 loops, best of 3: 1 ms per loop 

Полный скрипт, в котором определены функции:

 import numpy as np def improved(prob_matrix, items): # transpose here for better data locality later cdf = np.cumsum(prob_matrix.T, axis=1) # random numbers are expensive, so we'll get all of them at once ridx = np.random.random(size=n) # the one loop we can't avoid, made as simple as possible idx = np.zeros(n, dtype=int) for i, r in enumerate(ridx): idx[i] = np.searchsorted(cdf[i], r) # fancy indexing all at once is faster than indexing in a loop return items[idx] def original(prob_matrix, items): choices = np.zeros((n,)) # This is slow, because of the loop in Python for i in range(n): choices[i] = np.random.choice(items, p=prob_matrix[:,i]) return choices def vectorized(prob_matrix, items): s = prob_matrix.cumsum(axis=0) r = np.random.rand(prob_matrix.shape[1]) k = (s < r).sum(axis=0) return items[k] m = 10 n = 10000 # Or some very large number items = np.arange(m) prob_weights = np.random.rand(m, n) prob_matrix = prob_weights / prob_weights.sum(axis=0, keepdims=True) 

Я не думаю, что можно полностью векторизовать это, но вы все равно сможете добиться приличного ускорения, векторизируя столько, сколько сможете. Вот что я придумал:

 def improved(prob_matrix, items): # transpose here for better data locality later cdf = np.cumsum(prob_matrix.T, axis=1) # random numbers are expensive, so we'll get all of them at once ridx = np.random.random(size=n) # the one loop we can't avoid, made as simple as possible idx = np.zeros(n, dtype=int) for i, r in enumerate(ridx): idx[i] = np.searchsorted(cdf[i], r) # fancy indexing all at once is faster than indexing in a loop return items[idx] 

Тестирование против версии в вопросе:

 def original(prob_matrix, items): choices = np.zeros((n,)) # This is slow, because of the loop in Python for i in range(n): choices[i] = np.random.choice(items, p=prob_matrix[:,i]) return choices 

Вот ускорение (с использованием кода установки, заданного в вопросе):

 In [45]: %timeit original(prob_matrix, items) 100 loops, best of 3: 2.86 ms per loop In [46]: %timeit improved(prob_matrix, items) The slowest run took 4.15 times longer than the fastest. This could mean that an intermediate result is being cached 10000 loops, best of 3: 157 µs per loop 

Я не уверен, почему существует большое расхождение в настройках моей версии, но даже самый медленный (~ 650 мкс) по-прежнему почти в 5 раз быстрее.

Python - лучший язык программирования в мире.