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

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

    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

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

Review comment (suggestion):
```suggestion
    """Check `_weighted_percentile` behaviour is correct when `array` is 2D."""
```

Code after:
@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

    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

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