Improving stratification in StratifiedGroupKFold

### Describe the workflow you want to enable

I think I found two types of issues in StratifiedGroupKFold. First, StratifiedGroupKFold sometimes doesn't give ideally stratified splits when shuffle=True. Second, I think I found a general bug where the code just isn't doing what I think is intended when shuffle=True. My main change is something of a "new feature", but if this should be instead submitted as a bug report, I can do that.

### Improving the stratification

In the below example, StratifiedGroupKFold is working as intended. It gives train/test splits with identical balances of Y. For this example, there are six groups. For four groups, the an average Y = 0.6, but for two groups the average is Y = 0.4. The train/test split is 50%/50% so the ideal split has train Y = 0.533 and test Y = 0.533 (two groups from the Y = 0.6 and one from Y = 0.4).

However, when random_state=0 is changed to random_state=15, it no longer gives perfectly balanced splits. For random_state=15, the first fold has average Y = 0.6 and the second fold has average Y = 0.466. We generally want StratifiedGroupKFold to give the most ideal splits possible. My solution attempts to achieve this. My solution is also meant to just generally improve upon the author's original algorithm.


``` Python
import numpy as np
from sklearn.model_selection import StratifiedGroupKFold

subjects = 6
X = np.random.normal(size=(subjects*5, 1))
Y = [1, 1, 0, 0, 0]*(1)
Y += [1, 1, 0, 0, 1]*(2)
Y += [1, 1, 0, 0, 0]*(1)
Y += [1, 1, 0, 0, 1]*(2)
Y = np.array(Y)

groups = [[i]*5 for i in range(subjects)]
groups = [item for sublist in groups for item in sublist]

skf = StratifiedGroupKFold(n_splits=2, shuffle=True, random_state=0) # shift this to random_state=15

for train_index, test_index in skf.split(X, Y, groups=groups):
    Y_train, Y_test = Y[train_index], Y[test_index]
    Y_train_mean = np.mean(Y_train)
    Y_test_mean = np.mean(Y_test)
    print(f'{Y_train_mean=:.3f} | {Y_test_mean=:.3f}')
```

### Bug fix

I also found a bug in the original code related to shuffle=True. shuffle=true shuffles a matrix, y_counts_per_group. For this matrix, each row corresponded to a group and the columns were counts of how many times that group had a given value of y. So [[1, 2], [3, 1]] means group0 had one Y=0 & two Y=1, whereas group1 had three Y=0 and one Y=1.
This matrix was being shuffled. 

However, I think this creates a problem with another variable, group_inverse, so the shuffling should occur earlier. Group_inverse maps the index of the X or y arrays to a group number. That is, group_inverse[idxs] yields each idx's group. By shuffling y_counts_per_group, this causes the row-numbering for y_counts_per_group to become misaligned" with the mapping of group_inverse (e.g., row 1 in y_counts_per_group no longer corresponds to group 1 in group_inverse). 

I fixed this to do the shuffling earlier so the two are not misaligned. I noticed this issue because I needed to use y_counts_per_group for my solution below.


### Describe your proposed solution

### Improving stratification

I have a greedy algorithm solution. Basically, when non-perfect stratification is detected, I attempt to swap pairs of groups until the stratification is the best that it can be. I do this for all pairs of groups. If doing so improves the stratification, the swap is made and the process is restarted. The process keeps going until no pair of groups can be swapped to improve the stratification.

Notably, the current code already provides an approach that calculates the level of stratification:
``` Python
std_per_class = np.std(y_counts_per_fold / y_cnt.reshape(1, -1), axis=0)
```

I wrote some code to implement my solution:

