1"""
2Quantilization functions and related stuff
3"""
4from __future__ import annotations
5
6from typing import (
7 TYPE_CHECKING,
8 Any,
9 Callable,
10 Literal,
11)
12
13import numpy as np
14
15from pandas._libs import (
16 Timedelta,
17 Timestamp,
18 lib,
19)
20
21from pandas.core.dtypes.common import (
22 ensure_platform_int,
23 is_bool_dtype,
24 is_integer,
25 is_list_like,
26 is_numeric_dtype,
27 is_scalar,
28)
29from pandas.core.dtypes.dtypes import (
30 CategoricalDtype,
31 DatetimeTZDtype,
32 ExtensionDtype,
33)
34from pandas.core.dtypes.generic import ABCSeries
35from pandas.core.dtypes.missing import isna
36
37from pandas import (
38 Categorical,
39 Index,
40 IntervalIndex,
41)
42import pandas.core.algorithms as algos
43from pandas.core.arrays.datetimelike import dtype_to_unit
44
45if TYPE_CHECKING:
46 from pandas._typing import (
47 DtypeObj,
48 IntervalLeftRight,
49 )
50
51
52def cut(
53 x,
54 bins,
55 right: bool = True,
56 labels=None,
57 retbins: bool = False,
58 precision: int = 3,
59 include_lowest: bool = False,
60 duplicates: str = "raise",
61 ordered: bool = True,
62):
63 """
64 Bin values into discrete intervals.
65
66 Use `cut` when you need to segment and sort data values into bins. This
67 function is also useful for going from a continuous variable to a
68 categorical variable. For example, `cut` could convert ages to groups of
69 age ranges. Supports binning into an equal number of bins, or a
70 pre-specified array of bins.
71
72 Parameters
73 ----------
74 x : array-like
75 The input array to be binned. Must be 1-dimensional.
76 bins : int, sequence of scalars, or IntervalIndex
77 The criteria to bin by.
78
79 * int : Defines the number of equal-width bins in the range of `x`. The
80 range of `x` is extended by .1% on each side to include the minimum
81 and maximum values of `x`.
82 * sequence of scalars : Defines the bin edges allowing for non-uniform
83 width. No extension of the range of `x` is done.
84 * IntervalIndex : Defines the exact bins to be used. Note that
85 IntervalIndex for `bins` must be non-overlapping.
86
87 right : bool, default True
88 Indicates whether `bins` includes the rightmost edge or not. If
89 ``right == True`` (the default), then the `bins` ``[1, 2, 3, 4]``
90 indicate (1,2], (2,3], (3,4]. This argument is ignored when
91 `bins` is an IntervalIndex.
92 labels : array or False, default None
93 Specifies the labels for the returned bins. Must be the same length as
94 the resulting bins. If False, returns only integer indicators of the
95 bins. This affects the type of the output container (see below).
96 This argument is ignored when `bins` is an IntervalIndex. If True,
97 raises an error. When `ordered=False`, labels must be provided.
98 retbins : bool, default False
99 Whether to return the bins or not. Useful when bins is provided
100 as a scalar.
101 precision : int, default 3
102 The precision at which to store and display the bins labels.
103 include_lowest : bool, default False
104 Whether the first interval should be left-inclusive or not.
105 duplicates : {default 'raise', 'drop'}, optional
106 If bin edges are not unique, raise ValueError or drop non-uniques.
107 ordered : bool, default True
108 Whether the labels are ordered or not. Applies to returned types
109 Categorical and Series (with Categorical dtype). If True,
110 the resulting categorical will be ordered. If False, the resulting
111 categorical will be unordered (labels must be provided).
112
113 Returns
114 -------
115 out : Categorical, Series, or ndarray
116 An array-like object representing the respective bin for each value
117 of `x`. The type depends on the value of `labels`.
118
119 * None (default) : returns a Series for Series `x` or a
120 Categorical for all other inputs. The values stored within
121 are Interval dtype.
122
123 * sequence of scalars : returns a Series for Series `x` or a
124 Categorical for all other inputs. The values stored within
125 are whatever the type in the sequence is.
126
127 * False : returns an ndarray of integers.
128
129 bins : numpy.ndarray or IntervalIndex.
130 The computed or specified bins. Only returned when `retbins=True`.
131 For scalar or sequence `bins`, this is an ndarray with the computed
132 bins. If set `duplicates=drop`, `bins` will drop non-unique bin. For
133 an IntervalIndex `bins`, this is equal to `bins`.
134
135 See Also
136 --------
137 qcut : Discretize variable into equal-sized buckets based on rank
138 or based on sample quantiles.
139 Categorical : Array type for storing data that come from a
140 fixed set of values.
141 Series : One-dimensional array with axis labels (including time series).
142 IntervalIndex : Immutable Index implementing an ordered, sliceable set.
143
144 Notes
145 -----
146 Any NA values will be NA in the result. Out of bounds values will be NA in
147 the resulting Series or Categorical object.
148
149 Reference :ref:`the user guide <reshaping.tile.cut>` for more examples.
150
151 Examples
152 --------
153 Discretize into three equal-sized bins.
154
155 >>> pd.cut(np.array([1, 7, 5, 4, 6, 3]), 3)
156 ... # doctest: +ELLIPSIS
157 [(0.994, 3.0], (5.0, 7.0], (3.0, 5.0], (3.0, 5.0], (5.0, 7.0], ...
158 Categories (3, interval[float64, right]): [(0.994, 3.0] < (3.0, 5.0] ...
159
160 >>> pd.cut(np.array([1, 7, 5, 4, 6, 3]), 3, retbins=True)
161 ... # doctest: +ELLIPSIS
162 ([(0.994, 3.0], (5.0, 7.0], (3.0, 5.0], (3.0, 5.0], (5.0, 7.0], ...
163 Categories (3, interval[float64, right]): [(0.994, 3.0] < (3.0, 5.0] ...
164 array([0.994, 3. , 5. , 7. ]))
165
166 Discovers the same bins, but assign them specific labels. Notice that
167 the returned Categorical's categories are `labels` and is ordered.
168
169 >>> pd.cut(np.array([1, 7, 5, 4, 6, 3]),
170 ... 3, labels=["bad", "medium", "good"])
171 ['bad', 'good', 'medium', 'medium', 'good', 'bad']
172 Categories (3, object): ['bad' < 'medium' < 'good']
173
174 ``ordered=False`` will result in unordered categories when labels are passed.
175 This parameter can be used to allow non-unique labels:
176
177 >>> pd.cut(np.array([1, 7, 5, 4, 6, 3]), 3,
178 ... labels=["B", "A", "B"], ordered=False)
179 ['B', 'B', 'A', 'A', 'B', 'B']
180 Categories (2, object): ['A', 'B']
181
182 ``labels=False`` implies you just want the bins back.
183
184 >>> pd.cut([0, 1, 1, 2], bins=4, labels=False)
185 array([0, 1, 1, 3])
186
187 Passing a Series as an input returns a Series with categorical dtype:
188
189 >>> s = pd.Series(np.array([2, 4, 6, 8, 10]),
190 ... index=['a', 'b', 'c', 'd', 'e'])
191 >>> pd.cut(s, 3)
192 ... # doctest: +ELLIPSIS
193 a (1.992, 4.667]
194 b (1.992, 4.667]
195 c (4.667, 7.333]
196 d (7.333, 10.0]
197 e (7.333, 10.0]
198 dtype: category
199 Categories (3, interval[float64, right]): [(1.992, 4.667] < (4.667, ...
200
201 Passing a Series as an input returns a Series with mapping value.
202 It is used to map numerically to intervals based on bins.
203
204 >>> s = pd.Series(np.array([2, 4, 6, 8, 10]),
205 ... index=['a', 'b', 'c', 'd', 'e'])
206 >>> pd.cut(s, [0, 2, 4, 6, 8, 10], labels=False, retbins=True, right=False)
207 ... # doctest: +ELLIPSIS
208 (a 1.0
209 b 2.0
210 c 3.0
211 d 4.0
212 e NaN
213 dtype: float64,
214 array([ 0, 2, 4, 6, 8, 10]))
215
216 Use `drop` optional when bins is not unique
217
218 >>> pd.cut(s, [0, 2, 4, 6, 10, 10], labels=False, retbins=True,
219 ... right=False, duplicates='drop')
220 ... # doctest: +ELLIPSIS
221 (a 1.0
222 b 2.0
223 c 3.0
224 d 3.0
225 e NaN
226 dtype: float64,
227 array([ 0, 2, 4, 6, 10]))
228
229 Passing an IntervalIndex for `bins` results in those categories exactly.
230 Notice that values not covered by the IntervalIndex are set to NaN. 0
231 is to the left of the first bin (which is closed on the right), and 1.5
232 falls between two bins.
233
234 >>> bins = pd.IntervalIndex.from_tuples([(0, 1), (2, 3), (4, 5)])
235 >>> pd.cut([0, 0.5, 1.5, 2.5, 4.5], bins)
236 [NaN, (0.0, 1.0], NaN, (2.0, 3.0], (4.0, 5.0]]
237 Categories (3, interval[int64, right]): [(0, 1] < (2, 3] < (4, 5]]
238 """
239 # NOTE: this binning code is changed a bit from histogram for var(x) == 0
240
241 original = x
242 x_idx = _preprocess_for_cut(x)
243 x_idx, _ = _coerce_to_type(x_idx)
244
245 if not np.iterable(bins):
246 bins = _nbins_to_bins(x_idx, bins, right)
247
248 elif isinstance(bins, IntervalIndex):
249 if bins.is_overlapping:
250 raise ValueError("Overlapping IntervalIndex is not accepted.")
251
252 else:
253 bins = Index(bins)
254 if not bins.is_monotonic_increasing:
255 raise ValueError("bins must increase monotonically.")
256
257 fac, bins = _bins_to_cuts(
258 x_idx,
259 bins,
260 right=right,
261 labels=labels,
262 precision=precision,
263 include_lowest=include_lowest,
264 duplicates=duplicates,
265 ordered=ordered,
266 )
267
268 return _postprocess_for_cut(fac, bins, retbins, original)
269
270
271def qcut(
272 x,
273 q,
274 labels=None,
275 retbins: bool = False,
276 precision: int = 3,
277 duplicates: str = "raise",
278):
279 """
280 Quantile-based discretization function.
281
282 Discretize variable into equal-sized buckets based on rank or based
283 on sample quantiles. For example 1000 values for 10 quantiles would
284 produce a Categorical object indicating quantile membership for each data point.
285
286 Parameters
287 ----------
288 x : 1d ndarray or Series
289 q : int or list-like of float
290 Number of quantiles. 10 for deciles, 4 for quartiles, etc. Alternately
291 array of quantiles, e.g. [0, .25, .5, .75, 1.] for quartiles.
292 labels : array or False, default None
293 Used as labels for the resulting bins. Must be of the same length as
294 the resulting bins. If False, return only integer indicators of the
295 bins. If True, raises an error.
296 retbins : bool, optional
297 Whether to return the (bins, labels) or not. Can be useful if bins
298 is given as a scalar.
299 precision : int, optional
300 The precision at which to store and display the bins labels.
301 duplicates : {default 'raise', 'drop'}, optional
302 If bin edges are not unique, raise ValueError or drop non-uniques.
303
304 Returns
305 -------
306 out : Categorical or Series or array of integers if labels is False
307 The return type (Categorical or Series) depends on the input: a Series
308 of type category if input is a Series else Categorical. Bins are
309 represented as categories when categorical data is returned.
310 bins : ndarray of floats
311 Returned only if `retbins` is True.
312
313 Notes
314 -----
315 Out of bounds values will be NA in the resulting Categorical object
316
317 Examples
318 --------
319 >>> pd.qcut(range(5), 4)
320 ... # doctest: +ELLIPSIS
321 [(-0.001, 1.0], (-0.001, 1.0], (1.0, 2.0], (2.0, 3.0], (3.0, 4.0]]
322 Categories (4, interval[float64, right]): [(-0.001, 1.0] < (1.0, 2.0] ...
323
324 >>> pd.qcut(range(5), 3, labels=["good", "medium", "bad"])
325 ... # doctest: +SKIP
326 [good, good, medium, bad, bad]
327 Categories (3, object): [good < medium < bad]
328
329 >>> pd.qcut(range(5), 4, labels=False)
330 array([0, 0, 1, 2, 3])
331 """
332 original = x
333 x_idx = _preprocess_for_cut(x)
334 x_idx, _ = _coerce_to_type(x_idx)
335
336 quantiles = np.linspace(0, 1, q + 1) if is_integer(q) else q
337
338 bins = x_idx.to_series().dropna().quantile(quantiles)
339
340 fac, bins = _bins_to_cuts(
341 x_idx,
342 Index(bins),
343 labels=labels,
344 precision=precision,
345 include_lowest=True,
346 duplicates=duplicates,
347 )
348
349 return _postprocess_for_cut(fac, bins, retbins, original)
350
351
352def _nbins_to_bins(x_idx: Index, nbins: int, right: bool) -> Index:
353 """
354 If a user passed an integer N for bins, convert this to a sequence of N
355 equal(ish)-sized bins.
356 """
357 if is_scalar(nbins) and nbins < 1:
358 raise ValueError("`bins` should be a positive integer.")
359
360 if x_idx.size == 0:
361 raise ValueError("Cannot cut empty array")
362
363 rng = (x_idx.min(), x_idx.max())
364 mn, mx = rng
365
366 if is_numeric_dtype(x_idx.dtype) and (np.isinf(mn) or np.isinf(mx)):
367 # GH#24314
368 raise ValueError(
369 "cannot specify integer `bins` when input data contains infinity"
370 )
371
372 if mn == mx: # adjust end points before binning
373 if _is_dt_or_td(x_idx.dtype):
374 # using seconds=1 is pretty arbitrary here
375 # error: Argument 1 to "dtype_to_unit" has incompatible type
376 # "dtype[Any] | ExtensionDtype"; expected "DatetimeTZDtype | dtype[Any]"
377 unit = dtype_to_unit(x_idx.dtype) # type: ignore[arg-type]
378 td = Timedelta(seconds=1).as_unit(unit)
379 # Use DatetimeArray/TimedeltaArray method instead of linspace
380 # error: Item "ExtensionArray" of "ExtensionArray | ndarray[Any, Any]"
381 # has no attribute "_generate_range"
382 bins = x_idx._values._generate_range( # type: ignore[union-attr]
383 start=mn - td, end=mx + td, periods=nbins + 1, freq=None, unit=unit
384 )
385 else:
386 mn -= 0.001 * abs(mn) if mn != 0 else 0.001
387 mx += 0.001 * abs(mx) if mx != 0 else 0.001
388
389 bins = np.linspace(mn, mx, nbins + 1, endpoint=True)
390 else: # adjust end points after binning
391 if _is_dt_or_td(x_idx.dtype):
392 # Use DatetimeArray/TimedeltaArray method instead of linspace
393
394 # error: Argument 1 to "dtype_to_unit" has incompatible type
395 # "dtype[Any] | ExtensionDtype"; expected "DatetimeTZDtype | dtype[Any]"
396 unit = dtype_to_unit(x_idx.dtype) # type: ignore[arg-type]
397 # error: Item "ExtensionArray" of "ExtensionArray | ndarray[Any, Any]"
398 # has no attribute "_generate_range"
399 bins = x_idx._values._generate_range( # type: ignore[union-attr]
400 start=mn, end=mx, periods=nbins + 1, freq=None, unit=unit
401 )
402 else:
403 bins = np.linspace(mn, mx, nbins + 1, endpoint=True)
404 adj = (mx - mn) * 0.001 # 0.1% of the range
405 if right:
406 bins[0] -= adj
407 else:
408 bins[-1] += adj
409
410 return Index(bins)
411
412
413def _bins_to_cuts(
414 x_idx: Index,
415 bins: Index,
416 right: bool = True,
417 labels=None,
418 precision: int = 3,
419 include_lowest: bool = False,
420 duplicates: str = "raise",
421 ordered: bool = True,
422):
423 if not ordered and labels is None:
424 raise ValueError("'labels' must be provided if 'ordered = False'")
425
426 if duplicates not in ["raise", "drop"]:
427 raise ValueError(
428 "invalid value for 'duplicates' parameter, valid options are: raise, drop"
429 )
430
431 result: Categorical | np.ndarray
432
433 if isinstance(bins, IntervalIndex):
434 # we have a fast-path here
435 ids = bins.get_indexer(x_idx)
436 cat_dtype = CategoricalDtype(bins, ordered=True)
437 result = Categorical.from_codes(ids, dtype=cat_dtype, validate=False)
438 return result, bins
439
440 unique_bins = algos.unique(bins)
441 if len(unique_bins) < len(bins) and len(bins) != 2:
442 if duplicates == "raise":
443 raise ValueError(
444 f"Bin edges must be unique: {repr(bins)}.\n"
445 f"You can drop duplicate edges by setting the 'duplicates' kwarg"
446 )
447 bins = unique_bins
448
449 side: Literal["left", "right"] = "left" if right else "right"
450
451 try:
452 ids = bins.searchsorted(x_idx, side=side)
453 except TypeError as err:
454 # e.g. test_datetime_nan_error if bins are DatetimeArray and x_idx
455 # is integers
456 if x_idx.dtype.kind == "m":
457 raise ValueError("bins must be of timedelta64 dtype") from err
458 elif x_idx.dtype.kind == bins.dtype.kind == "M":
459 raise ValueError(
460 "Cannot use timezone-naive bins with timezone-aware values, "
461 "or vice-versa"
462 ) from err
463 elif x_idx.dtype.kind == "M":
464 raise ValueError("bins must be of datetime64 dtype") from err
465 else:
466 raise
467 ids = ensure_platform_int(ids)
468
469 if include_lowest:
470 ids[x_idx == bins[0]] = 1
471
472 na_mask = isna(x_idx) | (ids == len(bins)) | (ids == 0)
473 has_nas = na_mask.any()
474
475 if labels is not False:
476 if not (labels is None or is_list_like(labels)):
477 raise ValueError(
478 "Bin labels must either be False, None or passed in as a "
479 "list-like argument"
480 )
481
482 if labels is None:
483 labels = _format_labels(
484 bins, precision, right=right, include_lowest=include_lowest
485 )
486 elif ordered and len(set(labels)) != len(labels):
487 raise ValueError(
488 "labels must be unique if ordered=True; pass ordered=False "
489 "for duplicate labels"
490 )
491 else:
492 if len(labels) != len(bins) - 1:
493 raise ValueError(
494 "Bin labels must be one fewer than the number of bin edges"
495 )
496
497 if not isinstance(getattr(labels, "dtype", None), CategoricalDtype):
498 labels = Categorical(
499 labels,
500 categories=labels if len(set(labels)) == len(labels) else None,
501 ordered=ordered,
502 )
503 # TODO: handle mismatch between categorical label order and pandas.cut order.
504 np.putmask(ids, na_mask, 0)
505 result = algos.take_nd(labels, ids - 1)
506
507 else:
508 result = ids - 1
509 if has_nas:
510 result = result.astype(np.float64)
511 np.putmask(result, na_mask, np.nan)
512
513 return result, bins
514
515
516def _coerce_to_type(x: Index) -> tuple[Index, DtypeObj | None]:
517 """
518 if the passed data is of datetime/timedelta, bool or nullable int type,
519 this method converts it to numeric so that cut or qcut method can
520 handle it
521 """
522 dtype: DtypeObj | None = None
523
524 if _is_dt_or_td(x.dtype):
525 dtype = x.dtype
526 elif is_bool_dtype(x.dtype):
527 # GH 20303
528 x = x.astype(np.int64)
529 # To support cut and qcut for IntegerArray we convert to float dtype.
530 # Will properly support in the future.
531 # https://github.com/pandas-dev/pandas/pull/31290
532 # https://github.com/pandas-dev/pandas/issues/31389
533 elif isinstance(x.dtype, ExtensionDtype) and is_numeric_dtype(x.dtype):
534 x_arr = x.to_numpy(dtype=np.float64, na_value=np.nan)
535 x = Index(x_arr)
536
537 return Index(x), dtype
538
539
540def _is_dt_or_td(dtype: DtypeObj) -> bool:
541 # Note: the dtype here comes from an Index.dtype, so we know that that any
542 # dt64/td64 dtype is of a supported unit.
543 return isinstance(dtype, DatetimeTZDtype) or lib.is_np_dtype(dtype, "mM")
544
545
546def _format_labels(
547 bins: Index,
548 precision: int,
549 right: bool = True,
550 include_lowest: bool = False,
551):
552 """based on the dtype, return our labels"""
553 closed: IntervalLeftRight = "right" if right else "left"
554
555 formatter: Callable[[Any], Timestamp] | Callable[[Any], Timedelta]
556
557 if _is_dt_or_td(bins.dtype):
558 # error: Argument 1 to "dtype_to_unit" has incompatible type
559 # "dtype[Any] | ExtensionDtype"; expected "DatetimeTZDtype | dtype[Any]"
560 unit = dtype_to_unit(bins.dtype) # type: ignore[arg-type]
561 formatter = lambda x: x
562 adjust = lambda x: x - Timedelta(1, unit=unit).as_unit(unit)
563 else:
564 precision = _infer_precision(precision, bins)
565 formatter = lambda x: _round_frac(x, precision)
566 adjust = lambda x: x - 10 ** (-precision)
567
568 breaks = [formatter(b) for b in bins]
569 if right and include_lowest:
570 # adjust lhs of first interval by precision to account for being right closed
571 breaks[0] = adjust(breaks[0])
572
573 if _is_dt_or_td(bins.dtype):
574 # error: "Index" has no attribute "as_unit"
575 breaks = type(bins)(breaks).as_unit(unit) # type: ignore[attr-defined]
576
577 return IntervalIndex.from_breaks(breaks, closed=closed)
578
579
580def _preprocess_for_cut(x) -> Index:
581 """
582 handles preprocessing for cut where we convert passed
583 input to array, strip the index information and store it
584 separately
585 """
586 # Check that the passed array is a Pandas or Numpy object
587 # We don't want to strip away a Pandas data-type here (e.g. datetimetz)
588 ndim = getattr(x, "ndim", None)
589 if ndim is None:
590 x = np.asarray(x)
591 if x.ndim != 1:
592 raise ValueError("Input array must be 1 dimensional")
593
594 return Index(x)
595
596
597def _postprocess_for_cut(fac, bins, retbins: bool, original):
598 """
599 handles post processing for the cut method where
600 we combine the index information if the originally passed
601 datatype was a series
602 """
603 if isinstance(original, ABCSeries):
604 fac = original._constructor(fac, index=original.index, name=original.name)
605
606 if not retbins:
607 return fac
608
609 if isinstance(bins, Index) and is_numeric_dtype(bins.dtype):
610 bins = bins._values
611
612 return fac, bins
613
614
615def _round_frac(x, precision: int):
616 """
617 Round the fractional part of the given number
618 """
619 if not np.isfinite(x) or x == 0:
620 return x
621 else:
622 frac, whole = np.modf(x)
623 if whole == 0:
624 digits = -int(np.floor(np.log10(abs(frac)))) - 1 + precision
625 else:
626 digits = precision
627 return np.around(x, digits)
628
629
630def _infer_precision(base_precision: int, bins: Index) -> int:
631 """
632 Infer an appropriate precision for _round_frac
633 """
634 for precision in range(base_precision, 20):
635 levels = np.asarray([_round_frac(b, precision) for b in bins])
636 if algos.unique(levels).size == bins.size:
637 return precision
638 return base_precision # default