scikit-learn StratifiedShuffleSplit KeyError with index - python-3.x

Question

This is my pandas dataframe lots_not_preprocessed_usd:
<class 'pandas.core.frame.DataFrame'>
Index: 78718 entries, 2017-09-12T18-38-38-076065 to 2017-10-02T07-29-40-245031
Data columns (total 20 columns):
created_year 78718 non-null float64
price 78718 non-null float64
........
decade 78718 non-null int64
dtypes: float64(8), int64(1), object(11)
memory usage: 12.6+ MB
head(1):
artist_name_normalized house created_year description exhibited_in exhibited_in_museums height images max_estimated_price min_estimated_price price provenance provenance_estate_of sale_date sale_id sale_title style title width decade
key
2017-09-12T18-38-38-076065 NaN c11 1862.0 An Album and a small Quantity of unframed Draw... NaN NaN NaN NaN 535.031166 267.515583 845.349242 NaN NaN 1998-06-21 8033 OILS, WATERCOLOURS & DRAWINGS FROM 18TH - 20TH... watercolor painting An Album and a small Quantity of unframed Draw... NaN 186
My script:
from sklearn.model_selection import StratifiedShuffleSplit
split = StratifiedShuffleSplit(n_splits=1, test_size =0.2, random_state=42)
for train_index, test_index in split.split(lots_not_preprocessed_usd, lots_not_preprocessed_usd['decade']):
strat_train_set = lots_not_preprocessed_usd.loc[train_index]
strat_test_set = lots_not_preprocessed_usd.loc[test_index]
I'm getting the error message
KeyError Traceback (most recent call last)
<ipython-input-224-cee2389254f2> in <module>()
3 split = StratifiedShuffleSplit(n_splits=1, test_size =0.2, random_state=42)
4 for train_index, test_index in split.split(lots_not_preprocessed_usd, lots_not_preprocessed_usd['decade']):
----> 5 strat_train_set = lots_not_preprocessed_usd.loc[train_index]
6 strat_test_set = lots_not_preprocessed_usd.loc[test_index]
......
KeyError: 'None of [[32199 67509 69003 ..., 44204 2809 56726]] are in the [index]'
There seems to be a problem with my index (e.g. 2017-09-12T18-38-38-076065) which I don't understand. Where is the issue?
If I use another split it works as expected:
from sklearn.model_selection import train_test_split
train_set, test_set = train_test_split(lots_not_preprocessed_usd, test_size=0.2, random_state=42)

Answer 1

When you use .loc you need to pass same index for row_indexer so use .iloc when you want to use orindary numeric indexer instead of .loc. In the for loop train_index and text_index are not datetime since split.split(X,y) return array of random indices.
...
for train_index, test_index in split.split(lots_not_preprocessed_usd, lots_not_preprocessed_usd['decade']):
strat_train_set = lots_not_preprocessed_usd.iloc[train_index]
strat_test_set = lots_not_preprocessed_usd.iloc[test_index]
Sample example
lots_not_preprocessed_usd = pd.DataFrame({'some':np.random.randint(5,10,100),'decade':np.random.randint(5,10,100)},index= pd.date_range('5-10-15',periods=100))
for train_index, test_index in split.split(lots_not_preprocessed_usd, lots_not_preprocessed_usd['decade']):
strat_train_set = lots_not_preprocessed_usd.iloc[train_index]
strat_test_set = lots_not_preprocessed_usd.iloc[test_index]
Sample output :
strat_train_set.head()
decade some
2015-08-02 6 7
2015-06-14 7 6
2015-08-14 7 9
2015-06-25 9 5
2015-05-15 7 9