Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3from typing import TYPE_CHECKING
5import numpy as np
7from pandas.core.dtypes.missing import (
8 isna,
9 na_value_for_dtype,
10)
12if TYPE_CHECKING:
13 from pandas._typing import (
14 ArrayLike,
15 Scalar,
16 npt,
17 )
20def quantile_compat(
21 values: ArrayLike, qs: npt.NDArray[np.float64], interpolation: str
22) -> ArrayLike:
23 """
24 Compute the quantiles of the given values for each quantile in `qs`.
26 Parameters
27 ----------
28 values : np.ndarray or ExtensionArray
29 qs : np.ndarray[float64]
30 interpolation : str
32 Returns
33 -------
34 np.ndarray or ExtensionArray
35 """
36 if isinstance(values, np.ndarray):
37 fill_value = na_value_for_dtype(values.dtype, compat=False)
38 mask = isna(values)
39 return quantile_with_mask(
40 values,
41 mask,
42 fill_value, # pyright: ignore[reportArgumentType]
43 qs,
44 interpolation,
45 )
46 else:
47 return values._quantile(qs, interpolation)
50def quantile_with_mask(
51 values: np.ndarray,
52 mask: npt.NDArray[np.bool_],
53 fill_value: Scalar,
54 qs: npt.NDArray[np.float64],
55 interpolation: str,
56) -> np.ndarray:
57 """
58 Compute the quantiles of the given values for each quantile in `qs`.
60 Parameters
61 ----------
62 values : np.ndarray
63 For ExtensionArray, this is _values_for_factorize()[0]
64 mask : np.ndarray[bool]
65 mask = isna(values)
66 For ExtensionArray, this is computed before calling _value_for_factorize
67 fill_value : Scalar
68 The value to interpret fill NA entries with
69 For ExtensionArray, this is _values_for_factorize()[1]
70 qs : np.ndarray[float64]
71 interpolation : str
72 Type of interpolation
74 Returns
75 -------
76 np.ndarray
78 Notes
79 -----
80 Assumes values is already 2D. For ExtensionArray this means np.atleast_2d
81 has been called on _values_for_factorize()[0]
83 Quantile is computed along axis=1.
84 """
85 assert values.shape == mask.shape
86 if values.ndim == 1:
87 # unsqueeze, operate, re-squeeze
88 values = np.atleast_2d(values)
89 mask = np.atleast_2d(mask)
90 res_values = quantile_with_mask(values, mask, fill_value, qs, interpolation)
91 return res_values[0]
93 assert values.ndim == 2
95 is_empty = values.shape[1] == 0
97 if is_empty:
98 # create the array of na_values
99 # 2d len(values) * len(qs)
100 flat = np.full(len(qs), fill_value)
101 result = np.repeat(flat, len(values)).reshape(len(values), len(qs))
102 else:
103 result = _nanquantile(
104 values,
105 qs,
106 na_value=fill_value,
107 mask=mask,
108 interpolation=interpolation,
109 )
111 result = np.asarray(result)
112 result = result.T
114 return result
117def _nanquantile_1d(
118 values: np.ndarray,
119 mask: npt.NDArray[np.bool_],
120 qs: npt.NDArray[np.float64],
121 na_value: Scalar,
122 interpolation: str,
123) -> Scalar | np.ndarray:
124 """
125 Wrapper for np.quantile that skips missing values, specialized to
126 1-dimensional case.
128 Parameters
129 ----------
130 values : array over which to find quantiles
131 mask : ndarray[bool]
132 locations in values that should be considered missing
133 qs : np.ndarray[float64] of quantile indices to find
134 na_value : scalar
135 value to return for empty or all-null values
136 interpolation : str
138 Returns
139 -------
140 quantiles : scalar or array
141 """
142 # mask is Union[ExtensionArray, ndarray]
143 values = values[~mask]
145 if len(values) == 0:
146 # Can't pass dtype=values.dtype here bc we might have na_value=np.nan
147 # with values.dtype=int64 see test_quantile_empty
148 # equiv: 'np.array([na_value] * len(qs))' but much faster
149 return np.full(len(qs), na_value)
151 return np.quantile(
152 values,
153 qs,
154 # error: No overload variant of "percentile" matches argument
155 # types "ndarray[Any, Any]", "ndarray[Any, dtype[floating[_64Bit]]]"
156 # , "Dict[str, str]" [call-overload]
157 method=interpolation, # type: ignore[call-overload]
158 )
161def _nanquantile(
162 values: np.ndarray,
163 qs: npt.NDArray[np.float64],
164 *,
165 na_value: Scalar,
166 mask: npt.NDArray[np.bool_],
167 interpolation: str,
168) -> np.ndarray:
169 """
170 Wrapper for np.quantile that skips missing values.
172 Parameters
173 ----------
174 values : np.ndarray[ndim=2] over which to find quantiles
175 qs : np.ndarray[float64] of quantile indices to find
176 na_value : scalar
177 value to return for empty or all-null values
178 mask : np.ndarray[bool]
179 locations in values that should be considered missing
180 interpolation : str
182 Returns
183 -------
184 quantiles : scalar or array
185 """
187 if values.dtype.kind in "mM":
188 # need to cast to integer to avoid rounding errors in numpy
189 result = _nanquantile(
190 values.view("i8"),
191 qs=qs,
192 na_value=na_value.view("i8"), # type: ignore[union-attr, arg-type]
193 mask=mask,
194 interpolation=interpolation,
195 )
197 # Note: we have to do `astype` and not view because in general we
198 # have float result at this point, not i8
199 return result.astype(values.dtype)
201 if mask.any():
202 # Caller is responsible for ensuring mask shape match
203 assert mask.shape == values.shape
204 result_list = [
205 _nanquantile_1d(val, m, qs, na_value, interpolation=interpolation)
206 for (val, m) in zip(list(values), list(mask), strict=True)
207 ]
208 if values.dtype.kind == "f":
209 # preserve itemsize
210 result = np.asarray(result_list, dtype=values.dtype).T
211 else:
212 result = np.asarray(result_list).T
213 if (
214 result.dtype != values.dtype
215 and not mask.all()
216 and (result == result.astype(values.dtype, copy=False)).all()
217 ):
218 # mask.all() will never get cast back to int
219 # e.g. values id integer dtype and result is floating dtype,
220 # only cast back to integer dtype if result values are all-integer.
221 result = result.astype(values.dtype, copy=False)
222 return result
223 else:
224 return np.quantile(
225 values,
226 qs,
227 axis=1,
228 # error: No overload variant of "percentile" matches argument types
229 # "ndarray[Any, Any]", "ndarray[Any, dtype[floating[_64Bit]]]",
230 # "int", "Dict[str, str]" [call-overload]
231 method=interpolation, # type: ignore[call-overload]
232 )