Point wise indexing with only NumPy Arrays.
(x, dict_indexes)
| 5649 | |
| 5650 | |
| 5651 | def _vindex_array(x, dict_indexes): |
| 5652 | """Point wise indexing with only NumPy Arrays.""" |
| 5653 | |
| 5654 | token = tokenize(x, dict_indexes) |
| 5655 | try: |
| 5656 | broadcast_shape = np.broadcast_shapes( |
| 5657 | *(arr.shape for arr in dict_indexes.values()) |
| 5658 | ) |
| 5659 | except ValueError as e: |
| 5660 | # note: error message exactly matches numpy |
| 5661 | shapes_str = " ".join(str(a.shape) for a in dict_indexes.values()) |
| 5662 | raise IndexError( |
| 5663 | "shape mismatch: indexing arrays could not be " |
| 5664 | "broadcast together with shapes " + shapes_str |
| 5665 | ) from e |
| 5666 | npoints = math.prod(broadcast_shape) |
| 5667 | axes = [i for i in range(x.ndim) if i in dict_indexes] |
| 5668 | |
| 5669 | def _subset_to_indexed_axes(iterable): |
| 5670 | for i, elem in enumerate(iterable): |
| 5671 | if i in axes: |
| 5672 | yield elem |
| 5673 | |
| 5674 | bounds2 = tuple( |
| 5675 | np.array(cached_cumsum(c, initial_zero=True)) |
| 5676 | for c in _subset_to_indexed_axes(x.chunks) |
| 5677 | ) |
| 5678 | axis = _get_axis(tuple(i if i in axes else None for i in range(x.ndim))) |
| 5679 | out_name = "vindex-merge-" + token |
| 5680 | |
| 5681 | # Now compute indices of each output element within each input block |
| 5682 | # The index is relative to the block, not the array. |
| 5683 | block_idxs = tuple( |
| 5684 | np.searchsorted(b, ind, side="right") - 1 |
| 5685 | for b, ind in zip(bounds2, dict_indexes.values()) |
| 5686 | ) |
| 5687 | starts = (b[i] for i, b in zip(block_idxs, bounds2)) |
| 5688 | inblock_idxs = [] |
| 5689 | for idx, start in zip(dict_indexes.values(), starts): |
| 5690 | a = idx - start |
| 5691 | if len(a) > 0: |
| 5692 | dtype = np.min_scalar_type(np.max(a, axis=None)) |
| 5693 | inblock_idxs.append(a.astype(dtype, copy=False)) |
| 5694 | else: |
| 5695 | inblock_idxs.append(a) |
| 5696 | |
| 5697 | inblock_idxs = np.broadcast_arrays(*inblock_idxs) |
| 5698 | |
| 5699 | chunks = [c for i, c in enumerate(x.chunks) if i not in axes] |
| 5700 | # determine number of points in one single output block. |
| 5701 | # Use the input chunk size to determine this. |
| 5702 | max_chunk_point_dimensions = reduce( |
| 5703 | mul, map(cached_max, _subset_to_indexed_axes(x.chunks)) |
| 5704 | ) |
| 5705 | |
| 5706 | n_chunks, remainder = divmod(npoints, max_chunk_point_dimensions) |
| 5707 | chunks.insert( |
| 5708 | 0, |
no test coverage detected