| 1649 | |
| 1650 | |
| 1651 | def _custom_quantile( |
| 1652 | a, |
| 1653 | q, |
| 1654 | axis=None, |
| 1655 | method="linear", |
| 1656 | keepdims=False, |
| 1657 | **kwargs, |
| 1658 | ): |
| 1659 | if ( |
| 1660 | method != "linear" |
| 1661 | or len(axis) != 1 |
| 1662 | or axis[0] != len(a.shape) - 1 |
| 1663 | or len(a.shape) == 1 |
| 1664 | or a.shape[-1] > 1000 |
| 1665 | ): |
| 1666 | # bail to nanquantile. Assumptions are pretty strict for now but we |
| 1667 | # do cover the xarray.quantile case. |
| 1668 | return np.nanquantile( |
| 1669 | a, |
| 1670 | q, |
| 1671 | axis=axis, |
| 1672 | method=method, |
| 1673 | keepdims=keepdims, |
| 1674 | **kwargs, |
| 1675 | ) |
| 1676 | # nanquantile in NumPy is pretty slow if the quantile axis is slow because |
| 1677 | # each quantile has overhead. |
| 1678 | # This method works around this by calculating the quantile manually. |
| 1679 | # Steps: |
| 1680 | # 1. Sort the array along the quantile axis (this is the most expensive step |
| 1681 | # 2. Calculate which positions are the quantile positions |
| 1682 | # (respecting NaN values, so each quantile can have a different indexer) |
| 1683 | # 3. Get the neighboring values of the quantile positions |
| 1684 | # 4. Perform linear interpolation between the neighboring values |
| 1685 | # |
| 1686 | # The main advantage is that we get rid of the overhead, removing GIL blockage |
| 1687 | # and just generally making things faster. |
| 1688 | |
| 1689 | sorted_arr = np.sort(a, axis=-1) |
| 1690 | indexers = _span_indexers(a) |
| 1691 | nr_quantiles = len(indexers[0]) |
| 1692 | |
| 1693 | is_scalar = False |
| 1694 | if not isinstance(q, Iterable): |
| 1695 | is_scalar = True |
| 1696 | q = [q] |
| 1697 | |
| 1698 | quantiles = [] |
| 1699 | reshape_shapes = (1,) + tuple(sorted_arr.shape[:-1]) + ((1,) if keepdims else ()) |
| 1700 | for single_q in list(q): |
| 1701 | i = ( |
| 1702 | np.ones(nr_quantiles) * (a.shape[-1] - 1) |
| 1703 | - np.isnan(sorted_arr).sum(axis=-1).reshape(-1) |
| 1704 | ) * single_q |
| 1705 | lower_value, higher_value = np.floor(i).astype(int), np.ceil(i).astype(int) |
| 1706 | |
| 1707 | # Get neighboring values |
| 1708 | lower = sorted_arr[tuple(indexers) + (tuple(lower_value),)] |