如何在 Python/numpy 中创建嵌套加权分布？

Question

我正在尝试优化（矢量化？）蒙特卡洛式模拟的创建，但我无法弄清楚如何使用 numpy 或类似库创建嵌套加权随机值。考虑以下场景，受 interviewqs 问题的启发：“三个 class 教室的学生必须投票给两个 class 总统候选人中的一个。A 教室有 40% 的学生并且被分开候选人 X 和 Y 50/50。B 有 25% 的学生，分成 60/40。C 有 35% 的学生，分成 35/65

使用 vanilla Python 创建数据可能看起来像这样，

import random

nsimulations = 10_000_000

def choose_classroom():
    "returns A, B, or C based on percentages"
    value = random.random()
    if value < .4:
        return 'A'
    elif value < .65:
        return 'B'
    else:
        return 'C'
        
def choose_support(classroom):
    "return X or Y based on support percentage by classroom"
    value = random.random()
    if classroom == 'A':
        return "X" if value < .5 else "Y"
    elif classroom == 'B':
        return "X" if value < .6 else "Y"
    elif classroom == 'C':
        return "X" if value < .35 else "Y"
        
results = []
for i in range(nsimulations):
    classroom = choose_classroom()
    support = choose_support(classroom)
    results.append({'classroom': classroom, 'support': support})

在我的机器上运行 10M 模拟大约需要 10 秒。我想改善那个时间。我首先看到的是numpy.random.choice、fast_classrooms = np.random.choice(['A', 'B', 'C'], size=nsimulations, p=[.4, .25, .35])。这确实执行得很快，大约 350 毫秒。但是我不知道如何从那里生成后续 X/Y 分布。

我试过的一件事是 Pandas apply，它似乎在暗中进行了某种优化。下面的 Pandas 片段需要 ~2.5 秒到运行而列表理解（[choose_support(record) for record in fast_classrooms] 需要 ~4 秒。不过，感觉这不是“正确”的做事方式.

import pandas
import numpy as np

fast_classrooms = np.random.choice(['A', 'B', 'C'], size=nsimulations, p=[.4, .25, .35])
df = pandas.DataFrame({'classroom': fast_classrooms})
df['support'] = df.classroom.apply(choose_support)

我希望找到的是可以生成嵌套加权分布的东西，比如 - np.random.choice([['A', 'B', 'C'], ['X', 'Y']], p=[[.4, .25, .35], [[.5, .5], [.6, .4], [.35, .65]]])

生成这些数据的方法有哪些？

Answer 1

您可以显着减少 python 循环次数，使代码更加向量化：

import numpy as np

nsimulations = 12
uniquerooms = ['A','B','C']
supportprobs = [0.5, 0.6, 0.35]
classrooms = np.random.choice(uniquerooms, size=nsimulations, p=[.4, .25, .35])
supports = np.empty_like(classrooms, dtype=classrooms.dtype)
for room, prob in zip(uniquerooms, supportprobs):
    mask = classrooms == room
    supports[mask] = np.random.choice(['X','Y'], size=mask.sum(), p=[prob, 1-prob])

print(classrooms)
# ['C' 'A' 'C' 'A' 'C' 'C' 'A' 'C' 'B' 'C' 'B' 'A']
print(supports)
# ['X' 'Y' 'Y' 'Y' 'Y' 'X' 'Y' 'X' 'X' 'X' 'Y' 'X']

Answer 2

您可以对对使用 np.random.choice 而不是运行函数两次。这意味着您可以计算出一对 ('classroom', 'support') 的概率。例如，选择课堂 'A' 和获得支持 'X' 的概率是 0.4*0.5 = 0.2，依此类推。

下面的代码对我来说运行速度非常快。

import numpy as np
nsimulations = 10000000

#Construct the probabilities and pairs 
p = [.4*.5, .4*.5, .25*.6, .25*.4, .35*.35, .35*.65]
pairs = [{'classroom': 'A', 'support': 'X'}, 
         {'classroom': 'A', 'support': 'Y'},
         {'classroom': 'B', 'support': 'X'},
         {'classroom': 'B', 'support': 'Y'},
         {'classroom': 'C', 'support': 'X'},
         {'classroom': 'C', 'support': 'Y'}]

# Sample the pairs based on the probabilities
fast_classrooms = np.random.choice(pairs, size=nsimulations, p=p)

编辑：

与 OP 用 7.815439224243164 seconds 发布的方法相比，此代码用了 0.6193864345550537 seconds。 @Tom McLean 在评论中也证实了这一点。

Answer 3

我认为当前的最佳答案是最优雅的，但由于条件的可扩展性，我想将 numpy.piecewise 加入其中：

import numpy as np 

fast_classrooms = np.random.choice(['A', 'B', 'C'], size=nsimulations, p=[.4, .25, .35])

np.piecewise(fast_classrooms, [fast_classrooms == 'A', fast_classrooms =='B', fast_classrooms=='C'], 
             [lambda X: "X" if np.random.random() < .5 else "Y",
              lambda X: "X" if np.random.random() < .6 else "Y",
              lambda X: "X" if np.random.random() < .35 else "Y"
             ])

out: array(['X', 'X', 'Y', ..., 'Y', 'X', 'X']

~660ms 在我的机器上进行 1000 万次模拟，YMMV

如何在 Python/numpy 中创建嵌套加权分布？

How to create nested weighted distributions in Python/numpy?

python

numpy

montecarlo