Slice x with one or more dask arrays of bools This is a helper function of `Array.__getitem__`. Parameters ---------- x: Array index: tuple with as many elements as x.ndim, among which there are one or more Array's with dtype=bool Returns ------- tuple o
(x, index)
| 1059 | |
| 1060 | |
| 1061 | def slice_with_bool_dask_array(x, index): |
| 1062 | """Slice x with one or more dask arrays of bools |
| 1063 | |
| 1064 | This is a helper function of `Array.__getitem__`. |
| 1065 | |
| 1066 | Parameters |
| 1067 | ---------- |
| 1068 | x: Array |
| 1069 | index: tuple with as many elements as x.ndim, among which there are |
| 1070 | one or more Array's with dtype=bool |
| 1071 | |
| 1072 | Returns |
| 1073 | ------- |
| 1074 | tuple of (sliced x, new index) |
| 1075 | |
| 1076 | where the new index is the same as the input, but with slice(None) |
| 1077 | replaced to the original slicer when a filter has been applied. |
| 1078 | |
| 1079 | Note: The sliced x will have nan chunks on the sliced axes. |
| 1080 | """ |
| 1081 | from dask.array.core import Array, blockwise, elemwise |
| 1082 | |
| 1083 | out_index = [ |
| 1084 | slice(None) if isinstance(ind, Array) and ind.dtype == bool else ind |
| 1085 | for ind in index |
| 1086 | ] |
| 1087 | |
| 1088 | if len(index) == 1 and index[0].ndim == x.ndim: |
| 1089 | if not np.isnan(x.shape).any() and not np.isnan(index[0].shape).any(): |
| 1090 | x = x.ravel() |
| 1091 | index = tuple(i.ravel() for i in index) |
| 1092 | elif x.ndim > 1: |
| 1093 | warnings.warn( |
| 1094 | "When slicing a Dask array of unknown chunks with a boolean mask " |
| 1095 | "Dask array, the output array may have a different ordering " |
| 1096 | "compared to the equivalent NumPy operation. This will raise an " |
| 1097 | "error in a future release of Dask.", |
| 1098 | stacklevel=3, |
| 1099 | ) |
| 1100 | y = elemwise(getitem, x, *index, dtype=x.dtype) |
| 1101 | name = "getitem-" + tokenize(x, index) |
| 1102 | dsk = {(name, i): k for i, k in enumerate(core.flatten(y.__dask_keys__()))} |
| 1103 | chunks = ((np.nan,) * y.npartitions,) |
| 1104 | graph = HighLevelGraph.from_collections(name, dsk, dependencies=[y]) |
| 1105 | return Array(graph, name, chunks, x.dtype), out_index |
| 1106 | |
| 1107 | if any( |
| 1108 | isinstance(ind, Array) and ind.dtype == bool and ind.ndim != 1 for ind in index |
| 1109 | ): |
| 1110 | raise NotImplementedError( |
| 1111 | "Slicing with dask.array of bools only permitted when " |
| 1112 | "the indexer has only one dimension or when " |
| 1113 | "it has the same dimension as the sliced " |
| 1114 | "array" |
| 1115 | ) |
| 1116 | indexes = [ |
| 1117 | ind if isinstance(ind, Array) and ind.dtype == bool else slice(None) |
| 1118 | for ind in index |
no test coverage detected
searching dependent graphs…