Dask-version of bottleneck.push .. note:: Requires bottleneck to be installed.
(array, n, axis)
| 897 | |
| 898 | |
| 899 | def push(array, n, axis): |
| 900 | """ |
| 901 | Dask-version of bottleneck.push |
| 902 | |
| 903 | .. note:: |
| 904 | |
| 905 | Requires bottleneck to be installed. |
| 906 | """ |
| 907 | import_optional_dependency("bottleneck", min_version="1.3.7") |
| 908 | |
| 909 | if n is not None and 0 < n < array.shape[axis] - 1: |
| 910 | arr = broadcast_to( |
| 911 | arange( |
| 912 | array.shape[axis], chunks=array.chunks[axis], dtype=array.dtype |
| 913 | ).reshape( |
| 914 | tuple(size if i == axis else 1 for i, size in enumerate(array.shape)) |
| 915 | ), |
| 916 | array.shape, |
| 917 | array.chunks, |
| 918 | ) |
| 919 | valid_arange = where(notnull(array), arr, np.nan) |
| 920 | valid_limits = (arr - push(valid_arange, None, axis)) <= n |
| 921 | # omit the forward fill that violate the limit |
| 922 | return where(valid_limits, push(array, None, axis), np.nan) |
| 923 | |
| 924 | # The method parameter makes that the tests for python 3.7 fails. |
| 925 | return cumreduction( |
| 926 | func=_push, |
| 927 | binop=_fill_with_last_one, |
| 928 | ident=np.nan, |
| 929 | x=array, |
| 930 | axis=axis, |
| 931 | dtype=array.dtype, |
| 932 | ) |
| 933 | |
| 934 | |
| 935 | def _fill_with_last_one(a, b): |