Scikit-learn train test split with indices

Question

How do I get the original indices of the data when using train test split     What I have is the following  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    But this does not give the indices of the original data   One workaround is to add the indices to data  e g  data     i  d  for i  d in enumerate data    and then pass them inside train test split and then expand again   Are there any cleaner solutions

User · Accepted Answer

Scikit learn plays really well with Pandas, so I suggest you use it. Here's an example:

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

You can directly call any scikit functions on DataFrame/Series and it will work.

Let's say you wanted to do a LogisticRegression, here's how you could retrieve the coefficients in a nice way:

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

User · Answer

The docs mention train test split is just a convenience function on top of shuffle split   I just rearranged some of their code to make my own example  Note the actual solution is the middle block of code  The rest is imports  and setup for a runnable example   from sklearn model selection import ShuffleSplit from sklearn utils import safe indexing  indexable from itertools import chain import numpy as np X   np reshape np random randn 20   10 2     10 training examples y   np random randint 2  size 10    10 labels seed   1  cv   ShuffleSplit random state seed  test size 0 25  arrays   indexable X  y  train  test   next cv split X X   iterator   list chain from iterable       safe indexing a  train       safe indexing a  test       train      test       for a in arrays    X train  X test  train is  test is  y train  y test          iterator  print X  print train is  print X train    Now I have the actual indexes  train is  test is

User · Answer

If you are using pandas you can access the index by calling  index of whatever array you wish to mimic  The train test split carries over the pandas indices to the new dataframes  In your code you simply use x1 index and the returned array is the indexes relating to the original positions in x

User · Answer

You can use pandas dataframes or series as Julien said but if you want to restrict your-self to numpy you can pass an additional array of indices   from sklearn model selection import train test split import numpy as np n samples  n features  n classes   10  2  2 data   np random randn n samples  n features     10 training examples labels   np random randint n classes  size n samples     10 labels indices   np arange n samples  x1  x2  y1  y2  idx1  idx2   train test split      data  labels  indices  test size 0 2

User · Answer

Here s the simplest solution  Jibwa made it seem complicated in another answer   without having to generate indices yourself - just using the ShuffleSplit object to generate 1 split   import numpy as np  from sklearn model selection import ShuffleSplit   or StratifiedShuffleSplit sss   ShuffleSplit n splits 1  test size 0 1   data size   100 X   np reshape np random rand data size 2   data size 2   y   np random randint 2  size data size   sss get n splits X  y  train index  test index   next sss split X  y     X train  X test   X train index   X test index   y train  y test   y train index   y test index

[python] Scikit-learn train_test_split with indices

Examples related to python

Examples related to scipy

Examples related to scikit-learn

Examples related to classification