linting: ruff

Author: cakedev0

src/array_api_extra/_delegation.py

@@ -1055,7 +1055,8 @@ def quantile(
     if xp is None:
         xp = array_namespace(a)
     if is_pydata_sparse_namespace(xp):
-        raise ValueError('no supported')
+        msg = "Sparse backend not supported"
+        raise ValueError(msg)
 
     methods = {"linear", "inverted_cdf", "averaged_inverted_cdf"}
     if method not in methods:
@@ -1096,7 +1097,10 @@ def quantile(
                 msg = "Axis must be specified when shapes of `a` and ̀ weights` differ."
                 raise TypeError(msg)
             if weights.shape != eager_shape(a, axis):
-                msg = "Shape of weights must be consistent with shape of a along specified axis."
+                msg = (
+                    "Shape of weights must be consistent with shape"
+                    " of a along specified axis."
+                )
                 raise ValueError(msg)
         if axis is None and ndim == 2:
             msg = "When weights are provided, axis must be specified when `a` is 2d"
@@ -1119,7 +1123,9 @@ def quantile(
 
     # Delegate where possible.
     if is_numpy_namespace(xp) and nan_policy == "propagate":
-        return xp.quantile(a, q_arr, axis=axis, method=method, keepdims=keepdims, weights=weights)
+        return xp.quantile(
+            a, q_arr, 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 and nan_policy == "propagate"
     if (basic_case and is_jax_namespace(xp)) or is_cupy_namespace(xp):
@@ -1130,6 +1136,12 @@ def quantile(
     # XXX: I'm not sure we want to support dask, it seems uterly slow...
     # Otherwise call our implementation (will sort data)
     return _quantile.quantile(
-        a, q_arr, axis=axis, method=method, keepdims=keepdims,
-        nan_policy=nan_policy, weights=weights, xp=xp
+        a,
+        q_arr,
+        axis=axis,
+        method=method,
+        keepdims=keepdims,
+        nan_policy=nan_policy,
+        weights=weights,
+        xp=xp,
     )

src/array_api_extra/_lib/_quantile.py

@@ -49,10 +49,17 @@ def quantile(  # numpydoc ignore=PR01,RT01
             res = xp.moveaxis(res, axis, 0)
     else:
         weights_arr = xp.asarray(weights, dtype=xp.float64, device=device)
-        average = method == 'averaged_inverted_cdf'
+        average = method == "averaged_inverted_cdf"
         res = _weighted_quantile(
-            a, q, weights_arr, n, axis, average,
-            nan_policy=nan_policy, xp=xp, device=device
+            a,
+            q,
+            weights_arr,
+            n,
+            axis,
+            average,
+            nan_policy=nan_policy,
+            xp=xp,
+            device=device,
         )
 
     # reshaping to conform to doc/other libs' behavior
@@ -72,15 +79,17 @@ def _quantile(  # numpydoc ignore=GL08
     a = xp.sort(a, axis=axis, stable=False)
     mask_nan = xp.any(xp.isnan(a), axis=axis, keepdims=True)
     if xp.any(mask_nan):
-        # propogate NaNs:
+        # propagate NaNs:
         mask = xp.repeat(mask_nan, n, axis=axis)
         a = xp.where(mask, xp.nan, a)
         del mask
 
-    if method == "linear":
-        m = 1 - q
-    else: # method is "inverted_cdf" or "averaged_inverted_cdf"
-        m = xp.asarray(0, dtype=q.dtype)
+    m = (
+        1 - q
+        if method == "linear"
+        # method is "inverted_cdf" or "averaged_inverted_cdf"
+        else xp.asarray(0, dtype=q.dtype)
+    )
 
     jg = q * float(n) + m - 1
 
@@ -90,9 +99,9 @@ def _quantile(  # numpydoc ignore=GL08
     # `̀j` and `jp1` are 1d arrays
 
     g = jg % 1
-    if method == 'inverted_cdf':
+    if method == "inverted_cdf":
         g = xp.astype((g > 0), jg.dtype)
-    elif method == 'averaged_inverted_cdf':
+    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 readonly
@@ -106,8 +115,15 @@ def _quantile(  # numpydoc ignore=GL08
 
 
 def _weighted_quantile(
-    a: Array, q: Array, weights: Array, n: int, axis: int, average: bool, nan_policy: str,
-    xp: ModuleType, device: Device
+    a: Array,
+    q: Array,
+    weights: Array,
+    n: int,
+    axis: int,
+    average: bool,
+    nan_policy: str,
+    xp: ModuleType,
+    device: Device,
 ) -> Array:
     """
     a is expected to be 1d or 2d.
@@ -122,37 +138,45 @@ def _weighted_quantile(
         w = xp.take(weights, sorter)
         return _weighted_quantile_sorted_1d(x, q, w, n, average, nan_policy, xp, device)
 
-    d, = eager_shape(a, axis=0)
+    (d,) = eager_shape(a, axis=0)
     res = []
     for idx in range(d):
         w = weights if weights.ndim == 1 else weights[idx, ...]
         w = xp.take(w, sorter[idx, ...])
         x = xp.take(a[idx, ...], sorter[idx, ...])
-        res.append(_weighted_quantile_sorted_1d(x, q, w, n, average, nan_policy, xp, device))
-    res = xp.stack(res, axis=1)
-    return res
+        res.append(
+            _weighted_quantile_sorted_1d(x, q, w, n, average, nan_policy, xp, device)
+        )
+
+    return xp.stack(res, axis=1)
 
 
 def _weighted_quantile_sorted_1d(
-    x: Array, q: Array, w: Array, n: int, average: bool, nan_policy: str,
-    xp: ModuleType, device: Device
+    x: Array,
+    q: Array,
+    w: Array,
+    n: int,
+    average: bool,
+    nan_policy: str,
+    xp: ModuleType,
+    device: Device,
 ) -> Array:
     if nan_policy == "omit":
-        w = xp.where(xp.isnan(x), 0., w)
+        w = xp.where(xp.isnan(x), 0.0, w)
     elif xp.any(xp.isnan(x)):
         return xp.full(q.shape, xp.nan, dtype=x.dtype, device=device)
     cw = xp.cumulative_sum(w)
     t = cw[-1] * q
-    i = xp.searchsorted(cw, t, side='left')
-    j = xp.searchsorted(cw, t, side='right')
+    i = xp.searchsorted(cw, t, side="left")
+    j = xp.searchsorted(cw, t, side="right")
     i = xp.clip(i, 0, n - 1)
     j = xp.clip(j, 0, 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)
+    i = xp.where(q == 0.0, j, i)
     if average:
         # Ignore trailing `weights=0` observations when `q=1`
-        j = xp.where(q == 1., i, j)
+        j = xp.where(q == 1.0, i, j)
         return (xp.take(x, i) + xp.take(x, j)) / 2
     return xp.take(x, i)

tests/test_funcs.py

@@ -1572,9 +1572,11 @@ def test_against_numpy(self, xp: ModuleType, keepdims: bool):
             xp_assert_close(actual, expected, atol=1e-12)
 
     @pytest.mark.parametrize("keepdims", [True, False])
-    @pytest.mark.parametrize("nan_policy", ["omit", "no_nans", "propagate"])#, #["omit"])#["no_nans", "propagate"])
-    @pytest.mark.parametrize("q_np", [0.5, 0., 1., np.linspace(0, 1, num=11)])
-    def test_weighted_against_numpy(self, xp: ModuleType, keepdims: bool, q_np: Array | float, nan_policy: str):
+    @pytest.mark.parametrize("nan_policy", ["omit", "no_nans", "propagate"])
+    @pytest.mark.parametrize("q_np", [0.5, 0.0, 1.0, np.linspace(0, 1, num=11)])
+    def test_weighted_against_numpy(
+        self, xp: ModuleType, keepdims: bool, q_np: Array | float, nan_policy: str
+    ):
         rng = np.random.default_rng()
         n, d = 10, 20
         a_np = rng.random((n, d))
@@ -1590,7 +1592,7 @@ def test_weighted_against_numpy(self, xp: ModuleType, keepdims: bool, q_np: Arra
 
         a = xp.asarray(a_np, copy=True)
         q = xp.asarray(q_np, copy=True)
-        m = 'inverted_cdf'
+        m = "inverted_cdf"
 
         np_quantile = np.quantile
         if nan_policy == "omit":
@@ -1605,22 +1607,36 @@ def test_weighted_against_numpy(self, xp: ModuleType, keepdims: bool, q_np: Arra
             (rng.integers(0, 2, (n, d)), 1),
         ]:
             with warnings.catch_warnings(record=True) as warning:
-                warnings.filterwarnings("always", "invalid value encountered in divide", RuntimeWarning)
-                warnings.filterwarnings("ignore", "All-NaN slice encountered", RuntimeWarning)
+                divide_msg = "invalid value encountered in divide"
+                warnings.filterwarnings("always", divide_msg, RuntimeWarning)
+                nan_slice_msg = "All-NaN slice encountered"
+                warnings.filterwarnings("ignore", nan_slice_msg, RuntimeWarning)
                 try:
                     expected = np_quantile(  # type: ignore[call-overload]
-                        a_np, np.asarray(q_np),
-                        axis=axis, method=m, weights=w_np, keepdims=keepdims
+                        a_np,
+                        np.asarray(q_np),
+                        axis=axis,
+                        method=m,
+                        weights=w_np,
+                        keepdims=keepdims,
                     )
                 except IndexError:
                     continue
-                if warning:  # this means some weights sum was 0, in this case we skip calling xpx.quantile
+                if warning:
+                    # this means some weights sum was 0
+                    # in this case we skip calling xpx.quantile
                     continue
             expected = xp.asarray(expected)
 
             w = xp.asarray(w_np)
-            actual = quantile(  
-                a, q, axis=axis, method=m, weights=w, keepdims=keepdims, nan_policy=nan_policy
+            actual = quantile(
+                a,
+                q,
+                axis=axis,
+                method=m,
+                weights=w,
+                keepdims=keepdims,
+                nan_policy=nan_policy,
             )
             xp_assert_close(actual, expected, atol=1e-12)