Repository: scikit-learn/scikit-learn
File: sklearn/utils/tests/test_stats.py (Python)

Code before:

@pytest.mark.parametrize("average", [True, False])
@pytest.mark.parametrize("percentile_rank, expected_value", [(0, 2), (50, 3), (100, 5)])
def test_weighted_percentile_ignores_zero_weight(
    average, percentile_rank, expected_value
):
    """Check leading, trailing and middle 0 weights behave correctly.

    Check that leading zero-weight observations are ignored when `percentile_rank=0`.
    See #20528 for details.
    Check that when `average=True` and the `j+1` ('plus one') index has sample weight
    of 0, it is ignored. Also check that trailing zero weight observations are ignored
    (e.g., when `percentile_rank=100`).
    """
    y = np.array([0, 1, 2, 3, 4, 5, 6])
    sw = np.array([0, 0, 1, 1, 0, 1, 0])

    value = _weighted_percentile(
        np.vstack((y, y)).T, np.vstack((sw, sw)).T, percentile_rank, average=average
    )
    for idx in range(2):
        assert approx(value[idx]) == expected_value


def test_weighted_percentile_average_zero_weight_plateau():
    """Check average percentile ignores zero-weight plateaus."""
    score_without_zeros = _weighted_percentile([1, 3], [3, 3], average=True)
    score_with_zeros = _weighted_percentile([1, 2, 3], [3, 0, 3], average=True)
    assert approx(score_without_zeros) == score_with_zeros


@pytest.mark.parametrize("average", [True, False])
@pytest.mark.parametrize("percentile_rank", [20, 35, 50, 61])
def test_weighted_percentile_frequency_weight_semantics(
    global_random_seed, percentile_rank, average
):
    """Check integer weights give the same result as repeating values."""
    rng = np.random.RandomState(global_random_seed)
    x = rng.randint(20, size=10)
    weights = rng.choice(5, size=10)

    x_repeated = np.repeat(x, weights)
    percentile_weights = _weighted_percentile(
        x, weights, percentile_rank, average=average
    )
    percentile_repeated = _weighted_percentile(
        x_repeated, np.ones_like(x_repeated), percentile_rank, average=average
    )
    assert percentile_weights == approx(percentile_repeated)
    # Also check `percentile_rank=50` matches `median`
    if percentile_rank == 50 and average:

Review comment (suggestion):
This seems more accurate?
```suggestion
    """Check zero weights just before `max_index` handled correctly."""
```


Code after:

@pytest.mark.parametrize("average", [True, False])
@pytest.mark.parametrize("percentile_rank, expected_value", [(0, 2), (50, 3), (100, 5)])
def test_weighted_percentile_ignores_zero_weight(
    average, percentile_rank, expected_value
):
    """Check leading, trailing and middle 0 weights behave correctly.

    Check that leading zero-weight observations are ignored when `percentile_rank=0`.
    See #20528 for details.
    Check that when `average=True` and the `j+1` ('plus one') index has sample weight
    of 0, it is ignored. Also check that trailing zero weight observations are ignored
    (e.g., when `percentile_rank=100`).
    """
    y = np.array([0, 1, 2, 3, 4, 5, 6])
    sw = np.array([0, 0, 1, 1, 0, 1, 0])

    value = _weighted_percentile(
        np.vstack((y, y)).T, np.vstack((sw, sw)).T, percentile_rank, average=average
    )
    for idx in range(2):
        assert approx(value[idx]) == expected_value


def test_weighted_percentile_average_zero_weight_plateau():
    """Check zero weights just before `max_index` handled correctly."""
    score_without_zeros = _weighted_percentile([1, 3], [3, 3], average=True)
    score_with_zeros = _weighted_percentile([1, 2, 3], [3, 0, 3], average=True)
    assert approx(score_without_zeros) == score_with_zeros


@pytest.mark.parametrize("average", [True, False])
@pytest.mark.parametrize("percentile_rank", [20, 35, 50, 61])
def test_weighted_percentile_frequency_weight_semantics(
    global_random_seed, percentile_rank, average
):
    """Check integer weights give the same result as repeating values."""
    rng = np.random.RandomState(global_random_seed)
    x = rng.randint(20, size=10)
    weights = rng.choice(5, size=10)

    x_repeated = np.repeat(x, weights)
    percentile_weights = _weighted_percentile(
        x, weights, percentile_rank, average=average
    )
    percentile_repeated = _weighted_percentile(
        x_repeated, np.ones_like(x_repeated), percentile_rank, average=average
    )
    assert percentile_weights == approx(percentile_repeated)
    # Also check `percentile_rank=50` matches `median`
    if percentile_rank == 50 and average: