разделить данные теста на поезд, используя стратификацию на две колонки в scikit-learn

Question

У меня есть набор данных, который я хочу разделить на тренировку и тестирование, чтобы у меня были данные в наборе тестов от каждого источника данных (указанного в столбце «источник») и от каждого класса (указанного в столбце «класс»). Я читал об использовании параметра stratifiy с функцией scikitlearn train_test_split, но как я могу использовать его в двух столбцах?

Answer 1

Расслоение на несколько столбцов легко выполняется с помощью sklearn's train_test_split начиная с версии 19.0

Доказательство

from sklearn.model_selection import train_test_split
from sklearn.datasets import make_multilabel_classification

X, Y = make_multilabel_classification(1000000, 10, n_classes=2, n_labels=1)
train_X, test_X, train_Y, test_Y =train_test_split(X,Y,stratify=Y, train_size=.8, random_state=42)
Y.shape

(1000000, 2)

Затем вы можете сравнить простые средства столбцов результирующих стратификаций:

train_Y[:,0].mean(), test_Y[:,0].mean()
(0.45422, 0.45422)

train_Y[:,1].mean(), test_Y[:,1].mean()
(0.23472375, 0.234725)

Запустите статистику t-test по равенству средних значений:

from scipy.stats import ttest_ind
ttest_ind(train_Y[:,0],test_Y[:,0])

Ttest_indResult(statistic=0.0, pvalue=1.0)

И, наконец, сделайте то же самое для условных средств, чтобы доказать, что вы действительно достигли что ты хотел:

train_Y[train_Y[:,0].astype("bool"),1].mean(), test_Y[test_Y[:,0].astype("bool"),1].mean()
(0.43959149751221877, 0.43958874554180793)