Return a 2-D Array with ones on the diagonal and zeros elsewhere. Parameters ---------- N : int Number of rows in the output. chunks : int, str How to chunk the array. Must be one of the following forms: - A blocksize like 1000. - A size in by
(N, chunks="auto", M=None, k=0, dtype=float)
| 561 | |
| 562 | |
| 563 | def eye(N, chunks="auto", M=None, k=0, dtype=float): |
| 564 | """ |
| 565 | Return a 2-D Array with ones on the diagonal and zeros elsewhere. |
| 566 | |
| 567 | Parameters |
| 568 | ---------- |
| 569 | N : int |
| 570 | Number of rows in the output. |
| 571 | chunks : int, str |
| 572 | How to chunk the array. Must be one of the following forms: |
| 573 | |
| 574 | - A blocksize like 1000. |
| 575 | - A size in bytes, like "100 MiB" which will choose a uniform |
| 576 | block-like shape |
| 577 | - The word "auto" which acts like the above, but uses a configuration |
| 578 | value ``array.chunk-size`` for the chunk size |
| 579 | M : int, optional |
| 580 | Number of columns in the output. If None, defaults to `N`. |
| 581 | k : int, optional |
| 582 | Index of the diagonal: 0 (the default) refers to the main diagonal, |
| 583 | a positive value refers to an upper diagonal, and a negative value |
| 584 | to a lower diagonal. |
| 585 | dtype : data-type, optional |
| 586 | Data-type of the returned array. |
| 587 | |
| 588 | Returns |
| 589 | ------- |
| 590 | I : Array of shape (N,M) |
| 591 | An array where all elements are equal to zero, except for the `k`-th |
| 592 | diagonal, whose values are equal to one. |
| 593 | """ |
| 594 | if M is None: |
| 595 | M = N |
| 596 | if dtype is None: |
| 597 | dtype = float |
| 598 | |
| 599 | if not isinstance(chunks, (int, str)): |
| 600 | raise ValueError("chunks must be an int or string") |
| 601 | |
| 602 | vchunks, hchunks = normalize_chunks(chunks, shape=(N, M), dtype=dtype) |
| 603 | chunks = vchunks[0] |
| 604 | |
| 605 | token = tokenize(N, chunks, M, k, dtype) |
| 606 | name_eye = f"eye-{token}" |
| 607 | |
| 608 | dsk = {} |
| 609 | for i, vchunk in enumerate(vchunks): |
| 610 | for j, hchunk in enumerate(hchunks): |
| 611 | key = (name_eye, i, j) |
| 612 | if (j - i - 1) * chunks <= k <= (j - i + 1) * chunks: |
| 613 | t = Task( |
| 614 | key, |
| 615 | np.eye, |
| 616 | vchunk, |
| 617 | hchunk, |
| 618 | k - (j - i) * chunks, |
| 619 | dtype, |
| 620 | ) |
no test coverage detected