Return the indices to access (n, n) arrays, given a masking function. Assume `mask_func` is a function that, for a square array a of size ``(n, n)`` with a possible offset argument `k`, when called as ``mask_func(a, k)`` returns a new array with zeros in certain locations (func
(n, mask_func, k=0)
| 811 | |
| 812 | @set_module('numpy') |
| 813 | def mask_indices(n, mask_func, k=0): |
| 814 | """ |
| 815 | Return the indices to access (n, n) arrays, given a masking function. |
| 816 | |
| 817 | Assume `mask_func` is a function that, for a square array a of size |
| 818 | ``(n, n)`` with a possible offset argument `k`, when called as |
| 819 | ``mask_func(a, k)`` returns a new array with zeros in certain locations |
| 820 | (functions like `triu` or `tril` do precisely this). Then this function |
| 821 | returns the indices where the non-zero values would be located. |
| 822 | |
| 823 | Parameters |
| 824 | ---------- |
| 825 | n : int |
| 826 | The returned indices will be valid to access arrays of shape (n, n). |
| 827 | mask_func : callable |
| 828 | A function whose call signature is similar to that of `triu`, `tril`. |
| 829 | That is, ``mask_func(x, k)`` returns a boolean array, shaped like `x`. |
| 830 | `k` is an optional argument to the function. |
| 831 | k : scalar |
| 832 | An optional argument which is passed through to `mask_func`. Functions |
| 833 | like `triu`, `tril` take a second argument that is interpreted as an |
| 834 | offset. |
| 835 | |
| 836 | Returns |
| 837 | ------- |
| 838 | indices : tuple of arrays. |
| 839 | The `n` arrays of indices corresponding to the locations where |
| 840 | ``mask_func(np.ones((n, n)), k)`` is True. |
| 841 | |
| 842 | See Also |
| 843 | -------- |
| 844 | triu, tril, triu_indices, tril_indices |
| 845 | |
| 846 | Notes |
| 847 | ----- |
| 848 | .. versionadded:: 1.4.0 |
| 849 | |
| 850 | Examples |
| 851 | -------- |
| 852 | These are the indices that would allow you to access the upper triangular |
| 853 | part of any 3x3 array: |
| 854 | |
| 855 | >>> iu = np.mask_indices(3, np.triu) |
| 856 | |
| 857 | For example, if `a` is a 3x3 array: |
| 858 | |
| 859 | >>> a = np.arange(9).reshape(3, 3) |
| 860 | >>> a |
| 861 | array([[0, 1, 2], |
| 862 | [3, 4, 5], |
| 863 | [6, 7, 8]]) |
| 864 | >>> a[iu] |
| 865 | array([0, 1, 2, 4, 5, 8]) |
| 866 | |
| 867 | An offset can be passed also to the masking function. This gets us the |
| 868 | indices starting on the first diagonal right of the main one: |
| 869 | |
| 870 | >>> iu1 = np.mask_indices(3, np.triu, 1) |