Function _custom_nanquantile

dask/array/reductions.py:1660–1742

(
    a,
    q,
    axis=None,
    method="linear",
    keepdims=False,
    weights=None,
    **kwargs,
)

def _custom_nanquantile(
    a,
    q,
    axis=None,
    method="linear",
    keepdims=False,
    weights=None,
    **kwargs,
):
    if (
        method != "linear"
        or len(axis) != 1
        or axis[0] != len(a.shape) - 1
        or len(a.shape) == 1
        or a.shape[-1] > 1000
        or q.ndim > 1
        or weights is not None
    ):
        # bail to nanquantile. Assumptions are pretty strict for now but we
        # do cover the xarray.quantile case.
        if weights is not None and weights.ndim != a.ndim:
            # np.nanquantile requires weights to have the same shape as a,
            # unlike np.quantile which broadcasts 1-D weights automatically.
            expand_at = [i for i in range(a.ndim) if i not in axis]
            weights = np.broadcast_to(np.expand_dims(weights, axis=expand_at), a.shape)
            # NumPy <2.0 doesn't support the weights parameter
            kwargs["weights"] = weights

        return np.nanquantile(
            a,
            q,
            axis=axis,
            method=method,
            keepdims=keepdims,
            **kwargs,
        )
    # nanquantile in NumPy is pretty slow if the quantile axis is slow because
    # each quantile has overhead.
    # This method works around this by calculating the quantile manually.
    # Steps:
    # 1. Sort the array along the quantile axis (this is the most expensive step
    # 2. Calculate which positions are the quantile positions
    #    (respecting NaN values, so each quantile can have a different indexer)
    # 3. Get the neighboring values of the quantile positions
    # 4. Perform linear interpolation between the neighboring values
    #
    # The main advantage is that we get rid of the overhead, removing GIL blockage
    # and just generally making things faster.

    sorted_arr = np.sort(a, axis=-1)
    indexers = _span_indexers(a)
    nr_quantiles = len(indexers[0])

    is_scalar = q.ndim == 0
    if is_scalar:
        q = q.reshape((1,))

    quantiles = []