Question about sklearn's StratifiedShuffleSplit

Question

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

Answer 1

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.

Answer 2

The 'StratifiedShuffleSplit' function takes parameters on how the split needs to take place and returns a function to do the split.

The 'split' variable in the first line is used to store this function. In Python, functions/procedures can be stored as variables.

'n_splits' indicates the number of folds. 'test_size' indicates the proportion of the test data set out of the complete dataset.

The for loop allows you to iterate over the multiple splits when n_splits>1. Example, if n_splits=2 and test_size=0.2, the dataset 'housing' is split into two sets of (80% training and 20% test)

@Greg Rosen

Answer 3

I too had similar doubts as OP, then I went through the Scikit-Learn's GitHub repo. Repo Link to the file in discussion

And found out that:

1. split.split(housing, housing['income_cat']): Is nothing but the first split is the object of class StratifiedShuffleSplit() i.e can be any name you want to keep, and the second split is the function in this class (that refers to split function of base class BaseShuffleSplit()) the use of this is to Generate indices to split data into training and test set.

2. The reason for using a loop is best explained by @SparkZeus above.

Happy building models !!