python 中的快速向量化多项式
Fast vectorized multinomial in python
我目前正在使用 NumPy 来完成以下任务：我有一个很大的值网格，我需要在每个点取一个多项式样本。多项式的概率向量会因网格站点而异，因此 NumPy 多项式函数对我来说不太适用，因为它的所有绘制都来自同一分布。遍历每个站点似乎非常低效，我想知道是否有一种方法可以在 NumPy 中有效地执行此操作。如果我使用 Theano（请参阅 ），这样的事情似乎是可能的（而且速度很快），但这将需要相当多的重写，我希望避免这种情况。可以在基本 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
可能的解决方案
原则上你可以使用 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
- 根本没有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
完全没有numba
Time elapsed: 0.9460861682891846
Time elapsed: 0.9581060409545898
Time elapsed: 0.9654934406280518
Time elapsed: 0.9708254337310791
Time elapsed: 0.9757359027862549
正如你所看到的，numba 在这种情况下没有帮助（实际上它降低了性能）。对于不同大小的输入数组，结果是一致的。
第二次测试
- 没有类型签名的并行 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
完全没有numba
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，那么您将看不到好处（如果有的话）。
这是您可以做到的一种方法。它没有完全矢量化，但是 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
这是一个示例，其中 "grid" 的形状为 (2, 3)，多项式分布为四维（即每个 p 向量的长度为 4）。
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]])
在评论中，您说 "The p vector is of the form: p = [p_s, (1-p_s)/4, (1-p_s)/4, (1-p_s)/4, (1-p_s)/4], with p_s varying from site to site." 给定一个包含 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]])
我目前正在使用 NumPy 来完成以下任务：我有一个很大的值网格，我需要在每个点取一个多项式样本。多项式的概率向量会因网格站点而异，因此 NumPy 多项式函数对我来说不太适用，因为它的所有绘制都来自同一分布。遍历每个站点似乎非常低效，我想知道是否有一种方法可以在 NumPy 中有效地执行此操作。如果我使用 Theano（请参阅
编辑：
修改 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
可能的解决方案
原则上你可以使用 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
- 根本没有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
完全没有numba
Time elapsed: 0.9460861682891846
Time elapsed: 0.9581060409545898
Time elapsed: 0.9654934406280518
Time elapsed: 0.9708254337310791
Time elapsed: 0.9757359027862549
正如你所看到的，numba 在这种情况下没有帮助（实际上它降低了性能）。对于不同大小的输入数组，结果是一致的。
第二次测试
- 没有类型签名的并行 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
完全没有numba
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，那么您将看不到好处（如果有的话）。
这是您可以做到的一种方法。它没有完全矢量化，但是 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
这是一个示例，其中 "grid" 的形状为 (2, 3)，多项式分布为四维（即每个 p 向量的长度为 4）。
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]])
在评论中，您说 "The p vector is of the form: p = [p_s, (1-p_s)/4, (1-p_s)/4, (1-p_s)/4, (1-p_s)/4], with p_s varying from site to site." 给定一个包含 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]])