xorbits.numpy.mask_indices#
- xorbits.numpy.mask_indices(n, mask_func, k=0)#
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 – The n arrays of indices corresponding to the locations where
mask_func(np.ones((n, n)), k)
is True.- Return type
tuple of arrays.
See also
Notes
New in version 1.4.0(numpy).
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])
Warning
This method has not been implemented yet. Xorbits will try to execute it with numpy.
This docstring was copied from numpy.