有没有相当于STATA的splitsample命令的Python库？

Question

上下文： STATA 有一个名为“splitsample”的命令，它可以根据某些用户定义的值有效地将当前数据集拆分为不同的分区。 nsplit() 定义多少个组，gen() 为每个组创建组标识符。

还可以传递其他参数，例如：balance() 和 rround。

balance() 强制拆分组保留某些变量的半相等比例/平均值

rround 有助于随机舍入样本大小以满足样本大小比率

splitsample, gen(group) nsplit(2) balance(clicks networkrevenue) rround

这是 STATA 中的输出示例

由于我在 STATA 中指定了 rround，所以比例接近 50/50，但不准确

tab group

  group |      Freq.     Percent        Cum.
      1 |      1,903       50.44       50.44
      2 |      1,870       49.56      100.00
  Total |      3,773      100.00

因为我还指定了：

balance(clicks networkrevenue)

当我查看两组的汇总统计数据时，他们的平均值也很相似

tab group, sum(clicks)

        |          Summary of Clicks
  group |        Mean   Std. dev.       Freq.
      1 |    30.31319   108.58514       1,903
      2 |   29.740107   98.621471       1,870
  Total |   30.029155   103.75319       3,773

 tab group, sum(networkrevenue)

        |      Summary of NetworkRevenue
  group |        Mean   Std. dev.       Freq.
      1 |   44194.721   242880.98       1,903
      2 |   44970.651   220502.77       1,870
  Total |   44579.293   232029.24       3,773

我正在寻找的是与此命令等效的 Python，其行为与balance() 和 rround 参数类似。

我对可用的 Python 采样库做了一些粗略的研究，但它们的行为都与 STATA 的 splitsample 不同。

我研究了 Imbalanced-learn 的过度和欠采样方法 + scikit 的采样方法。

非常感谢任何帮助！

Answer 1

您可以很容易地编写自己的代码。可能没有完全相同的实现：

import numpy as np
import pandas as pd

def gr(df, n):
    a = np.resize(range(n),df.shape[0])
    np.random.shuffle(a)
    return a

def splitsample(df, gen, nsplit=2, balance=None):
    if balance is None:
        df[gen] = gr(df,nsplit)
    else:
        df[gen] = df.groupby(balance)[balance[0]].transform(lambda d: gr(d, n=nsplit))
    return df

用途：

splitsample(df, gen="group", nsplit=2, balance=['clicks','network_revenue'])