Получите одну случайную выборку для каждой группы и закончите стратифицированной выборкой pandas

Question

Я работаю с таким фреймом данных:

group    period
  A      20130101
  A      20130201
  .          .
  E      20130901
  E      20131001

Допустим, у меня есть 100 различных групп и 10 возможных дат, которые распределены следующим образом: [.1,.05,.2,.05,.1,.1,.2,.05,.05,.1]. Мне нужно получить один образец для каждой группы, поэтому 10% окончательного образца получают из первого периода, 5% из второго периода, 20% из третьего периода и так далее. Мне удалось получить случайную выборку для каждой группы, но она сильно искажена, например:

fn = lambda obj: obj.loc[np.random.choice(obj.index, 1, replace=False),:]
dfrd = df[['group','period']].groupby('group', as_index=False).apply(fn)
dfrd.index = [index[1] for index in dfrd.index]

Итак, есть ли способ сделать что-то похожее, но стратифицированное? Спасибо

Answer 1

Вы можете использовать p параметр из np.random.choice:

df1 = (df
      .groupby('grp')
      .apply(lambda x: np.random.choice(x['period'].values, size=1, p=prob)[0])
      .reset_index()
      .rename(columns={0:'period'}))

  grp      period
0   A  2013-01-03
1   B  2013-01-04
2   C  2013-01-04
3   D  2013-01-03

Пример данных

period = list(map(str, pd.date_range(start='20130101', freq='D', periods=10).date))
grp = sorted(['A','B','C','D']*10)
prob = [.1,.05,.2,.05,.1,.1,.2,.05,.05,.1]

df = pd.DataFrame({'grp': grp, 'period': period*4})