numpy.mask_indices¶

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 liketriu
ortril
do precisely this). Then this function returns the indices where the nonzero values would be located. Parameters
 nint
The returned indices will be valid to access arrays of shape (n, n).
 mask_funccallable
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. kscalar
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
 indicestuple of arrays.
The n arrays of indices corresponding to the locations where
mask_func(np.ones((n, n)), k)
is True.
See also
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])