# Statistical Sampling with Scikit-learn

Setting a test set is one of the early stages of developing a machine learning model. Thus creating a representative sample will determine how the model will perform in production.

The following example demonstrates the difference between stratified sampling and random sampling.

**0. Step Nada, imports:**

`import numpy as np`

import pandas as pd

from sklearn.model_selection import train_test_split

from sklearn.model_selection import StratifiedShuffleSplit

**First we create a population to work on:**

np.random.seed(42) # (1)p = 0.8

a = ["Circle", "Square"]

size = 100data = np.random.choice(a=a, size=size, p=[p, 1-p]) # (2)

- a seed, to make random numbers reproducible
- the population of circles and squares, the sample size is 100 with a population of 0.8 and 0.2

**2. Then, we convert the data to a data frame and process it:**

df = pd.DataFrame(data=data, columns=['X']) # (1)df.replace({'Circle': 0, 'Square': 1}, inplace=True) # (2)df['y'] = df.apply(lambda row: int(row["X"] == 0), axis=1) # (3)

- converting the data to a data frame
- encoding the data
- adding a label column

**3. The random sampling step:**

X_train, X_test, y_train, y_test = train_test_split(df['X'], df['y'], test_size=0.2) # (1)p_random = X_test.value_counts() / len(X_test) # (2)

- sample size of 20%
- p_random result is 0.85 and 0.15, a skewed sample

**4. The stratified sampling step:**

split = StratifiedShuffleSplit(n_splits=1, test_size=0.2) # (1)for train_index, test_index in split.split(df, df['X']):

strat_train = df.loc[train_index]

strat_test = df.loc[test_index]p_strat = strat_test['X'].value_counts() / len(strat_test) # (2)

- sample size of 20%
- p_strat result is 0.8 and 0.2, just like the original population.

In conclusion, the result is as expected. The stratified sampling was more representative than random sampling; this is only because the play was staged to demonstrates the differences. Increase the size of the population, and see what happens.

- The overall population is: 0.8, 0.2
- The random sample population is: 0.85, 0.15
- The stratified sample populations is: 0.8, 0.2

The stratified sampling function (StratifiedShuffleSplit) will return a sample from each stratum (in this case, our strata was column ‘X’).

Happy cross-validating!

The complete code is here https://github.com/booletic/medium/blob/main/strata.ipynb