Stratified Sampling in Pandas

Question

I've looked at the Sklearn stratified sampling docs as well as the pandas docs and also Stratified samples from Pandas and sklearn stratified sampling based on a column but they do not address this issue.

Im looking for a fast pandas/sklearn/numpy way to generate stratified samples of size n from a dataset. However, for rows with less than the specified sampling number, it should take all of the entries.

Concrete example:

Thank you! :)

Answer 1

Use min when passing the number to sample. Consider the dataframe df

df = pd.DataFrame(dict(         A=[1, 1, 1, 2, 2, 2, 2, 3, 4, 4],         B=range(10)     ))  df.groupby('A', group_keys=False).apply(lambda x: x.sample(min(len(x), 2)))     A  B 1  1  1 2  1  2 3  2  3 6  2  6 7  3  7 9  4  9 8  4  8 

Answer 2

Extending the groupby answer, we can make sure that sample is balanced. To do so, when for all classes the number of samples is >= n_samples, we can just take n_samples for all classes (previous answer). When minority class contains < n_samples, we can take the number of samples for all classes to be the same as of minority class.

def stratified_sample_df(df, col, n_samples):     n = min(n_samples, df[col].value_counts().min())     df_ = df.groupby(col).apply(lambda x: x.sample(n))     df_.index = df_.index.droplevel(0)     return df_