对随机矩阵的所有行进行快速随机加权选择

Question

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) 

以概率 0.2 选择 1，以概率 0.5 选择 2，以概率 0.3 选择 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.])

This post 建议 cumsum 和 bisect 可能是一种潜在的方法，而且速度很快。但是，虽然 numpy.cumsum(arr, axis=1) 可以沿着 numpy 数组的一个轴执行此操作，但 bisect.bisect 函数一次只能在一个数组上工作。同样，numpy.searchsorted 也仅适用于一维数组。

有没有一种仅使用矢量化操作的快速方法？

Answer 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 得到的时间（使用@perimosocordiae 的答案中的original 和improved 函数）：

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)

Answer 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 µs) 仍然快近 5 倍。