numpy.mask_indices — NumPy v1.15 Manual (original) (raw)
numpy. mask_indices(n, mask_func, k=0)[source]¶
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 asmask_func(a, k) returns a new array with zeros in certain locations (functions like triu or tril do precisely this). Then this function returns the indices where the non-zero values would be located.
| Parameters: | n : int The returned indices will be valid to access arrays of shape (n, n). mask_func : callable A function whose call signature is similar to that of triu, tril. That is, mask_func(x, k) returns a boolean array, shaped like x.k is an optional argument to the function. k : scalar An optional argument which is passed through to mask_func. Functions like triu, tril take a second argument that is interpreted as an offset. |
|---|---|
| Returns: | indices : tuple of arrays. The n arrays of indices corresponding to the locations wheremask_func(np.ones((n, n)), k) is True. |
Notes
New in version 1.4.0.
Examples
These are the indices that would allow you to access the upper triangular part of any 3x3 array:
iu = np.mask_indices(3, np.triu)
For example, if a is a 3x3 array:
a = np.arange(9).reshape(3, 3) a array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) a[iu] array([0, 1, 2, 4, 5, 8])
An offset can be passed also to the masking function. This gets us the indices starting on the first diagonal right of the main one:
iu1 = np.mask_indices(3, np.triu, 1)
with which we now extract only three elements:
a[iu1] array([1, 2, 5])