python中的快速向量化多项式

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

编辑： 很容易修改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

可能的解决办法 原则上，您可以使用

numba

实现，因为支持

多项式

分布

Numba允许您使用

Numba.njit

decorator简单地修饰numpy（甚至更重要的是标准Python函数），以获得显著的性能提升

更详细地了解这种方法。特别是

2.7.4

，因为它支持

np.random

（也支持多项式分布）

缺点：

大小

参数当前不受支持。您可以在嵌套循环中多次调用

np.random.multinomial

，但如果使用

numba.njit

进行修饰，它应该会更快

最后但并非最不重要的一点：您可以使用前面提到的decorator的

numba.prange

和

parallel

参数并行化外循环

性能测试 第一次测试：

• 具有签名类型的未并行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

Time elapsed: 1.0510437488555908
Time elapsed: 1.0691254138946533
Time elapsed: 1.065258264541626
Time elapsed: 1.0559568405151367
Time elapsed: 1.0446960926055908

完全没有麻木

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

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

可以看出，

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: 0.9460861682891846
Time elapsed: 0.9581060409545898
Time elapsed: 0.9654934406280518
Time elapsed: 0.9708254337310791
Time elapsed: 0.9757359027862549

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

部分结论 正如评论中正确指出的那样，我过早得出结论。并行化（如果可能的话）似乎对您的情况有最大的帮助，而

numba

（至少在这个仍然简单且不太全面的测试中）并没有带来很大的改进

---

总之，您应该检查一下您的具体情况，根据经验，您使用的Python代码越多，使用

numba

可能会得到更好的结果。如果它主要是基于numpy的，那么你不会看到任何好处（如果有的话）。

这里有一种方法可以做到这一点。它没有完全矢量化，但是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）

---

在一篇评论中，您说“p向量的形式是：p=[p_s，（1-p_s）/4，（1-p_s）/4，（1-p_s）/4，（1-p_s）/4]，p_s因站点而异。”鉴于包含

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，（1-p_s）/4，（1-p_s）/4]创建包含多项式概率的数组
：

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更改为适合您的问题的值）：

正如我所提到的，由于我可能需要更改每个站点的输入概率，因此在这里使用
size
参数实际上没有帮助，尽管它会影响代码的其他一些方面。我已经在考虑Numba或Theano是否适合我的需要，所以很高兴知道。添加了一个性能测试示例（希望与您的问题类似）。哦，在代码方面，它比Theano好得多（不再维护它，需要重写逻辑）。您可以检查与Python非常相似的
pytorch
，尽管
numba
似乎非常适合这种情况。您正在测量编译（大约0.6s）和运行时（低于10µs）。由于没有输出，Numba优化了您想要基准测试的代码部分。（因此10µs计时）1）决不要对任何编译代码中没有意义（例如没有输出）的函数进行基准测试。通常，编译器足够聪明（启用了一些优化标志），可以优化掉无用的部分。2） 如果你想得到运行时，不要测量第一次调用。你感兴趣的多项式分布的维数是多少？也就是说，就k而言，对于我的应用程序，每个阵列站点的k=5。p向量的形式是：p=[p_s，（1-p_s）/4，（1-p_s）/4，（1-p_s）/4，（1-p_s）/4]，不同站点的p_不同。我发现我需要在不同站点实现不同的计数，所以我在主要答案中发布了此结果的修改版本。我试着把它贴在这里，但是太长了，无法发表评论。