``` Python
def _optimize_stratification(self, y_counts_per_fold, y_cnt, y_counts_per_group, groups_per_fold):
    groups = np.array(range(y_counts_per_group.shape[0]))
    np.random.shuffle(groups)
    for group0 in groups:
        for fold0, groups_ in groups_per_fold.items():
            if group0 in groups_:
                break
        for group1 in groups:
            if group0 == group1:
                continue
            for fold1, groups_ in groups_per_fold.items():
                if group1 in groups_:
                    break
            std_per_class = np.std(y_counts_per_fold / y_cnt.reshape(1, -1),
                                   axis=0)

            y_counts_per_fold[fold0] -= y_counts_per_group[group0]
            y_counts_per_fold[fold1] += y_counts_per_group[group0]
            y_counts_per_fold[fold1] -= y_counts_per_group[group1]
            y_counts_per_fold[fold0] += y_counts_per_group[group1]

            std_per_class_new = np.std(
                y_counts_per_fold / y_cnt.reshape(1, -1), axis=0)
            if np.mean(std_per_class_new) < np.mean(std_per_class):
                groups_per_fold[fold0].remove(group0)
                groups_per_fold[fold1].remove(group1)
                groups_per_fold[fold1].add(group0)
                groups_per_fold[fold0].add(group1)
                if self._check_stratification_ideal(y_counts_per_fold,
                                                    warn=False):
                    # If the stratification is ideal, then looping stops
                    return
                else:
                    # If the stratification is not, it runs again
                    return self._optimize_stratification(self,
                                                         y_counts_per_fold,
                                                         y_counts_per_group,
                                                         y_cnt,
                                                         groups_per_fold)
            else:
                y_counts_per_fold[fold0] += y_counts_per_group[group0]
                y_counts_per_fold[fold1] -= y_counts_per_group[group0]
                y_counts_per_fold[fold1] += y_counts_per_group[group1]
                y_counts_per_fold[fold0] -= y_counts_per_group[group1]

``` 

Regardless, even without any changes to how the code functions and even aside from pursuing ideal splits, I think a warning would be helpful to say that the stratification isn't perfect (this function is also called by the one above):

``` Python
def _check_stratification_ideal(self, y_counts_per_fold, warn=False):
        if y_counts_per_fold.shape[0] > 1:
            first_fold = y_counts_per_fold[0]
            for fold in y_counts_per_fold[1:]:
                if not np.array_equal(first_fold, fold):
                    if warn:
                        warnings.warn('Folds are not perfectly stratified.',
                                  UserWarning)
                    return False
        return True
```

Then, in `_split`, these lines are added to check if the stratification is ideal and if not, to run the optimization:
``` Python
if not self._check_stratification_ideal(y_counts_per_fold, warn=False):
    self._optimize_stratification(y_counts_per_fold, y_cnt, y_counts_per_group,
                 groups_per_fold)
```

### Shuffling bug fix

As I mentioned above, I needed to shift the shuffling to fix a general bug. I commented out:
``` Python
# if self.shuffle:
#    rng.shuffle(y_counts_per_group)
```

And added earlier (the original code uses group_inv to calculate y_counts_per_groups in the subsequent lines):
``` Python
if self.shuffle:
    group2shuffled = np.array(range(len(groups_unique)))
    rng.shuffle(group2shuffled)
    groups_inv = group2shuffled[groups_inv]
```


### Describe alternatives you've considered, if relevant

I think my greedy algorithm finds the best solution in cases where groups all have equal sizes and where Y = 0 or 1. I'm not sure about cases where the groups have unequal sizes or where Y can be outside 0 or 1. The optimization would still make the stratification better, but I'm not sure if it will always find the best solution.

A downside to my solution is that users may not care about ideal stratification, and this function could waste their time. However, I don't think the time cost is too major relative to the time needed to actually fit models. I guess this just depends on how many groups people are using. It may be problematic if the number of groups is in the thousands. Implementing something like this may require an optional parameter asking users whether they want this optimization and if so, for how many loops they want to do.

### Additional context

Recently, a new cross-validation function was proposed (issue #24247), RepeatedStratifiedGroupKFold, which uses StratifiedGroupKFold. Having shuffling working properly is needed for RepeatedStratifiedGroupKFold.

Most of the historic conversations surrounding Stratified Group K-fold have been around biomedical/psychological applications. This is where I'm coming from. Having perfect stratification is sometimes vital in these cases.

I haven't submitted a pull request, but please let me know if I should.

COMMENT:
Very sorry, @glemaitre! This completely slipped my mind and I missed your reply. I will get back to you this weekend and look into the actual cost to run it. 

About your comment, "it means that you can finish up with a shuffle parameter that does not shuffle, isn't it?" Do you mean that in the case where there is a single ideal solution (requires that two groups have the same breakdown of examples), then it will always return the same solution even when shuffle = True? This indeed seems like a downside. However, my solution overall seems like an improvement because optimal stratification seems like a more common use case than "potentially suboptimal stratification but with differences every time it is run".

@katrinedam, I'm not sure I understand your quick-fix solution, although I believe the code I provided would address your issue. My code should detect that the stratification was far from optimal and then swap groups around to make it better.

Labels: New Feature, module:model_selection, Needs Decision - Include Feature