Python中的快速向量化多项式

Question

我目前正在使用NumPy执行以下任务：我有一个很大的值网格，并且需要在每个点取一个多项式样本。多项式的概率向量在网格位置之间会有所不同，因此NumPy多项式函数对我而言并不起作用，因为它从相同的分布进行所有绘制。遍历每个站点似乎效率极低，我想知道NumPy中是否有一种有效的方法。如果使用Theano（请参阅this answer），看起来这样的事情是可能的（并且很快），但是这需要大量的重写，我希望避免这种情况。可以在基本的NumPy中有效地对多项采样进行矢量化吗？

编辑： 正如我发现我所需要的那样，很容易修改Warren的代码以允许在不同的站点进行不同的计数：所有需要做的就是传递完整的count数组并删除第一个数组，如下所示： >

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参数并行化。

性能测试

第一次测试：

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

结果：

具有类型签名的无并行数词

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

可以看到，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循环位于syslog值之上。如果p个向量的长度不太大，那么对于您来说这可能足够快。

多项式分布是通过重复调用p来实现的，该实现实现了其参数的广播。

np.random.binomial

在此示例中，“网格”的形状为（2，3），多项式分布为四维（即每个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 向量的长度为4）。

p

---

在评论中，您说“ p向量的形式为：p = [p_s，（1-p_s）/ 4，（1-p_s）/ 4，（1-p_s）/ 4，（1- p_s）/ 4]，而p_s随站点而异。”给定一个包含In [182]: p = np.array([[[0.25, 0.25, 0.25, 0.25], ...: [0.01, 0.02, 0.03, 0.94], ...: [0.75, 0.15, 0.05, 0.05]], ...: [[0.01, 0.99, 0.00, 0.00], ...: [1.00, 0.00, 0.00, 0.00], ...: [0.00, 0.25, 0.25, 0.50]]]) In [183]: sample = multinomial_rvs(1000, p) In [184]: sample Out[184]: array([[[ 249, 260, 233, 258], [ 3, 21, 33, 943], [ 766, 131, 55, 48]], [[ 5, 995, 0, 0], [1000, 0, 0, 0], [ 0, 273, 243, 484]]]) In [185]: sample.sum(axis=-1) Out[185]: array([[1000, 1000, 1000], [1000, 1000, 1000]]) 值的数组，这就是使用上述函数的方法。

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

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 = [p_s, (1-p_s)/4, (1-p_s)/4, (1-p_s)/4, (1-p_s)/4]

现在与生成样本相同（将值1000更改为适合您的问题的值）：

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 ]]])