python中的快速向量多项式

Question

我目前正在使用NumPy来完成以下任务：我有一个大的值网格，我需要在每个点上取一个多项式样本。多项式的概率向量会随着网格位置的不同而变化，所以NumPy多项式函数对我来说不太适用，因为它的所有绘制都来自同一个分布。对每个站点进行迭代似乎效率极低，我想知道是否有一种方法可以在NumPy中高效地完成这项工作。如果我使用Theano（参见this answer），这样的事情似乎是可能的（而且很快），但这需要大量重写，这是我理想中希望避免的。多项式抽样能有效地在基元中进行矢量化吗？在

编辑： 修改Warren的代码以允许不同站点的不同计数是很容易的，因为我发现我需要做的就是传入完整的count数组并删除第一个like，如下所示：

import numpy as np


def multinomial_rvs(count, p):
    """
    Sample from the multinomial distribution with multiple p vectors.

    * count must be an (n-1)-dimensional numpy array.
    * p must an n-dimensional numpy array, n >= 1.  The last axis of p
      holds the sequence of probabilities for a multinomial distribution.

    The return value has the same shape as p.
    """
    out = np.zeros(p.shape, dtype=int)
    ps = p.cumsum(axis=-1)
    # Conditional probabilities
    with np.errstate(divide='ignore', invalid='ignore'):
        condp = p / ps
    condp[np.isnan(condp)] = 0.0
    for i in range(p.shape[-1]-1, 0, -1):
        binsample = np.random.binomial(count, condp[..., i])
        out[..., i] = binsample
        count -= binsample
    out[..., 0] = count
    return out

Answer 1

可能的解决方案

原则上，您可以使用numba，因为multinomial分布是受支持的。在

Numba允许您简单地用numba.njit装饰器装饰numpy（更重要的是标准Python函数）来显著提高性能。在

Check their documentation以更详细地查看此方法。尤其是2.7.4，因为它是关于支持np.random（也支持多项式分布）。在

下行：当前不支持size参数。您可以在嵌套循环中多次调用np.random.multinomial，如果用numba.njit修饰，它应该会更快。在

最后但并非最不重要的一点：您可以使用numba.prange和parallel参数将外部循环并行化到前面提到的decorator。在

性能试验

第一次试验：

• 带类型签名的未调用numba
• 一点都不麻木

测试规范：

import sys
from functools import wraps
from time import time

import numba
import numpy as np


def timing(function):
    @wraps(function)
    def wrap(*args, **kwargs):
        start = time()
        result = function(*args, **kwargs)
        end = time()
        print(f"Time elapsed: {end - start}", file=sys.stderr)
        return result

    return wrap


@timing
@numba.njit(numba.int64(numba.int64[:, :], numba.int64))
def my_multinomial(probabilities, output):
    experiments: int = 5000
    output_array = []
    for i in numba.prange(probabilities.shape[0]):
        probability = probabilities[i] / np.sum(probabilities[i])
        result = np.random.multinomial(experiments, pvals=probability)
        if i % output == 0:
            output_array.append(result)

    return output_array[-1][-1]


if __name__ == "__main__":
    np.random.seed(0)
    probabilities = np.random.randint(low=1, high=100, size=(10000, 1000))
    for _ in range(5):
        output = my_multinomial(probabilities, np.random.randint(low=3000, high=10000))

结果：

带类型签名的未调用numba

^{pr2}$

一点都不麻木

Time elapsed: 0.9460861682891846
Time elapsed: 0.9581060409545898
Time elapsed: 0.9654934406280518
Time elapsed: 0.9708254337310791
Time elapsed: 0.9757359027862549

可以看出，numba在这种情况下没有任何帮助（实际上它会降低性能）。对于不同大小的输入阵列，结果是一致的。在

第二次试验

• 无类型签名的并行numba
• 一点都不麻木

测试代码：

import sys
from functools import wraps
from time import time

import numba
import numpy as np


def timing(function):
    @wraps(function)
    def wrap(*args, **kwargs):
        start = time()
        result = function(*args, **kwargs)
        end = time()
        print(f"Time elapsed: {end - start}", file=sys.stderr)
        return result

    return wrap


@timing
@numba.njit(parallel=True)
def my_multinomial(probabilities, output):
    experiments: int = 5000
    for i in range(probabilities.shape[0]):
        probability = probabilities[i] / np.sum(probabilities[i])
        result = np.random.multinomial(experiments, pvals=probability)
        if i % output == 0:
            print(result)


if __name__ == "__main__":
    np.random.seed(0)
    probabilities = np.random.randint(low=1, high=100, size=(10000, 1000))
    for _ in range(5):
        my_multinomial(probabilities, np.random.randint(low=3000, high=10000))

