# How to stratify a dataset to keep groups of data together in Python?

I have a dataset I want to use for machine learning that looks something like this:

Group_ID | Column_1 | Column_2 | Column_3 ...
==========================================
A        | 1        | 2        | 33
A        | 2        | 2        | 3765
A        | 3        | 6        | 3436
A        | 4        | 8        | 32
B        | 5        | 9        | 33
B        | 3        | 34       | 385
B        | 7        | 25       | 3
B        | 3        | 1        | 38
C        | 6        | 2        | 3
C        | 8        | 2        | 4
D        | 7        | 1        | 5
D        | 6        | 9        | 11


Where Group_ID is an identifier which isn't used by the model, it's just a reference I have, and the rest of the columns are the predictors.

I want to create a training and test set where the rows are split say 70-30, but all rows belonging to a group are always in the same dataset.

So if we aggregate the counts of rows, and get something like this:

A | 145
B | 110
C | 60
D | 35


The final split should be ~70-30, so approx. 245-105, so I'd want the final training-test split to be:

Training set: Group A, C and D
Test set: Group B


I want to keep them together to see if there is a difference between the groups in terms of how well my classifiers perform.

First of all:

• statistically speaking, is there anything wrong with splitting like this?

and second:

• is there a simple way of doing this kind of stratified sampling (in Python)?

This is what I ended up doing, which is hopefully valid. I'm making the assumption that there is no difference between the various groups.

1) Aggregate the group counts (as in the question)

A | 145
B | 110
C | 60
D | 35


2) Create a sample 70% the size of the original dataset by sampling from the groups with a probability proportional to their counts (using numpy.choice), which in the above case would be:

A | 41.43%
B | 31.43%
C | 17.14%
D | 10%


3) Take all the rows from the original dataset that are in this sampled list of groups to create the training set (and the test set is what's left over).

This results in an approximate 70-30 split where rows belonging to a group are always in the same dataset.