Use sklearn's train_test_split on the first level of a multi-indexed pandas DataFrame

pandas_multindex_train_test_split.py

import pandas as pd
from sklearn.model_selection import train_test_split

df = pd.DataFrame(data=[[1], [2], [3]], columns=['foo'], index=pd.MultiIndex.from_tuples(((0,'a'),(0, 'b'), (1, 'a'))))
train_ix, test_ix = train_test_split(df.index.levels[0])
train_df = df.loc[train_ix]
test_df = df.loc[test_ix]

Thanks for this!

Thanks for the code!