结果：

无类型签名的并行numba：

Time elapsed: 1.0705969333648682                                                                                                                                                          
Time elapsed: 0.18749785423278809                                                                                                                                                         
Time elapsed: 0.1877145767211914                                                                                                                                                          
Time elapsed: 0.18813610076904297                                                                                                                                                         
Time elapsed: 0.18747472763061523 

一点都不麻木

Time elapsed: 1.0142333507537842                                                                                                                                                          
Time elapsed: 1.0311956405639648                                                                                                                                                          
Time elapsed: 1.022024154663086                                                                                                                                                           
Time elapsed: 1.0191617012023926                                                                                                                                                          
Time elapsed: 1.0144879817962646

部分结论

正如max9111在评论中正确指出的那样，我过早地得出结论。似乎并行化（如果可能的话）对您的情况是最大的帮助，而numba（至少在这个仍然简单且不太全面的测试中）并没有带来很大的改进。在

总之，您应该检查一下您的确切情况，根据经验，您使用的Python代码越多，使用numba可能会得到更好的结果。如果它主要是基于numpy的，那么你不会看到任何好处（如果有的话）。在

Answer 2

这里有一个方法你可以做到。它不是完全矢量化的，但是Python循环在p值之上。如果你的p向量的长度不是太大，这可能对你来说足够快了。在

多项式分布是通过重复调用np.random.binomial来实现的，它实现了参数的广播。在

import numpy as np


def multinomial_rvs(n, p):
    """
    Sample from the multinomial distribution with multiple p vectors.

    * n must be a scalar.
    * p must an n-dimensional numpy array, n >= 1.  The last axis of p
      holds the sequence of probabilities for a multinomial distribution.

    The return value has the same shape as p.
    """
    count = np.full(p.shape[:-1], n)
    out = np.zeros(p.shape, dtype=int)
    ps = p.cumsum(axis=-1)
    # Conditional probabilities
    with np.errstate(divide='ignore', invalid='ignore'):
        condp = p / ps
    condp[np.isnan(condp)] = 0.0
    for i in range(p.shape[-1]-1, 0, -1):
        binsample = np.random.binomial(count, condp[..., i])
        out[..., i] = binsample
        count -= binsample
    out[..., 0] = count
    return out

下面是一个例子，其中“网格”具有形状（2，3），多项式分布是四维的（即每个p向量的长度为4）。在

^{pr2}$

---

在一篇评论中，你说“p向量的形式是：p=[p_s，（1-p_s）/4，（1-p_s）/4，（1-p_s）/4，（1-p_s）/4]，每个站点的p_都不同。”下面是如何使用上面的函数，给定一个包含p_s值的数组。在

首先为示例创建一些数据：

In [73]: p_s = np.random.beta(4, 2, size=(2, 3))                                                                                                        

In [74]: p_s                                                                                                                                            
Out[74]: 
array([[0.61662208, 0.6072323 , 0.62208711],
       [0.86848938, 0.58959038, 0.47565799]])

根据公式p = [p_s, (1-p_s)/4, (1-p_s)/4, (1-p_s)/4, (1-p_s)/4]创建包含多项式概率的数组p：

In [75]: p = np.expand_dims(p_s, -1) * np.array([1, -0.25, -0.25, -0.25, -0.25]) + np.array([0, 0.25, 0.25, 0.25, 0.25])                                

In [76]: p                                                                                                                                              
Out[76]: 
array([[[0.61662208, 0.09584448, 0.09584448, 0.09584448, 0.09584448],
        [0.6072323 , 0.09819192, 0.09819192, 0.09819192, 0.09819192],
        [0.62208711, 0.09447822, 0.09447822, 0.09447822, 0.09447822]],

       [[0.86848938, 0.03287765, 0.03287765, 0.03287765, 0.03287765],
        [0.58959038, 0.1026024 , 0.1026024 , 0.1026024 , 0.1026024 ],
        [0.47565799, 0.1310855 , 0.1310855 , 0.1310855 , 0.1310855 ]]])

现在执行与之前相同的操作以生成样本（将值1000更改为适合您的问题的任何值）：

In [77]: sample = multinomial_rvs(1000, p)                                                                                                              

In [78]: sample                                                                                                                                         
Out[78]: 
array([[[618,  92, 112,  88,  90],
        [597, 104, 103, 101,  95],
        [605, 100,  95,  98, 102]],

       [[863,  32,  43,  27,  35],
        [602, 107, 108,  94,  89],
        [489, 130, 129, 129, 123]]])

In [79]: sample.sum(axis=-1)                                                                                                                            
Out[79]: 
array([[1000, 1000, 1000],
       [1000, 1000, 1000]])