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

Code before:
from pytest import approx

from sklearn._config import config_context
from sklearn.utils._array_api import (
    _convert_to_numpy,
    get_namespace,
    yield_namespace_device_dtype_combinations,
)
from sklearn.utils._array_api import device as array_device
from sklearn.utils.estimator_checks import _array_api_for_tests
from sklearn.utils.fixes import np_version, parse_version
from sklearn.utils.stats import _weighted_percentile


@pytest.mark.parametrize("average", [True, False])
@pytest.mark.parametrize("size", [10, 15])
def test_weighted_percentile_matches_median(size, average):
    """Ensure `_weighted_percentile` matches `median` when expected.

    With unit `sample_weight`, `_weighted_percentile` should match median except
    when `average=False` and the number of samples is odd.
    When number of samples is odd, `_weighted_percentile(average=False)` always falls
    on a single observation (not between 2 values, in which case the lower value would
    be taken) and is thus equal to `np.median`.
    For an even number of samples, `median` gives the average between the 2 middle
    samples, `_weighted_percentile(average=False)` gives the higher (right) sample.
    """
    y = np.arange(size)
    sample_weight = np.ones_like(y)

    score = _weighted_percentile(y, sample_weight, 50, average=average)

    # `_weighted_percentile(average=False)` does not match `median` when n is even
    if size == 10 and average is False:
        assert score != np.median(y)
    else:
        assert score == np.median(y)


# test 2D?
@pytest.mark.parametrize("average", [True, False])
@pytest.mark.parametrize("percentile_rank", [20, 35, 61])
@pytest.mark.parametrize("size", [10, 15])
def test_weighted_percentile_matches_numpy(
    global_random_seed, size, percentile_rank, average
):
    """Check `_weighted_percentile` with unit weights is correct.

    `average=True` results should be the same as `np.percentile`'s
    'averaged_inverted_cdf'.
    `average=False` results should be the same as `np.percentile`'s

Review comment (suggestion):
I think the more concise suggestion below is enough and simpler to understand.

```suggestion
    When number of samples is odd, `_weighted_percentile(average=False)` always falls
    on a single observation and is thus equal to `np.median`.
    For an even number of samples, `median` gives the average between the 2 middle
    samples, `_weighted_percentile(average=False)` gives the higher (right) sample.
```

Code after:
from pytest import approx

from sklearn._config import config_context
from sklearn.utils._array_api import (
    _convert_to_numpy,
    get_namespace,
    yield_namespace_device_dtype_combinations,
)
from sklearn.utils._array_api import device as array_device
from sklearn.utils.estimator_checks import _array_api_for_tests
from sklearn.utils.fixes import np_version, parse_version
from sklearn.utils.stats import _weighted_percentile


@pytest.mark.parametrize("average", [True, False])
@pytest.mark.parametrize("size", [10, 15])
def test_weighted_percentile_matches_median(size, average):
    """Ensure `_weighted_percentile` matches `median` when expected.

    With unit `sample_weight`, `_weighted_percentile` should match the median except
    when `average=False` and the number of samples is even.
    For an even array and `average=False`, `percentile_rank=50` gives the lower
    of the two 'middle' values, that are averaged when calculating the `median`.
    """
    y = np.arange(size)
    sample_weight = np.ones_like(y)

    score = _weighted_percentile(y, sample_weight, 50, average=average)

    # `_weighted_percentile(average=False)` does not match `median` when n is even
    if size % 2 == 0 and average is False:
        assert score != np.median(y)
    else:
        assert approx(score) == np.median(y)


@pytest.mark.parametrize("average", [True, False])
@pytest.mark.parametrize("percentile_rank", [20, 35, 61])
@pytest.mark.parametrize("size", [10, 15])
def test_weighted_percentile_matches_numpy(
    global_random_seed, size, percentile_rank, average
):
    """Check `_weighted_percentile` with unit weights is correct.

    `average=True` results should be the same as `np.percentile`'s
    'averaged_inverted_cdf'.
    `average=False` results should be the same as `np.percentile`'s
    'inverted_cdf'.
    Note `np.percentile` is the same as `np.quantile` except `q` is in range [0, 100].

    We parametrize through different `percentile_rank` and `size` to