[解決済み] Scikit-learnのtrain_test_splitとindexの関係

質問

train_test_split()で元のデータのインデックスを取得するにはどうしたらよいですか？

私が持っているのは以下のものです。

from sklearn.cross_validation import train_test_split
import numpy as np
data = np.reshape(np.randn(20),(10,2)) # 10 training examples
labels = np.random.randint(2, size=10) # 10 labels
x1, x2, y1, y2 = train_test_split(data, labels, size=0.2)

しかし、これでは元データのインデックスを得ることができません。 回避策として、データにインデックスを追加する方法があります（例えば data = [(i, d) for i, d in enumerate(data)] のように) データにインデックスを追加し、それを train_test_split の中に渡して、再び展開します。 よりクリーンな解決策はありますか？

どのように解決するのですか？

Scikit learnはPandasと非常に相性が良いので、ぜひ使ってみてください。以下はその例です。

In [1]: 
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
data = np.reshape(np.random.randn(20),(10,2)) # 10 training examples
labels = np.random.randint(2, size=10) # 10 labels

In [2]: # Giving columns in X a name
X = pd.DataFrame(data, columns=['Column_1', 'Column_2'])
y = pd.Series(labels)

In [3]:
X_train, X_test, y_train, y_test = train_test_split(X, y, 
                                                    test_size=0.2, 
                                                    random_state=0)

In [4]: X_test
Out[4]:

     Column_1    Column_2
2   -1.39       -1.86
8    0.48       -0.81
4   -0.10       -1.83

In [5]: y_test
Out[5]:

2    1
8    1
4    1
dtype: int32

DataFrame/Series上で直接scikitの関数を呼び出しても動作します。

例えば、LogisticRegressionを行いたい場合、どのように係数を取得すれば良いかを説明します。

In [6]: 
from sklearn.linear_model import LogisticRegression

model = LogisticRegression()
model = model.fit(X_train, y_train)

# Retrieve coefficients: index is the feature name (['Column_1', 'Column_2'] here)
df_coefs = pd.DataFrame(model.coef_[0], index=X.columns, columns = ['Coefficient'])
df_coefs
Out[6]:
            Coefficient
Column_1    0.076987
Column_2    -0.352463