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

Code before:
            [np.nan, 2],
            [np.nan, np.nan],
        ]
    )
    weights = np.ones_like(array)
    percentile_rank = 90

    values = _weighted_percentile(array, weights, percentile_rank)

    # The percentile of the second column should be `5` even though there are many nan
    # values present; the percentile of the first column can only be nan, since there
    # are no other possible values:
    assert np.array_equal(values, np.array([np.nan, 5]), equal_nan=True)


@pytest.mark.skipif(
    np_version < parse_version("2.0"),
    reason="np.quantile only accepts weights since version 2.0",
)
@pytest.mark.parametrize("percentile", [66, 10, 50])
@pytest.mark.parametrize("average", [False, True])
@pytest.mark.parametrize("uniform_weight", [False, True])
def test_weighted_percentile_like_numpy_quantile(
    percentile, average, uniform_weight, global_random_seed
):
    """Check `_weighted_percentile` equivalent to `np.quantile` with weights."""
    # TODO: remove the following skip once no longer applicable.
    if average and not uniform_weight:
        pytest.skip(
            "np.quantile does not support weights with method='averaged_inverted_cdf'"
        )

    rng = np.random.RandomState(global_random_seed)
    array = rng.rand(10, 100)
    if uniform_weight:
        sample_weight = np.ones_like(array) * rng.randint(1, 6, size=1)
    else:
        sample_weight = rng.randint(1, 6, size=(10, 100))

    percentile_weighted_percentile = _weighted_percentile(
        array, sample_weight, percentile, average=average
    )
    percentile_numpy_quantile = np.quantile(
        array,
        percentile / 100,
        weights=sample_weight if not uniform_weight else None,
        method="averaged_inverted_cdf" if average else "inverted_cdf",
        axis=0,
    )

    assert_array_equal(percentile_weighted_percentile, percentile_numpy_quantile)

Review comment (suggestion):
```suggestion
    """Check `_weighted_percentile` is equivalent to `np.quantile` with weights."""
```

Code after:
            [np.nan, 2],
            [np.nan, np.nan],
        ]
    )
    weights = np.ones_like(array)
    percentile_rank = 90

    values = _weighted_percentile(array, weights, percentile_rank)

    # The percentile of the second column should be `5` even though there are many nan
    # values present; the percentile of the first column can only be nan, since there
    # are no other possible values:
    assert np.array_equal(values, np.array([np.nan, 5]), equal_nan=True)


@pytest.mark.skipif(
    np_version < parse_version("2.0"),
    reason="np.quantile only accepts weights since version 2.0",
)
@pytest.mark.parametrize("percentile", [66, 10, 50])
@pytest.mark.parametrize("average", [False, True])
@pytest.mark.parametrize("uniform_weight", [False, True])
def test_weighted_percentile_like_numpy_quantile(
    percentile, average, uniform_weight, global_random_seed
):
    """Check `_weighted_percentile` is equivalent to `np.quantile` with weights."""
    # TODO: remove the following skip once no longer applicable.
    if average and not uniform_weight:
        pytest.skip(
            "np.quantile does not support weights with method='averaged_inverted_cdf'"
        )

    rng = np.random.RandomState(global_random_seed)
    array = rng.rand(10, 100)
    if uniform_weight:
        sample_weight = np.ones_like(array) * rng.randint(1, 6, size=1)
    else:
        sample_weight = rng.randint(1, 6, size=(10, 100))

    percentile_weighted_percentile = _weighted_percentile(
        array, sample_weight, percentile, average=average
    )
    percentile_numpy_quantile = np.quantile(
        array,
        percentile / 100,
        weights=sample_weight if not uniform_weight else None,
        method="averaged_inverted_cdf" if average else "inverted_cdf",
        axis=0,
    )

    assert_array_equal(percentile_weighted_percentile, percentile_numpy_quantile)