Skip to content

Commit f35d3bd

Browse files
authored
BUG: weighted quantile for some zero weights (numpy#27549)
This PR fixed weighted quantiles (and percentiles) for a corner case: * at least one weight is zero * q=0 (0-quantile equals minimum) Then: ``` np.quantile(np.arange(3), 0, weights=[0, 0, 1], method="inverted_cdf") ``` should return 2, the minimum when neglecting zero weight values. Current main returns 0.
1 parent 83c34d5 commit f35d3bd

File tree

2 files changed

+18
-0
lines changed

2 files changed

+18
-0
lines changed

numpy/lib/_function_base_impl.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4821,6 +4821,13 @@ def _quantile(
48214821
# returns 2 instead of 1 because 0.4 is not binary representable.
48224822
if quantiles.dtype.kind == "f":
48234823
cdf = cdf.astype(quantiles.dtype)
4824+
# Weights must be non-negative, so we might have zero weights at the
4825+
# beginning leading to some leading zeros in cdf. The call to
4826+
# np.searchsorted for quantiles=0 will then pick the first element,
4827+
# but should pick the first one larger than zero. We
4828+
# therefore simply set 0 values in cdf to -1.
4829+
if np.any(cdf[0, ...] == 0):
4830+
cdf[cdf == 0] = -1
48244831

48254832
def find_cdf_1d(arr, cdf):
48264833
indices = np.searchsorted(cdf, quantiles, side="left")

numpy/lib/tests/test_function_base.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4008,6 +4008,17 @@ def test_quantile_with_weights_and_axis(self, method):
40084008
)
40094009
assert_allclose(q, q_res)
40104010

4011+
@pytest.mark.parametrize("method", methods_supporting_weights)
4012+
def test_quantile_weights_min_max(self, method):
4013+
# Test weighted quantile at 0 and 1 with leading and trailing zero
4014+
# weights.
4015+
w = [0, 0, 1, 2, 3, 0]
4016+
y = np.arange(6)
4017+
y_min = np.quantile(y, 0, weights=w, method="inverted_cdf")
4018+
y_max = np.quantile(y, 1, weights=w, method="inverted_cdf")
4019+
assert y_min == y[2] # == 2
4020+
assert y_max == y[4] # == 4
4021+
40114022
def test_quantile_weights_raises_negative_weights(self):
40124023
y = [1, 2]
40134024
w = [-0.5, 1]

0 commit comments

Comments
 (0)