Question about sklearn's StratifiedShuffleSplit I'm reading through the book Hands-On Machine Learning with Scikit-Learn and Tensorflow by Aurélien Géron. In a regression project on California Housing Prices, he goes over the concept of stratified sampling.

I think I understand the concept as his explanation "the population is divided into homogeneous subgroups called strata, and the right number of instances is sampled from each stratum to guarantee that the test set is representative of the overall population."

So in my own words, simply splitting the dataset with sklearn's train_test_split leaves the train and test set vulnerable to misrepresenting the ratios of categorical variables (ie population has 40% one category, 60% another, but the train/test set are totally different ratios of these categories), so stratifying ensures the sample is 'random', but still maintaining proper ratios within test and train splits. Please correct me if I'm wrong.

Here's the code to his stratified sampling based on income categories (housing is the main dataframe):

split = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
for train_index, test_index in split.split(housing, housing['income_cat']):
    strat_train_set = housing.loc[train_index]
    strat_test_set = housing.loc[test_index]


I'm pretty confused with this code:

1) What does the variable 'split' represent? Does it comprise both the train and test split...?

2) In the 2nd line of code, what does split.split mean? I guess I'm confused with most of how StratifiedShuffleSplit divides the train and test set and why he needed to create this 'for' loop in order to create strat_train_set and strat_test_set.

Thanks,

Greg
split


Is the object that allows us to do stratified split, and split.split is the split’s object method/function called split that can be used to perform stratified split.

Depending on how many splits you want using the n_split parameter you can split, hence the for loop.