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 · 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

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

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

User · 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

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

[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