Grouped and Stratified K-Fold CV

[gkiar]

Hi!

I'm trying to prepare dataset splits for a problem I'm working on, and would ultimately like a hybrid of Stratified K-Fold and Grouped K-Fold. Is there a way to accomplish this using logic already built into sklearn? If not, where would be the right place for me to add it/do you have any suggestions for how to get started before I give it a go?

For a bit more context, my dataset is a considerably larger version of the following structure:

groups = [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4]  # you can think of this as a sample-id
y = [1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2]  # this is some trait of the samples 

The objective would be to balance y across the train and test sets, and have each group only represented on one side of the folds. Thanks in advance!

[NicolasHug]

Closing as a duplicate of #13621