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

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.

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

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-массивами.

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

Ответ 1

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

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), поэтому оба vectorized и @perimosocordiae improved равны 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

Полный script, где определены функции:

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)

Ответ 2

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

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 раз быстрее.