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

Code before:
    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


@pytest.mark.parametrize("average", [True, False])
@pytest.mark.parametrize("percentile_rank", [20, 35, 61])
def test_weighted_median_frequency_weights(
    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)


@pytest.mark.parametrize("average", [True, False])
def test_weighted_percentile_2d(global_random_seed, average):
    """Check `_weighted_percentile` behaviour correct when `array` is 2D."""
    # Check for when array 2D and sample_weight 1D
    rng = np.random.RandomState(global_random_seed)
    x1 = rng.randint(10, size=10)
    w1 = rng.choice(5, size=10)

    x2 = rng.randint(20, size=10)
    x_2d = np.vstack((x1, x2)).T

    w_median = _weighted_percentile(x_2d, w1, average=average)
    p_axis_0 = [
        _weighted_percentile(x_2d[:, i], w1, average=average)
        for i in range(x_2d.shape[1])
    ]
    assert_allclose(w_median, p_axis_0)

    # Check when array and sample_weight both 2D
    w2 = rng.choice(5, size=10)
    w_2d = np.vstack((w1, w2)).T

Review comment (suggestion):
Please (re-)include the `percentile_rank=50` case and checks that it also matches the results of `np.median(x_repeated)` when `average is True`.

Code after:

    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


@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:
        assert percentile_weights == approx(np.median(x_repeated))


@pytest.mark.parametrize("constant", [5, 8])
@pytest.mark.parametrize("average", [True, False])
@pytest.mark.parametrize("percentile_rank", [20, 35, 50, 61])
def test_weighted_percentile_constant_multiplier(
    global_random_seed, percentile_rank, average, constant
):
    """Check multiplying weights by a constant does not change the result.

    Note scale invariance does not always hold when multiplying by a
    float due to cumulative sum numerical error (which grows proportional to n).
    """
    rng = np.random.RandomState(global_random_seed)
    x = rng.randint(20, size=20)
    weights = rng.choice(5, size=20)
    weights_multiplied = weights * constant

    percentile = _weighted_percentile(x, weights, percentile_rank, average=average)
    percentile_multiplier = _weighted_percentile(
        x, weights_multiplied, percentile_rank, average=average
    )