Sklearn.cross_validation.StratifiedShuffleSplit - ошибка: "индексы вне пределов"

Я пытался разбить образец набора данных, используя Scikit-learn Stratified Shuffle Split. Я последовал примеру, показанному в документации Scikit-learn здесь

import pandas as pd
import numpy as np
# UCI wine dataset
wine = pd.read_csv("https://s3.amazonaws.com/demo-datasets/wine.csv")

# separate target variable from dataset
target = wine['quality']
data = wine.drop('quality',axis = 1)

# Stratified Split of train and test data
from sklearn.cross_validation import StratifiedShuffleSplit
sss = StratifiedShuffleSplit(target, n_iter=3, test_size=0.2)

for train_index, test_index in sss:
    xtrain, xtest = data[train_index], data[test_index]
    ytrain, ytest = target[train_index], target[test_index]

# Check target series for distribution of classes
ytrain.value_counts()
ytest.value_counts()

Однако при запуске этого script я получаю следующую ошибку:

IndexError: indices are out-of-bounds

Может ли кто-нибудь указать, что я делаю неправильно здесь? Спасибо!

Ответ 1

Вы используете разные соглашения для индексации Pandas DataFrame индексации и индексирования NumPy ndarray. Массивы train_index и test_index являются наборами индексов строк. Но data - это объект Pandas DataFrame, и когда вы используете один индекс в этом объекте, как в data[train_index], Pandas ожидает, что train_index будет содержать метки столбцов, а не индексы строк. Вы можете преобразовать dataframe в массив NumPy, используя .values:

data_array = data.values
for train_index, test_index in sss:
    xtrain, xtest = data_array[train_index], data_array[test_index]
    ytrain, ytest = target[train_index], target[test_index]

или используйте Pandas .iloc:

for train_index, test_index in sss:
    xtrain, xtest = data.iloc[train_index], data.iloc[test_index]
    ytrain, ytest = target[train_index], target[test_index]

Я предпочитаю второй подход, поскольку он дает xtrain и xtest типа DataFrame, а не ndarray, и поэтому сохраняет метки столбцов.