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

Code before:
    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 constant does not change 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
    )
    assert percentile == approx(percentile_multiplier)


@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


Review comment (suggestion):
```suggestion
    Note: scale invariance does not always hold when multiplying by a
```

Code after:
    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
    )
    assert percentile == approx(percentile_multiplier)


@pytest.mark.parametrize("average", [True, False])
def test_weighted_percentile_2d(global_random_seed, average):
    """Check `_weighted_percentile` behaviour is 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