WIP: adding support for weights

Author: cakedev0

src/array_api_extra/_delegation.py

@@ -905,6 +905,7 @@ def quantile(
     method: str = "linear",
     keepdims: bool = False,
     *,
+    weights: Array | None = None,
     xp: ModuleType | None = None,
 ) -> Array:
     """
@@ -943,6 +944,16 @@ def quantile(
         the result as dimensions with size one. With this option, the
         result will broadcast correctly against the original array `a`.
 
+    weights : array_like, optional
+        An array of weights associated with the values in `a`. Each value in
+        `a` contributes to the quantile according to its associated weight.
+        The weights array can either be 1-D (in which case its length must be
+        the size of `a` along the given axis) or of the same shape as `a`.
+        If `weights=None`, then all data in `a` are assumed to have a
+        weight equal to one.
+        Only `method="inverted_cdf"` or `method="averaged_inverted_cdf"`
+        support weights. See the notes for more details.
+
     xp : array_namespace, optional
         The standard-compatible namespace for `a` and `q`. Default: infer.
 
@@ -1040,7 +1051,7 @@ def quantile(
        "Sample quantiles in statistical packages,"
        The American Statistician, 50(4), pp. 361-365, 1996
     """
-    methods = {"linear"}
+    methods = {"linear", "inverted_cdf", "averaged_inverted_cdf"}
 
     if method not in methods:
         msg = f"`method` must be one of {methods}"
