从具有定义概率的 Pandas 组内进行抽样

Question

考虑以下 Pandas 数据框，

df = pd.DataFrame(
    [
         ['X', 0, 0.5],
         ['X', 1, 0.5],

         ['Y', 0, 0.25],
         ['Y', 1, 0.3],
         ['Y', 2, 0.45],

         ['Z', 0, 0.6],
         ['Z', 1, 0.1],
         ['Z', 2, 0.3]
    ], columns=['NAME', 'POSITION', 'PROB'])

请注意，df 为每个唯一的 NAME 值定义了离散概率分布，即

assert ((df.groupby('NAME')['PROB'].sum() - 1)**2 < 1e-10).all()

我想做的是从这些概率分布中抽样。

我们可以将POSITION 视为对应于概率的值。因此，当考虑X 时，样本将是0，概率为0.5 和1，概率为0.5。

我想创建一个新的数据框，其中包含代表这些样本的列 ['NAME', 'POSITION', 'PROB', 'SAMPLE']。每个唯一的SAMPLE 值代表一个新样本。 PROB 列现在始终为 0 或 1，表示在给定样本中是否选择了给定行。例如，如果我要选择 3 个样本，下面是一个示例结果，

df_samples = pd.DataFrame(
    [
         ['X', 0, 1, 0],
         ['X', 1, 0, 0],
         ['X', 0, 0, 1],
         ['X', 1, 1, 1],
         ['X', 0, 1, 2],
         ['X', 1, 0, 2],

         ['Y', 0, 1, 0],
         ['Y', 1, 0, 0],
         ['Y', 2, 0, 0],
         ['Y', 0, 0, 1],
         ['Y', 1, 0, 1],
         ['Y', 2, 1, 1],
         ['Y', 0, 1, 2],
         ['Y', 1, 0, 2],
         ['Y', 2, 0, 2],

         ['Z', 0, 0, 0],
         ['Z', 1, 0, 0],
         ['Z', 2, 1, 0],
         ['Z', 0, 0, 1],
         ['Z', 1, 0, 1],
         ['Z', 2, 1, 1],
         ['Z', 0, 1, 2],
         ['Z', 1, 0, 2],
         ['Z', 2, 0, 2],
    ], columns=['NAME', 'POSITION', 'PROB', 'SAMPLE'])

当然，由于涉及到随机性，这只是众多可能结果之一。

该程序的单元测试是随着样本的增加，根据大数定律，每个(NAME, POSITION) 对的样本平均数应该趋于实际概率。可以根据使用的总样本计算置信区域，然后确保真实概率在其中。例如，使用 normal approximation to binomial outcomes（要求总样本 n_samples 为“大”）（-4 sd，4 sd）区域测试将是：

z = 4

p_est = df_samples.groupby(['NAME', 'POSITION'])['PROB'].mean()
p_true = df.set_index(['NAME', 'POSITION'])['PROB']

CI_lower = p_est - z*np.sqrt(p_est*(1-p_est)/n_samples)
CI_upper = p_est + z*np.sqrt(p_est*(1-p_est)/n_samples)

assert p_true < CI_upper
assert p_true > CI_lower

在 Pandas 中最有效的方法是什么？我觉得我想将一些sample 函数应用于df.groupby('NAME') 对象。

附言

更明确地说，这是使用 Numpy 执行此操作的一种非常冗长的方法。

n_samples = 3
df_list = []
for name in ['X', 'Y', 'Z']:
    idx = df['NAME'] == name
    position_samples = np.random.choice(df.loc[idx, 'POSITION'], 
                                        n_samples, 
                                        p=df.loc[idx, 'PROB'])
    prob = np.zeros([idx.sum(), n_samples])
    prob[position_samples, np.arange(n_samples)] = 1
    position = np.tile(np.arange(idx.sum())[:, None], n_samples)
    sample = np.tile(np.arange(n_samples)[:,None], idx.sum()).T

    df_list.append(pd.DataFrame(
        [[name, prob.ravel()[i], position.ravel()[i], 
          sample.ravel()[i]] 
         for i in range(n_samples*idx.sum())], 
        columns=['NAME', 'PROB', 'POSITION', 'SAMPLE']))

df_samples = pd.concat(df_list)

Answer 1

如果我理解正确，您正在寻找 groupby + sample 然后是一些索引内容

概率的第一个样本：

n_samples = 3
df_samples = df.groupby('NAME').apply(lambda x: x[['NAME', 'POSITION']] \
                               .sample(n_samples, replace=True,
                                       weights=x.PROB)) \
                               .reset_index(drop=True)

现在添加额外的列：

df_samples['SAMPLE'] = df_samples.groupby('NAME').cumcount()
df_samples['PROB'] = 1


print(df_samples)

  NAME  POSITION  SAMPLE  PROB
0    X         1       0     1
1    X         0       1     1
2    X         1       2     1
3    Y         1       0     1
4    Y         1       1     1
5    Y         1       2     1
6    Z         2       0     1
7    Z         0       1     1
8    Z         0       2     1

请注意，这不包括初始问题中要求的每个样本的 0 概率位置，但它是一种更简洁的信息存储方式。

如果我们还想包含概率为 0 的位置，我们可以合并到其他位置，如下所示：

domain = df[['NAME', 'POSITION']].drop_duplicates()
df_samples.drop('PROB', axis=1, inplace=True)
df_samples = pd.merge(df_samples, domain, on='NAME', 
                      suffixes=['_sample', ''])
df_samples['PROB'] = (df_samples['POSITION'] ==
                     df_samples['POSITION_sample']).astype(int)
df_samples.drop('POSITION_sample', axis=1, inplace=True)