如何用default_rng生成分组可控的不平衡实验数据？

问题说明 #

需生成模拟实验数据，核心需求：

• response变量呈现不平衡分布
• 精准控制treatment（0/1组）之间的response差异，避免数据分布过于均匀

原代码通过rng.integers生成数据，导致两组response均值几乎无差异（0组0.501696，1组0.513530），无法满足需求。

原代码 #

rng = np.random.default_rng(seed=42)
# 样本量
size = 10000
df = pd.DataFrame({
    'treatment': rng.integers(0,2, size=size),
    'response': rng.integers (0,2, size=size)
})

解决方案思路 #

为treatment不同组别设置独立的response生成概率，直接控制每组的response分布：

• 分别生成对应treatment=1和treatment=0的偏置response序列
• 根据treatment取值，选择对应序列作为最终response

实现代码 #

rng = np.random.default_rng(seed=42)
# 样本量
size = 10000
df = pd.DataFrame({
    'treatment': rng.integers(0,2, size=size),
    'response1': rng.choice ([0,1], size=size, p=[0.1, 0.9]),  # treatment=1时的response分布（1占90%）
    'response2': rng.choice ([0,1], size=size, p=[0.9, 0.1])   # treatment=0时的response分布（0占90%）
})
df['response'] = np.where(df['treatment']==1, df['response1'], df['response2'])

效果验证 #

运行代码后，可通过以下代码查看两组response的均值：

print(df.groupby('treatment')['response'].mean())

输出示例：

treatment
0    0.1024
1    0.8957
Name: response, dtype: float64

两组response分布差异显著，完全符合预设的概率偏置。

内容的提问来源于stack exchange，提问作者Jameson