flopscope.

flopscope.numpy.where

fnp.where(condition, x=None, y=None)[flopscope source]

Return elements chosen from `x` or `y` depending on `condition`.

Adapted from NumPy docs np.where

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

Select elements based on condition. Cost: numel(input).

Return elements chosen from x or y depending on condition.

Note.

When only condition is provided, this function is a shorthand for flops.asarray(condition).nonzero(). Using nonzero directly should be preferred, as it behaves correctly for subclasses. The rest of this documentation covers only the case where all three arguments are provided.

Parameters

condition:array_like, bool

Where True, yield x, otherwise yield y.

x, y:array_like

Values from which to choose. x, y and condition need to be broadcastable to some shape.

Returns

out:ndarray

An array with elements from x where condition is True, and elements from y elsewhere.

See also

Notes

If all the arrays are 1-D, where is equivalent to:

[xv if c else yv
 for c, xv, yv in zip(condition, x, y)]

Examples

>>> import flopscope.numpy as fnp
>>> a = flops.arange(10)
>>> a
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> flops.where(a < 5, a, 10*a)
array([ 0,  1,  2,  3,  4, 50, 60, 70, 80, 90])

This can be used on multidimensional arrays too:

>>> flops.where([[True, False], [True, True]],
... [[1, 2], [3, 4]],
... [[9, 8], [7, 6]])
array([[1, 8],
       [3, 4]])

The shapes of x, y, and the condition are broadcast together:

>>> x, y = flops.ogrid[:3, :4]
>>> flops.where(x < y, x, 10 + y)  # both x and 10+y are broadcast
array([[10,  0,  0,  0],
       [10, 11,  1,  1],
       [10, 11, 12,  2]])
>>> a = flops.array([[0, 1, 2],
... [0, 2, 4],
... [0, 3, 6]])
>>> flops.where(a < 4, a, -1)  # -1 is broadcast
array([[ 0,  1,  2],
       [ 0,  2, -1],
       [ 0,  3, -1]])