在 python 中生成截断的负二项分布

Question

我正在尝试按照由数字组成的截断负二项分布生成数据集，以使数字集具有最大值。

def truncated_Nbinom(n, p, max_value, size):
    import scipy.stats as sct
    temp_size = size
    while True:
        temp_size *= 2
        temp = sct.nbinom.rvs(n, p, size=temp_size)
        truncated = temp[temp <= max_value]
        if len(truncated) >= size:
            return truncated[:size]

当 max_value 和 n 较小时，我能够得到结果。 但是，当我尝试：

input_1= truncated_Nbinom(99, 0.3, 99, 5000).tolist()

kernel 不断死亡。 我尝试更改 python 的端口并提高递归限制，但它们不起作用。 你有什么想法可以让我的代码更快吗？

Answer 1

这是一种方法。 您可以计算在负二项式下选择x的概率，然后将x s 低于max_value的概率归一化为总和为 1。 现在，您可以简单地以适当的概率调用np.random.choice 。

import numpy as np
import pandas as pd
from scipy import stats


def truncated_Nbinom2(n, p, max_value, size):
  support = np.arange(max_value + 1)
  probs = stats.nbinom.pmf(support, n, p)
  probs /= probs.sum()
  return np.random.choice(support, size=size, p=probs)

这是一个插图：

arr1 = truncated_Nbinom(9, 0.3, 9, 50000)
arr2 = truncated_Nbinom2(9, 0.3, 9, 50000)

df_counts = pd.DataFrame({
    "version_1": pd.Series(arr1).value_counts(),
    "version_2": pd.Series(arr2).value_counts(),
})