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)
| 864 | |
| 865 | @set_module('numpy') |
| 866 | def mask_indices(n, mask_func, k=0): |
| 867 | """ |
| 868 | Return the indices to access (n, n) arrays, given a masking function. |
| 869 | |
| 870 | Assume `mask_func` is a function that, for a square array a of size |
| 871 | ``(n, n)`` with a possible offset argument `k`, when called as |
| 872 | ``mask_func(a, k)`` returns a new array with zeros in certain locations |
| 873 | (functions like `triu` or `tril` do precisely this). Then this function |
| 874 | returns the indices where the non-zero values would be located. |
| 875 | |
| 876 | Parameters |
| 877 | ---------- |
| 878 | n : int |
| 879 | The returned indices will be valid to access arrays of shape (n, n). |
| 880 | mask_func : callable |
| 881 | A function whose call signature is similar to that of `triu`, `tril`. |
| 882 | That is, ``mask_func(x, k)`` returns a boolean array, shaped like `x`. |
| 883 | `k` is an optional argument to the function. |
| 884 | k : scalar |
| 885 | An optional argument which is passed through to `mask_func`. Functions |
| 886 | like `triu`, `tril` take a second argument that is interpreted as an |
| 887 | offset. |
| 888 | |
| 889 | Returns |
| 890 | ------- |
| 891 | indices : tuple of arrays. |
| 892 | The `n` arrays of indices corresponding to the locations where |
| 893 | ``mask_func(np.ones((n, n)), k)`` is True. |
| 894 | |
| 895 | See Also |
| 896 | -------- |
| 897 | triu, tril, triu_indices, tril_indices |
| 898 | |
| 899 | Examples |
| 900 | -------- |
| 901 | >>> import numpy as np |
| 902 | |
| 903 | These are the indices that would allow you to access the upper triangular |
| 904 | part of any 3x3 array: |
| 905 | |
| 906 | >>> iu = np.mask_indices(3, np.triu) |
| 907 | |
| 908 | For example, if `a` is a 3x3 array: |
| 909 | |
| 910 | >>> a = np.arange(9).reshape(3, 3) |
| 911 | >>> a |
| 912 | array([[0, 1, 2], |
| 913 | [3, 4, 5], |
| 914 | [6, 7, 8]]) |
| 915 | >>> a[iu] |
| 916 | array([0, 1, 2, 4, 5, 8]) |
| 917 | |
| 918 | An offset can be passed also to the masking function. This gets us the |
| 919 | indices starting on the first diagonal right of the main one: |
| 920 | |
| 921 | >>> iu1 = np.mask_indices(3, np.triu, 1) |
| 922 | |
| 923 | with which we now extract only three elements: |
searching dependent graphs…