Normalize slicing indexes 1. Replaces ellipses with many full slices 2. Adds full slices to end of index 3. Checks bounding conditions 4. Replace multidimensional numpy arrays with dask arrays 5. Replaces numpy arrays with lists 6. Posify's integers and lists 7. N
(idx, shape)
| 80 | |
| 81 | |
| 82 | def normalize_index(idx, shape): |
| 83 | """Normalize slicing indexes |
| 84 | |
| 85 | 1. Replaces ellipses with many full slices |
| 86 | 2. Adds full slices to end of index |
| 87 | 3. Checks bounding conditions |
| 88 | 4. Replace multidimensional numpy arrays with dask arrays |
| 89 | 5. Replaces numpy arrays with lists |
| 90 | 6. Posify's integers and lists |
| 91 | 7. Normalizes slices to canonical form |
| 92 | |
| 93 | Examples |
| 94 | -------- |
| 95 | >>> normalize_index(1, (10,)) |
| 96 | (1,) |
| 97 | >>> normalize_index(-1, (10,)) |
| 98 | (9,) |
| 99 | >>> normalize_index([-1], (10,)) |
| 100 | (array([9]),) |
| 101 | >>> normalize_index(slice(-3, 10, 1), (10,)) |
| 102 | (slice(7, None, None),) |
| 103 | >>> normalize_index((Ellipsis, None), (10,)) |
| 104 | (slice(None, None, None), None) |
| 105 | >>> normalize_index(np.array([[True, False], [False, True], [True, True]]), (3, 2)) |
| 106 | (dask.array<array, shape=(3, 2), dtype=bool, chunksize=(3, 2), chunktype=numpy.ndarray>,) |
| 107 | """ |
| 108 | from dask.array._array_expr._collection import Array, from_array |
| 109 | |
| 110 | if not isinstance(idx, tuple): |
| 111 | idx = (idx,) |
| 112 | |
| 113 | # if a > 1D numpy.array is provided, cast it to a dask array |
| 114 | if len(idx) > 0 and len(shape) > 1: |
| 115 | i = idx[0] |
| 116 | if is_arraylike(i) and not isinstance(i, Array) and i.shape == shape: |
| 117 | idx = (from_array(i), *idx[1:]) |
| 118 | |
| 119 | idx = replace_ellipsis(len(shape), idx) |
| 120 | n_sliced_dims = 0 |
| 121 | for i in idx: |
| 122 | if hasattr(i, "ndim") and i.ndim >= 1: |
| 123 | n_sliced_dims += i.ndim |
| 124 | elif i is None: |
| 125 | continue |
| 126 | else: |
| 127 | n_sliced_dims += 1 |
| 128 | |
| 129 | idx = idx + (slice(None),) * (len(shape) - n_sliced_dims) |
| 130 | if len([i for i in idx if i is not None]) > len(shape): |
| 131 | raise IndexError("Too many indices for array") |
| 132 | |
| 133 | none_shape = [] |
| 134 | i = 0 |
| 135 | for ind in idx: |
| 136 | if ind is not None: |
| 137 | none_shape.append(shape[i]) |
| 138 | i += 1 |
| 139 | else: |
searching dependent graphs…