Formatting

Author: betatim

src/array_api_extra/_delegation.py

@@ -270,18 +270,18 @@ def quantile(
         Probability or sequence of probabilities of the quantiles to compute.
         Values must be between 0 and 1 (inclusive). Must have length 1 along
         `axis` unless ``keepdims=True``.
-    method : str, default: 'linear'
-        The method to use for estimating the quantile. The available options are:
-        'inverted_cdf', 'averaged_inverted_cdf', 'closest_observation',
-        'interpolated_inverted_cdf', 'hazen', 'weibull', 'linear' (default),
-        'median_unbiased', 'normal_unbiased', 'harrell-davis'.
     axis : int or None, default: None
         Axis along which the quantiles are computed. ``None`` ravels both `x`
         and `q` before performing the calculation.
     keepdims : bool, optional
         If this is set to True, the axes which are reduced are left in the
         result as dimensions with size one. With this option, the result will
         broadcast correctly against the original array `x`.
+    method : str, default: 'linear'
+        The method to use for estimating the quantile. The available options are:
+        'inverted_cdf', 'averaged_inverted_cdf', 'closest_observation',
+        'interpolated_inverted_cdf', 'hazen', 'weibull', 'linear' (default),
+        'median_unbiased', 'normal_unbiased', 'harrell-davis'.
     xp : array_namespace, optional
         The standard-compatible namespace for `x` and `q`. Default: infer.
 
@@ -306,9 +306,12 @@ def quantile(
     try:
         import scipy
         from packaging import version
-        # The quantile function in scipy 1.17 supports array API directly, no need to delegate
+
+        # The quantile function in scipy 1.17 supports array API directly, no need
+        # to delegate
         if version.parse(scipy.__version__) >= version.parse("1.17"):
             from scipy.stats import quantile as scipy_quantile
+
             return scipy_quantile(x, p=q, axis=axis, keepdims=keepdims, method=method)
     except (ImportError, AttributeError):
         pass

src/array_api_extra/_lib/_funcs.py

@@ -1004,147 +1004,162 @@ def quantile(
     """See docstring in `array_api_extra._delegation.py`."""
     if xp is None:
         xp = array_namespace(x, q)
-    
+
     # Convert q to array if it's a scalar
-    q_is_scalar = isinstance(q, (int, float))
+    q_is_scalar = isinstance(q, int | float)
     if q_is_scalar:
         q = xp.asarray(q, dtype=xp.float64, device=_compat.device(x))
-    
+
     # Validate inputs
     if not xp.isdtype(x.dtype, ("integral", "real floating")):
-        raise ValueError("`x` must have real dtype.")
+        raise ValueError("`x` must have real dtype.")  # noqa: EM101
     if not xp.isdtype(q.dtype, "real floating"):
-        raise ValueError("`q` must have real floating dtype.")
-    
+        raise ValueError("`q` must have real floating dtype.")  # noqa: EM101
+
     # Promote to common dtype
     x = xp.astype(x, xp.float64)
     q = xp.astype(q, xp.float64)
     q = xp.asarray(q, device=_compat.device(x))
-    
+
     dtype = x.dtype
     axis_none = axis is None
     ndim = max(x.ndim, q.ndim)
-    
+
     if axis_none:
         x = xp.reshape(x, (-1,))
         q = xp.reshape(q, (-1,))
         axis = 0
     elif not isinstance(axis, int):
-        raise ValueError("`axis` must be an integer or None.")
+        raise ValueError("`axis` must be an integer or None.")  # noqa: EM101
     elif axis >= ndim or axis < -ndim:
-        raise ValueError("`axis` is not compatible with the shapes of the inputs.")
+        raise ValueError("`axis` is not compatible with the shapes of the inputs.")  # noqa: EM101
     else:
         axis = int(axis)
-    
+
     # Validate method
     methods = {
-        'inverted_cdf', 'averaged_inverted_cdf', 'closest_observation',
-        'hazen', 'interpolated_inverted_cdf', 'linear', 'median_unbiased',
-        'normal_unbiased', 'weibull', 'harrell-davis'
+        "inverted_cdf",
+        "averaged_inverted_cdf",
+        "closest_observation",
+        "hazen",
+        "interpolated_inverted_cdf",
+        "linear",
+        "median_unbiased",
+        "normal_unbiased",
+        "weibull",
+        "harrell-davis",
     }
     if method not in methods:
-        raise ValueError(f"`method` must be one of {methods}")
-    
+        raise ValueError(f"`method` must be one of {methods}")  # noqa: EM102
+
     # Handle keepdims parameter
     if keepdims not in {None, True, False}:
-        raise ValueError("If specified, `keepdims` must be True or False.")
-    
+        raise ValueError("If specified, `keepdims` must be True or False.")  # noqa: EM101
+
     # Handle empty arrays
     if x.shape[axis] == 0:
         shape = list(x.shape)
         shape[axis] = 1
         x = xp.full(shape, xp.nan, dtype=dtype, device=_compat.device(x))
-    
+
     # Sort the data
     y = xp.sort(x, axis=axis)
-    
+
     # Move axis to the end for easier processing
     y = xp.moveaxis(y, axis, -1)
     if not (q_is_scalar or q.ndim == 0):
         q = xp.moveaxis(q, axis, -1)
-    
+
     # Get the number of elements along the axis
     n = xp.asarray(y.shape[-1], dtype=dtype, device=_compat.device(y))
-    
+
     # Apply quantile calculation based on method
-    if method in {'inverted_cdf', 'averaged_inverted_cdf', 'closest_observation',
-                  'hazen', 'interpolated_inverted_cdf', 'linear', 'median_unbiased',
-                  'normal_unbiased', 'weibull'}:
+    if method in {
+        "inverted_cdf",
+        "averaged_inverted_cdf",
+        "closest_observation",
+        "hazen",
+        "interpolated_inverted_cdf",
+        "linear",
+        "median_unbiased",
+        "normal_unbiased",
+        "weibull",
+    }:
         res = _quantile_hf(y, q, n, method, xp)
-    elif method == 'harrell-davis':
+    elif method == "harrell-davis":
         res = _quantile_hd(y, q, n, xp)
     else:
-        raise ValueError(f"Unknown method: {method}")
-    
+        raise ValueError(f"Unknown method: {method}")  # noqa: EM102
+
     # Handle NaN output for invalid q values
     p_mask = (q > 1) | (q < 0) | xp.isnan(q)
     if xp.any(p_mask):
         res = xp.asarray(res, copy=True)
         res = at(res, p_mask).set(xp.nan)
-    
+
     # Reshape per axis/keepdims
     if axis_none and keepdims:
         shape = (1,) * (ndim - 1) + res.shape
         res = xp.reshape(res, shape)
         axis = -1
-    
+
     # Move axis back to original position
     res = xp.moveaxis(res, -1, axis)
-    
+
     # Handle keepdims
     if not keepdims and res.shape[axis] == 1:
         res = xp.squeeze(res, axis=axis)
-    
+
     # For scalar q, ensure we return a scalar result
-    if q_is_scalar:
-        if hasattr(res, 'shape') and res.shape != ():
-            res = res[()]
-    
+    if q_is_scalar and hasattr(res, "shape") and res.shape != ():
+        res = res[()]
+
     return res
 
 
 def _quantile_hf(y: Array, p: Array, n: Array, method: str, xp: ModuleType) -> Array:
     """Helper function for Hyndman-Fan quantile methods."""
     ms = {
-        'inverted_cdf': 0,
-        'averaged_inverted_cdf': 0,
-        'closest_observation': -0.5,
-        'interpolated_inverted_cdf': 0,
-        'hazen': 0.5,
-        'weibull': p,
-        'linear': 1 - p,
-        'median_unbiased': p/3 + 1/3,
-        'normal_unbiased': p/4 + 3/8
+        "inverted_cdf": 0,
+        "averaged_inverted_cdf": 0,
+        "closest_observation": -0.5,
+        "interpolated_inverted_cdf": 0,
+        "hazen": 0.5,
+        "weibull": p,
+        "linear": 1 - p,
+        "median_unbiased": p / 3 + 1 / 3,
+        "normal_unbiased": p / 4 + 3 / 8,
     }
     m = ms[method]
-    
+
     jg = p * n + m - 1
     j = xp.astype(jg // 1, xp.int64)  # Convert to integer
     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
-    elif method == 'closest_observation':
-        g = (1 - xp.astype((g == 0) & (j % 2 == 1), jg.dtype))
-    if method in {'inverted_cdf', 'averaged_inverted_cdf', 'closest_observation'}:
+    elif method == "closest_observation":
+        g = 1 - xp.astype((g == 0) & (j % 2 == 1), jg.dtype)
+    if method in {"inverted_cdf", "averaged_inverted_cdf", "closest_observation"}:
         g = xp.asarray(g)
         g = at(g, jg < 0).set(0)
         g = at(g, j < 0).set(0)
     j = xp.clip(j, 0, n - 1)
     jp1 = xp.clip(j + 1, 0, n - 1)
-    
+
     # Broadcast indices to match y shape except for the last axis
     if y.ndim > 1:
         # Create broadcast shape for indices
-        broadcast_shape = list(y.shape[:-1]) + [1]
+        broadcast_shape = list(y.shape[:-1]) + [1]  # noqa: RUF005
         j = xp.broadcast_to(j, broadcast_shape)
         jp1 = xp.broadcast_to(jp1, broadcast_shape)
         g = xp.broadcast_to(g, broadcast_shape)
-    
-    return ((1 - g) * xp.take_along_axis(y, j, axis=-1) +
-            g * xp.take_along_axis(y, jp1, axis=-1))
+
+    return (1 - g) * xp.take_along_axis(y, j, axis=-1) + g * xp.take_along_axis(
+        y, jp1, axis=-1
+    )
 
 
 def _quantile_hd(y: Array, p: Array, n: Array, xp: ModuleType) -> Array:

tests/test_funcs.py

@@ -1193,7 +1193,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", "hazen", "weibull"]
         for method in methods:
             actual = quantile(x, 0.5, method=method)
             # All methods should give reasonable results
@@ -1205,7 +1205,7 @@ def test_edge_cases(self, xp: ModuleType):
         actual = quantile(x, 0.0)
         expect = xp.asarray(1.0)
         xp_assert_close(actual, expect)
-        
+
         # q = 1 should give maximum
         actual = quantile(x, 1.0)
         expect = xp.asarray(5.0)
@@ -1216,7 +1216,7 @@ def test_invalid_q(self, xp: ModuleType):
         # q > 1 should return NaN
         actual = quantile(x, 1.5)
         assert xp.isnan(actual)
-        
+
         # q < 0 should return NaN
         actual = quantile(x, -0.5)
         assert xp.isnan(actual)