dpnp.mask_indices

dpnp.mask_indices(n, mask_func, k=0, device=None, usm_type='device', sycl_queue=None)[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 as mask_func(a, k=k) returns a new array with zeros in certain locations (functions like dpnp.triu or dpnp.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 dpnp.triu, dpnp.tril. That is, mask_func(x, k=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 dpnp.triu, dpnp.tril take a second argument that is interpreted as an offset. Default: 0.

  • device ({None, string, SyclDevice, SyclQueue}, optional) -- An array API concept of device where the output array is created. The device can be None (the default), an OneAPI filter selector string, an instance of dpctl.SyclDevice corresponding to a non-partitioned SYCL device, an instance of dpctl.SyclQueue, or a Device object returned by dpnp.dpnp_array.dpnp_array.device property.

  • usm_type ({"device", "shared", "host"}, optional) -- The type of SYCL USM allocation for the output array.

  • sycl_queue ({None, SyclQueue}, optional) -- A SYCL queue to use for output array allocation and copying.

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 dpnp.ndarray

See also

dpnp.tril

Return lower triangle of an array.

dpnp.triu

Return upper triangle of an array.

dpnp.triu_indices

Return the indices for the upper-triangle of an (n, m) array.

dpnp.tril_indices

Return the indices for the lower-triangle of an (n, m) array.

Examples

These are the indices that would allow you to access the upper triangular part of any 3x3 array:

>>> import dpnp as np
>>> 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])