How to split train/test datasets having equal classes proportion

Question

I would like to know how I can split in an equal number the following

Target
0    1586
1     318

in order to have the same proportion of 0 and 1 classes in a dataset to train, if my dataset is called df and includes 10 columns, both numerical and categorical.

I would consider the following

y=df['Target']
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.1, stratify=y)

so to do a stratification, but I do not know if it is right and I would appreciate if you could confirm it or provide an alternative to do that.

Sample of data

Fin                 Eco   Target
High percentage     12      1
Low percentage      5       0
Medium percentage   48      0
NA                  3       1
TBC                 NA      1
Low percentage      25      0
Medium percentage   12      0

How can I check if it is actually splitting in equal classes proportion my dataset? I think the best way to train a model should be having an equal proportion of 0 and 1 values. Right now I have 5 times data with Target=0.

Answer 1

If what you want is to have the same proportion of classes, 50% - 0 and 50% - 1. Then, there are two techniques oversampling (sampling more data from the smaller class) and undersampling (sampling less data from the bigger class). But I won't recommend you this for your problem. Your label seems fairly well balanced 1/5 it's a great proportion.

Still, this library has some implementations to do that

In this blog, you can see an overview of imbalanced datasets, but yours is not. Choosing a proper metric is more important.

If what you want to do is keep the same proportions of the label across the splits, what you are doing is right.

To validate your model properly, the class distribution and the different splits (train, validation, test) should be similar.

In the train test split documentation , you can find the argument:

stratifyarray-like, default=None If not None, data is split in a stratified fashion, using this as the class labels.

One step beyond will be using Stratified K-Folds cross-validator.

This cross-validation object is a variation of KFold that returns stratified folds. The folds are made by preserving the percentage of samples for each class.

There are more splitting techniques in Scikit Learn that you can use, have a look.

To test if the function is doing what you want just calculate the percentages in the splits:

 np.unique(y_train, return_counts=True)
 np.unique(y_val, return_counts=True)

But this will make you have the same proportions across the whole data, if your original label proportion is 1/5, then you will have 1/5 in train and 1/5 in test.

Answer 2

you can try stratified sampling method

from sklearn.model_selection import StratifiedShuffleSplit
split=StratifiedShuffleSplit(n_split=1, test_size=0.2, random_state=9)

Answer 3

This is the function I am using. You can adapt it to your needs:

# Returns a Test dataset that contains an equal amounts of each class
# y should contain only two classes 0 and 1
def TrainSplitEqualBinary(X, y, samples_n): #samples_n per class
    
    indicesClass1 = []
    indicesClass2 = []
    
    for i in range(0, len(y)):
        if y[i] == 0 and len(indicesClass1) < samples_n:
            indicesClass1.append(i)
        elif y[i] == 1 and len(indicesClass2) < samples_n:
            indicesClass2.append(i)
            
        if len(indicesClass1) == samples_n and len(indicesClass2) == samples_n:
            break
    
    X_test_class1 = X[indicesClass1]
    X_test_class2 = X[indicesClass2]
    
    X_test = np.concatenate((X_test_class1,X_test_class2), axis=0)
    
    #remove x_test from X
    X_train = np.delete(X, indicesClass1 + indicesClass2, axis=0)
    
    Y_test_class1 = y[indicesClass1]
    Y_test_class2 = y[indicesClass2]
    
    y_test = np.concatenate((Y_test_class1,Y_test_class2), axis=0)
    
    #remove y_test from y
    y_train = np.delete(y, indicesClass1 + indicesClass2, axis=0)
    
    if (X_test.shape[0] != 2 * samples_n or y_test.shape[0] != 2 * samples_n):
        raise Exception("Problem with split 1!")
        
    if (X_train.shape[0] + X_test.shape[0] != X.shape[0] or y_train.shape[0] + y_test.shape[0] != y.shape[0]):
        raise Exception("Problem with split 2!")
    
    return X_train, X_test, y_train, y_test