Implements NumPy's ``indices`` for Dask Arrays. Generates a grid of indices covering the dimensions provided. The final array has the shape ``(len(dimensions), *dimensions)``. The chunks are used to specify the chunking for axis 1 up to ``len(dimensions)``. The 0th axis always
(dimensions, dtype=int, chunks="auto")
| 507 | |
| 508 | |
| 509 | def indices(dimensions, dtype=int, chunks="auto"): |
| 510 | """ |
| 511 | Implements NumPy's ``indices`` for Dask Arrays. |
| 512 | |
| 513 | Generates a grid of indices covering the dimensions provided. |
| 514 | |
| 515 | The final array has the shape ``(len(dimensions), *dimensions)``. The |
| 516 | chunks are used to specify the chunking for axis 1 up to |
| 517 | ``len(dimensions)``. The 0th axis always has chunks of length 1. |
| 518 | |
| 519 | Parameters |
| 520 | ---------- |
| 521 | dimensions : sequence of ints |
| 522 | The shape of the index grid. |
| 523 | dtype : dtype, optional |
| 524 | Type to use for the array. Default is ``int``. |
| 525 | chunks : sequence of ints, str |
| 526 | The size of each block. Must be one of the following forms: |
| 527 | |
| 528 | - A blocksize like (500, 1000) |
| 529 | - A size in bytes, like "100 MiB" which will choose a uniform |
| 530 | block-like shape |
| 531 | - The word "auto" which acts like the above, but uses a configuration |
| 532 | value ``array.chunk-size`` for the chunk size |
| 533 | |
| 534 | Note that the last block will have fewer samples if ``len(array) % chunks != 0``. |
| 535 | |
| 536 | Returns |
| 537 | ------- |
| 538 | grid : dask array |
| 539 | """ |
| 540 | dimensions = tuple(dimensions) |
| 541 | dtype = np.dtype(dtype) |
| 542 | chunks = normalize_chunks(chunks, shape=dimensions, dtype=dtype) |
| 543 | |
| 544 | if len(dimensions) != len(chunks): |
| 545 | raise ValueError("Need same number of chunks as dimensions.") |
| 546 | |
| 547 | xi = [] |
| 548 | for i in range(len(dimensions)): |
| 549 | xi.append(arange(dimensions[i], dtype=dtype, chunks=(chunks[i],))) |
| 550 | |
| 551 | grid = [] |
| 552 | if all(dimensions): |
| 553 | grid = meshgrid(*xi, indexing="ij") |
| 554 | |
| 555 | if grid: |
| 556 | grid = stack(grid) |
| 557 | else: |
| 558 | grid = empty((len(dimensions),) + dimensions, dtype=dtype, chunks=(1,) + chunks) |
| 559 | |
| 560 | return grid |
| 561 | |
| 562 | |
| 563 | def eye(N, chunks="auto", M=None, k=0, dtype=float): |