flopscope.

flopscope.numpy.mask_indices

fnp.mask_indices(n, mask_func, k=0)[flopscope source][numpy source]

Return the indices to access (n, n) arrays, given a masking function.

Adapted from NumPy docs np.mask_indices

Areacore
Typecustom
NumPy Refnp.mask_indices
Cost
4×per-operation
Flopscope Context

Return indices of mask for n x n array. Cost: numel(output).

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 (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 where mask_func(flops.ones((n, n)), k) is True.

See also

Examples

>>> import flopscope.numpy as fnp

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

>>> iu = flops.mask_indices(3, flops.triu)

For example, if a is a 3x3 array:

>>> a = flops.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 = flops.mask_indices(3, flops.triu, 1)

with which we now extract only three elements:

>>> a[iu1]
array([1, 2, 5])