| 1621 | |
| 1622 | |
| 1623 | def _custom_quantile( |
| 1624 | a, |
| 1625 | q, |
| 1626 | axis=None, |
| 1627 | method="linear", |
| 1628 | interpolation=None, |
| 1629 | keepdims=False, |
| 1630 | **kwargs, |
| 1631 | ): |
| 1632 | if ( |
| 1633 | not {method, interpolation}.issubset({"linear", None}) |
| 1634 | or len(axis) != 1 |
| 1635 | or axis[0] != len(a.shape) - 1 |
| 1636 | or len(a.shape) == 1 |
| 1637 | or a.shape[-1] > 1000 |
| 1638 | ): |
| 1639 | # bail to nanquantile. Assumptions are pretty strict for now but we |
| 1640 | # do cover the xarray.quantile case. |
| 1641 | return np.nanquantile( |
| 1642 | a, |
| 1643 | q, |
| 1644 | axis=axis, |
| 1645 | method=method, |
| 1646 | interpolation=interpolation, |
| 1647 | keepdims=keepdims, |
| 1648 | **kwargs, |
| 1649 | ) |
| 1650 | # nanquantile in NumPy is pretty slow if the quantile axis is slow because |
| 1651 | # each quantile has overhead. |
| 1652 | # This method works around this by calculating the quantile manually. |
| 1653 | # Steps: |
| 1654 | # 1. Sort the array along the quantile axis (this is the most expensive step |
| 1655 | # 2. Calculate which positions are the quantile positions |
| 1656 | # (respecting NaN values, so each quantile can have a different indexer) |
| 1657 | # 3. Get the neighboring values of the quantile positions |
| 1658 | # 4. Perform linear interpolation between the neighboring values |
| 1659 | # |
| 1660 | # The main advantage is that we get rid of the overhead, removing GIL blockage |
| 1661 | # and just generally making things faster. |
| 1662 | |
| 1663 | sorted_arr = np.sort(a, axis=-1) |
| 1664 | indexers = _span_indexers(a) |
| 1665 | nr_quantiles = len(indexers[0]) |
| 1666 | |
| 1667 | is_scalar = False |
| 1668 | if not isinstance(q, Iterable): |
| 1669 | is_scalar = True |
| 1670 | q = [q] |
| 1671 | |
| 1672 | quantiles = [] |
| 1673 | reshape_shapes = (1,) + tuple(sorted_arr.shape[:-1]) + ((1,) if keepdims else ()) |
| 1674 | for single_q in list(q): |
| 1675 | i = ( |
| 1676 | np.ones(nr_quantiles) * (a.shape[-1] - 1) |
| 1677 | - np.isnan(sorted_arr).sum(axis=-1).reshape(-1) |
| 1678 | ) * single_q |
| 1679 | lower_value, higher_value = np.floor(i).astype(int), np.ceil(i).astype(int) |
| 1680 | |