Pandas стратифицирует разделение на обучающий, тестовый и проверочный наборы на основе целевой переменной и ее кластера.

У меня есть фреймворк данных с некоторыми функциями и целевой столбец, принадлежащий {0,1}. Мне нужно разделить этот набор данных на наборы для обучения, тестирования и проверки. Часть проверки должна составлять 20% набора данных, а оставшиеся 80% должны быть разделены так, чтобы 80% вошли в обучающий набор. И это может быть легко достигнуто с помощью sklearntrain_test_split

Моя проблема в том, что разделение должно быть выполнено стратифицированным способом на основе кластеров, которые я вычислил для обоих целевых значений.

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

ones = df[df_numerical['Target'] == 1].copy()
zeroes = df[df_numerical['Target'] == 1].copy()

Затем для каждого подмножества я использовал kmeans для вычисления их кластеров и добавлял кластеры в кадр данных, например:

# the number of clusters for both variables is not the same
clusters_1 = kmeans_1.predict(ones[NUMERICAL_FEATURES])
ones['Cluster'] = clusters_1
clusters_0 = kmeans_0.predict(zeroes[NUMERICAL_FEATURES])
zeroes['Cluster'] = clusters_0

Теперь, как я могу разделить наборы данных, чтобы они были стратифицированы по размеру кластера?

Нужное мне разбиение должно быть выполнено следующим образом: при условии наличия 100 записей, 80 из класса 1 и 20 из класса 0, мне нужно разделить эти записи на 70/30%, поэтому мне нужно иметь 56 (70% от 80) записей класса 1 и 14 (70% от 20) класса 0. И я знаю, что это можно сделать с помощью stratifyпараметра train_test_split, но моя проблема в том, что кроме этого разбиение должно быть расслоено еще и по кластерам каждое целевое значение.

Я подумал, что одним из решений было бы извлечь индексы элементов для обоих классов, поместить их в списки, извлечь из них нужное количество элементов, а затем повторно объединить кадры данных:

cluster_indices_0 = zeroes.groupby(['Cluster']).apply(lambda x: x.index)
cluster_indices_1 = ones.groupby(['Cluster']).apply(lambda x: x.index)

Но таким образом мне пришлось бы вручную вычислять для каждого кластера количество элементов, которые нужно извлечь, и я искал способ сделать это автоматически.

Есть ли функция в sklearn или pandas для достижения того, что я ищу, без получения списка при вычислении количества элементов для извлечения?

Решение проблемы

Поскольку у вас уже есть данные, разделенные по целевым объектам, вам просто нужно вызвать train_test_splitкаждое подмножество и использовать столбец кластера для стратификации.

train_test_0, validation_0 = train_test_split(zeroes, train_size=0.8, stratify=zeroes['Cluster'])
train_0, test_0 = train_test_split(train_test_0, train_size=0.7, stratify=train_test_0['Cluster'])

затем сделайте то же самое для целевого и объедините все подмножества