# Dealing with the test set of imbalanced data

I am working on a problem dealing with unbalanced data that has a very specific request.

I would like to know the following:

When I have an imbalanced dataset and I do train test split, the test samples also have an imbalance. How can I at least balance the test set so that when I do the confusion matrix equal numbers are tested?

Here's an example (Classification):

If I have three groups: A 60,B 30,C 20. For simplicity, the test set is 10%. So, instead of having a train/test set of 54/6, 27/3, 18/2, I would like 58/2, 28/2 and 18/2 OR 54/2, 27/2, 18/2.

How can I do this?

For a case like this, it would probably be the easiest to do the train/test split manually. Here is a quick example with Pandas:

import pandas as pd
import numpy as np

# Create an empty dataframe with 60*A, 30*B, 20*C and some random target values
df = pd.DataFrame(columns=['category'])
df['category'] = (["A"] * 60 + ["B"] * 30 + ["C"] * 20)
df['target'] = np.random.rand(len(df))
df.category.value_counts()

> A    60
> B    30
> C    20
> Name: category, dtype: int64

# Obtain the size of the least represented category
smallest_category_size = np.min(df.category.value_counts().values)

# Size of the test split of the smallest category
test_size = 0.1

# Test count (2 in this case)
test_count = int(test_size*smallest_category_size)

# New train/test dataframes
train_df = pd.DataFrame(columns=df.columns)
test_df = pd.DataFrame(columns=df.columns)

for cat in df.category.unique():
# Select only a part of the dataset containing one category
cat_df = df[df.category == cat]

# Create and shuffle boolean mask
mask = np.array((len(cat_df) - test_count) * [True] + test_count * [False])

# Append the selected values to the train/test

train_df.category.value_counts()

> A    58
> B    28
> C    18
> Name: category, dtype: int64

test_df.category.value_counts()

> A    2
> B    2
> C    2
> Name: category, dtype: int64


I would still use scikit-learn's built-in mechanism for train/test splitting but downsample after.

import numpy as np
from sklearn.model_selection import train_test_split

# Create mock data
n = 110
X = np.arange(n).reshape((n, 1))
y = [0]*60 + [1]*30 + [2]*20 # Uneven categories

# Standard train test split
X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y)

# Fixed number of samples from each category
n_test = 2

# A category sampling
a_indices = [i for i, n in enumerate(y_test) if n == 0]
a_indices_sub_sample = np.random.choice(a_indices, n_test, replace=False)
X_test_a_category = X_test[a_indices_sub_sample]

# B category sampling
b_indices = [i for i, n in enumerate(y_test) if n == 1]
b_indices_sub_sample = np.random.choice(b_indices, n_test, replace=False)
X_test_b_category = X_test[b_indices_sub_sample]

# C category sampling
c_indices = [i for i, n in enumerate(y_test) if n == 2]
c_indices_sub_sample = np.random.choice(c_indices, n_test, replace=False)
X_test_c_category = X_test[c_indices_sub_sample]