@@ -1084,12 +1095,12 @@ def quantile(
 
     # Delegate where possible.
     if is_numpy_namespace(xp):
+        return xp.quantile(a, q, axis=axis, method=method, keepdims=keepdims, weights=weights)
+    # No delegation for dask: I couldn't make it work
+    basic_case = method == "linear" and weights is None
+    if (basic_case and is_jax_namespace(xp)) or is_cupy_namespace(xp):
         return xp.quantile(a, q, axis=axis, method=method, keepdims=keepdims)
-    # No delegating for dask: I couldn't make it work
-    is_linear = method == "linear"
-    if (is_linear and is_jax_namespace(xp)) or is_cupy_namespace(xp):
-        return xp.quantile(a, q, axis=axis, method=method, keepdims=keepdims)
-    if is_linear and is_torch_namespace(xp):
+    if basic_case and is_torch_namespace(xp):
         return xp.quantile(a, q, dim=axis, interpolation=method, keepdim=keepdims)
 
     # Otherwise call our implementation (will sort data)

src/array_api_extra/_lib/_quantile.py

@@ -9,24 +9,22 @@
 
 def quantile(  # numpydoc ignore=PR01,RT01
     a: Array,
-    q: Array | float,
+    q: Array,
     /,
-    method: str = "linear",  # noqa: ARG001
+    method: str = "linear",
     axis: int | None = None,
     keepdims: bool = False,
     *,
+    weights: Array | None = None,
     xp: ModuleType,
 ) -> Array:
     """See docstring in `array_api_extra._delegation.py`."""
     device = get_device(a)
-    floating_dtype = xp.float64  # xp.result_type(a, xp.asarray(q))
-    a = xp.asarray(a, dtype=floating_dtype, device=device)
     a_shape = list(a.shape)
-    p: Array = xp.asarray(q, dtype=floating_dtype, device=device)
 
-    q_scalar = p.ndim == 0
+    q_scalar = q.ndim == 0
     if q_scalar:
-        p = xp.reshape(p, (1,))
+        q = xp.reshape(q, (1,))
 
     axis_none = axis is None
     a_ndim = a.ndim
@@ -42,33 +40,41 @@ def quantile(  # numpydoc ignore=PR01,RT01
     if n == 0:
         a_shape[axis] = 1
         n = 1
-        a = xp.full(tuple(a_shape), xp.nan, dtype=floating_dtype, device=device)
+        a = xp.full(tuple(a_shape), xp.nan, dtype=a.dtype, device=device)
 
-    a = xp.sort(a, axis=axis, stable=False)
+    if weights is None:
+        res = _quantile(a, q, float(n), axis, method, xp)
+    else:
+        average = method == 'averaged_inverted_cdf'
+        res = _weighted_quantile(a, q, weights, n, axis, average, xp)
     # to support weights, the main thing would be to
     # argsort a, and then use it to sort a and w.
     # The hard part will be dealing with 0-weights and NaNs
     # But maybe a proper use of searchsorted + left/right side will work?
 
-    res = _quantile_hf(a, p, float(n), axis, xp)
-
     # reshaping to conform to doc/other libs' behavior
     if axis_none:
         if keepdims:
-            res = xp.reshape(res, p.shape + (1,) * a_ndim)
+            res = xp.reshape(res, q.shape + (1,) * a_ndim)
     else:
         res = xp.moveaxis(res, axis, 0)
         if keepdims:
             a_shape[axis] = 1
-            res = xp.reshape(res, p.shape + tuple(a_shape))
+            res = xp.reshape(res, q.shape + tuple(a_shape))
 
     return res[0, ...] if q_scalar else res
 
 
-def _quantile_hf(  # numpydoc ignore=GL08
-    a: Array, q: Array, n: float, axis: int, xp: ModuleType
+def _quantile(  # numpydoc ignore=GL08
+    a: Array, q: Array, n: float, axis: int, method: str, xp: ModuleType
 ) -> Array:
-    m = 1 - q
+    a = xp.sort(a, axis=axis, stable=False)
+
+    if method == "linear":
+        m = 1 - q       
+    else: # method is "inverted_cdf" or "averaged_inverted_cdf"
+        m = 0
+
     jg = q * n + m - 1
 
     j = jg // 1
@@ -77,6 +83,11 @@ def _quantile_hf(  # numpydoc ignore=GL08
     # `̀j` and `jp1` are 1d arrays
 
     g = jg % 1
+    if method == 'inverted_cdf':
+        g = xp.astype((g > 0), jg.dtype)
+    elif method == 'averaged_inverted_cdf':
+        g = (1 + xp.astype((g > 0), jg.dtype)) / 2
+
     g = xp.where(j < 0, 0, g)  # equivalent to g[j < 0] = 0, but works with strictest
     new_g_shape = [1] * a.ndim
     new_g_shape[axis] = g.shape[0]
@@ -85,3 +96,39 @@ def _quantile_hf(  # numpydoc ignore=GL08
     return (1 - g) * xp.take(a, xp.astype(j, xp.int64), axis=axis) + g * xp.take(
         a, xp.astype(jp1, xp.int64), axis=axis
     )
+
+
+def _weighted_quantile(a: Array, q: Array, weights: Array, n: int, axis, average: bool, xp: ModuleType):
+    a = xp.moveaxis(a, axis, -1)
+    sorter = xp.argsort(a, axis=-1, stable=False)
+    a = xp.take_along_axis(a, sorter, axis=-1)
+
+    if a.ndim == 1:
+        return _weighted_quantile_sorted_1d(a, q, weights, n, )
+
+    d, = eager_shape(a, axis=0)
+    res = xp.empty((q.shape[0], d))
+    for idx in range(d):
+        w = weights if weights.ndim == 1 else weights[idx, ...]
+        w = xp.take(w, sorter[idx, ...])
+        res[..., idx] = _weighted_quantile_sorted_1d(a[idx, ...], q, w, n, average)
+    return res
+
+
+def _weighted_quantile_sorted_1d(a, q, w, n, average: bool, xp: ModuleType):
+    cw = xp.cumsum(w)
+    t = cw[-1] * q
+    i = xp.searchsorted(cw, t)
+    j = xp.searchsorted(cw, t, side='right')
+    i = xp.minimum(i, float(n - 1))
+    j = xp.minimum(j, float(n - 1))
+
+    # Ignore leading `weights=0` observations when `q=0`
+    # see https://github.com/scikit-learn/scikit-learn/pull/20528
+    i = xp.where(q == 0., j, i)   
+    if average:
+        # Ignore trailing `weights=0` observations when `q=1`
+        j = xp.where(q == 1., i, j)
+        return (xp.take(a, i) + xp.take(a, j)) / 2
+    else:
+        return xp.take(a, i)

tests/test_funcs.py

@@ -1585,7 +1585,7 @@ def test_2d_axis_keepdims(self, xp: ModuleType):
 
     def test_methods(self, xp: ModuleType):
         x = xp.asarray([1, 2, 3, 4, 5])
-        methods = ["linear"]  # "hazen", "weibull"]
+        methods = ["linear", "inverted_cdf", "averaged_inverted_cdf"]
         for method in methods:
             actual = quantile(x, 0.5, method=method)
             # All methods should give reasonable results