Repository: scikit-learn/scikit-learn
File: sklearn/model_selection/tests/test_split.py (Python)

Code before:
            np.array([0] * 9 + [1] * 3),
            np.array([1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 5, 6]),
            np.asarray([[0.75, 0.25], [0.75, 0.25], [0.75, 0.25]]),
        ),
    ],
)
def test_stratified_group_kfold_homogeneous_groups(y, groups, expected):
    sgkf = StratifiedGroupKFold(n_splits=3)
    X = np.ones_like(y).reshape(-1, 1)
    for (train, test), expect_dist in zip(sgkf.split(X, y, groups), expected):
        # check group constraint
        assert np.intersect1d(groups[train], groups[test]).size == 0
        split_dist = np.bincount(y[test]) / len(test)
        assert_allclose(split_dist, expect_dist, atol=0.001)


def test_stratified_group_kfold_shuffle_preserves_stratification():
    # Check StratifiedGroupKFold with shuffle=True preserves stratification:
    # shuffling only affects tie-breaking among groups with identical
    # standard deviation of class distribution
    y = np.array([0] * 12 + [1] * 6)
    X = np.ones((len(y), 1))
    # Groups are arranged so perfect stratification across 3 folds is
    # achievable
    groups = np.array([1, 1, 3, 3, 3, 4, 5, 5, 5, 5, 7, 7, 2, 2, 6, 6, 8, 8])
    distr = np.asarray([2.0 / 3, 1.0 / 3])

    # Run multiple seeds to ensure the property holds regardless of the
    # tie-breaking order among groups with identical std of class distribution
    n_iters = 100
    for seed in range(n_iters):
        sgkf = StratifiedGroupKFold(n_splits=3, shuffle=True, random_state=seed)
        test_sizes = []
        for train, test in sgkf.split(X, y, groups):
            # check group constraint
            assert np.intersect1d(groups[train], groups[test]).size == 0
            # check y distribution
            assert_allclose(np.bincount(y[train]) / len(train), distr, atol=1e-8)
            assert_allclose(np.bincount(y[test]) / len(test), distr, atol=1e-8)
            test_sizes.append(len(test))
        assert np.ptp(test_sizes) <= 1


@pytest.mark.parametrize("cls_distr", [(0.4, 0.6), (0.3, 0.7), (0.2, 0.8), (0.8, 0.2)])
@pytest.mark.parametrize("n_groups", [5, 30, 70])
def test_stratified_group_kfold_against_group_kfold(cls_distr, n_groups):
    # Check that given sufficient amount of samples StratifiedGroupKFold
    # produces better stratified folds than regular GroupKFold
    n_splits = 5
    sgkf = StratifiedGroupKFold(n_splits=n_splits)
    gkf = GroupKFold(n_splits=n_splits)

Review comment (suggestion):
I think this variable could be renamed `expected_class_ratios` to be more explicit.

Code after:
            np.array([0] * 9 + [1] * 3),
            np.array([1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 5, 6]),
            np.asarray([[0.75, 0.25], [0.75, 0.25], [0.75, 0.25]]),
        ),
    ],
)
def test_stratified_group_kfold_homogeneous_groups(y, groups, expected):
    sgkf = StratifiedGroupKFold(n_splits=3)
    X = np.ones_like(y).reshape(-1, 1)
    for (train, test), expect_dist in zip(sgkf.split(X, y, groups), expected):
        # check group constraint
        assert np.intersect1d(groups[train], groups[test]).size == 0
        split_dist = np.bincount(y[test]) / len(test)
        assert_allclose(split_dist, expect_dist, atol=0.001)


def test_stratified_group_kfold_shuffle_preserves_stratification():
    # Check StratifiedGroupKFold with shuffle=True preserves stratification:
    # shuffling only affects tie-breaking among groups with identical
    # standard deviation of class distribution (see #32478)
    y = np.array([0] * 12 + [1] * 6)
    X = np.ones((len(y), 1))
    # Groups are arranged so perfect stratification across 3 folds is
    # achievable
    groups = np.array([1, 1, 3, 3, 3, 4, 5, 5, 5, 5, 7, 7, 2, 2, 6, 6, 8, 8])
    expected_class_ratios = np.asarray([2.0 / 3, 1.0 / 3])

    # Run multiple seeds to ensure the property holds regardless of the
    # tie-breaking order among groups with identical std of class distribution
    n_iters = 100
    for seed in range(n_iters):
        sgkf = StratifiedGroupKFold(n_splits=3, shuffle=True, random_state=seed)
        test_sizes = []
        for train, test in sgkf.split(X, y, groups):
            # check group constraint
            assert np.intersect1d(groups[train], groups[test]).size == 0
            # check y distribution
            assert_allclose(
                np.bincount(y[train]) / len(train), expected_class_ratios, atol=1e-8
            )
            assert_allclose(
                np.bincount(y[test]) / len(test), expected_class_ratios, atol=1e-8
            )
            test_sizes.append(len(test))
        assert np.ptp(test_sizes) <= 1


@pytest.mark.parametrize("cls_distr", [(0.4, 0.6), (0.3, 0.7), (0.2, 0.8), (0.8, 0.2)])
@pytest.mark.parametrize("n_groups", [5, 30, 70])
def test_stratified_group_kfold_against_group_kfold(cls_distr, n_groups):
    # Check that given sufficient amount of samples StratifiedGroupKFold