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

Question

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

使用普通 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 sn-p 运行大约需要 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

您可以在对上使用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 2

可以显着减少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 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']

在我的机器 YMMV 上进行 1000 万次模拟时约为 660 毫秒