sklearn train_test_split on pandas stratify by multiple columns

user_8419599 0
Viewed 54377

I'm a relatively new user to sklearn and have run into some unexpected behavior in train_test_split from sklearn.model_selection. I have a pandas dataframe that I would like to split into a training and test set. I would like to stratify my data by at least 2, but ideally 4 columns in my dataframe.

There were no warnings from sklearn when I tried to do this, however I found later that there were repeated rows in my final data set. I created a sample test to show this behavior:

from sklearn.model_selection import train_test_split
a = np.array([i for i in range(1000000)])
b = [i%10 for i in a]
c = [i%5 for i in a]
df = pd.DataFrame({'a':a, 'b':b, 'c':c})

It seems to work as expected if I stratify by either column:

train, test = train_test_split(df, test_size=0.2, random_state=0, stratify=df[['b']])
print(len(train.a.values))  # prints 800000
print(len(set(train.a.values)))  # prints 800000

train, test = train_test_split(df, test_size=0.2, random_state=0, stratify=df[['c']])
print(len(train.a.values))  # prints 800000
print(len(set(train.a.values)))  # prints 800000

But when I try to stratify by both columns, I get repeated values:

train, test = train_test_split(df, test_size=0.2, random_state=0, stratify=df[['b', 'c']])
print(len(train.a.values))  # prints 800000
print(len(set(train.a.values)))  # prints 640000
scikit-learn pandas python
4 Answers

If you want train_test_split to behave as you expected (stratify by multiple columns with no duplicates), create a new column that is a concatenation of the values in your other columns and stratify on the new column.

df['bc'] = df['b'].astype(str) + df['c'].astype(str)
train, test = train_test_split(df, test_size=0.2, random_state=0, stratify=df[['bc']])

If you're worried about collision due to values like 11 and 3 and 1 and 13 both creating a concatenated value of 113, then you can add some arbitrary string in the middle:

df['bc'] = df['b'].astype(str) + "_" + df['c'].astype(str)
user_56324220

You need to iteratively split your data. There's a class for it in scikit-multilearn. Bit annoying that it only works on NumPy arrays but what can you do?

Here's a function that should do what you are asking for:

import pandas as pd
from skmultilearn.model_selection import IterativeStratification

def iterative_split(df, test_size, stratify_columns):
    """Custom iterative train test split which
    'maintains balanced representation with respect
    to order-th label combinations.'

    From https://madewithml.com/courses/mlops/splitting/#stratified-split
    """
    # One-hot encode the stratify columns and concatenate them
    one_hot_cols = [pd.get_dummies(df[col]) for col in stratify_columns]
    one_hot_cols = pd.concat(one_hot_cols, axis=1).to_numpy()
    stratifier = IterativeStratification(
        n_splits=2, order=len(stratify_columns), sample_distribution_per_fold=[test_size, 1-test_size])
    train_indices, test_indices = next(stratifier.split(df.to_numpy(), one_hot_cols))
    # Return the train and test set dataframes
    train, test = df.iloc[train_indices], df.iloc[test_indices]
    return train, test

example = pd.DataFrame({'a': [1, 2, 3]*8*2, 'b': [4, 5, 6, 7]*6*2, 'c': [7, 8]*12*2})
train, test = iterative_split(example, 0.4, ['a', 'b'])
# print(f'{train =}')
# print(f'{test =}')

print(f'{train[["a"]].value_counts() =}')
print(f'{test[["a"]].value_counts()  =}')
print(f'{train[["b"]].value_counts() =}')
print(f'{test[["b"]].value_counts()  =}')

Output

train[["a"]].value_counts() =a
1    10
2    10
3    10
dtype: int64
test[["a"]].value_counts()  =a
1    6
2    6
3    6
dtype: int64
train[["b"]].value_counts() =b
5    8
6    8
4    7
7    7
dtype: int64
test[["b"]].value_counts()  =b
4    5
7    5
5    4
6    4
dtype: int64

And for your example we can add this code:

import numpy as np

a = np.array([i for i in range(10_000)])
b = [i%10 for i in a]
c = [i%5 for i in a]
df = pd.DataFrame({'a':a, 'b':b, 'c':c})

train, test = iterative_split(df, test_size=0.2, stratify_columns=['b', 'c'])

print(len(train.a.values))  # prints 8000
print(len(set(train.a.values)))  # prints 8000

The one_hot_cols becomes a matrix of 1e6 x 3e5 in your example and that was a bit much. If someone comes up with a better way then I am all ears.

user_34061890
Related