Function push

dask/array/overlap.py:895–928

Dask-version of bottleneck.push .. note:: Requires bottleneck to be installed.

(array, n, axis)

def push(array, n, axis):
    """
    Dask-version of bottleneck.push

    .. note::

        Requires bottleneck to be installed.
    """
    import_optional_dependency("bottleneck", min_version="1.3.7")

    if n is not None and 0 < n < array.shape[axis] - 1:
        arr = broadcast_to(
            arange(
                array.shape[axis], chunks=array.chunks[axis], dtype=array.dtype
            ).reshape(
                tuple(size if i == axis else 1 for i, size in enumerate(array.shape))
            ),
            array.shape,
            array.chunks,
        )
        valid_arange = where(notnull(array), arr, np.nan)
        valid_limits = (arr - push(valid_arange, None, axis)) <= n
        # omit the forward fill that violate the limit
        return where(valid_limits, push(array, None, axis), np.nan)

    # The method parameter makes that the tests for python 3.7 fails.
    return cumreduction(
        func=_push,
        binop=_fill_with_last_one,
        ident=np.nan,
        x=array,
        axis=axis,
        dtype=array.dtype,
    )


def _fill_with_last_one(a, b):