diff options
| -rw-r--r-- | numpy/add_newdocs.py | 77 | ||||
| -rw-r--r-- | numpy/core/fromnumeric.py | 21 | ||||
| -rw-r--r-- | numpy/lib/arraysetops.py | 8 | ||||
| -rw-r--r-- | numpy/ma/core.py | 25 |
4 files changed, 72 insertions, 59 deletions
diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py index 09cae54b1..62b63454e 100644 --- a/numpy/add_newdocs.py +++ b/numpy/add_newdocs.py @@ -1577,71 +1577,72 @@ add_newdoc('numpy.core.multiarray', 'where', """ where(condition, [x, y]) - Return elements, either from `x` or `y`, depending on `condition`. + Return elements chosen from `x` or `y` depending on `condition`. - If only `condition` is given, return ``condition.nonzero()``. + .. note:: + When only `condition` is provided, this function is a shorthand for + ``np.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 - When True, yield `x`, otherwise yield `y`. - x, y : array_like, optional + 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 or tuple of ndarrays - If both `x` and `y` are specified, the output array contains - elements of `x` where `condition` is True, and elements from - `y` elsewhere. - - If only `condition` is given, return the tuple - ``condition.nonzero()``, the indices where `condition` is True. + out : ndarray + An array with elements from `x` where `condition` is True, and elements + from `y` elsewhere. See Also -------- - nonzero, choose + choose + nonzero : The function that is called when x and y are omitted Notes ----- - If `x` and `y` are given and input arrays are 1-D, `where` is - equivalent to:: + 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)] + [xv if c else yv + for c, xv, yv in zip(condition, x, y)] Examples -------- + >>> a = np.arange(10) + >>> a + array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + >>> np.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: + >>> np.where([[True, False], [True, True]], ... [[1, 2], [3, 4]], ... [[9, 8], [7, 6]]) array([[1, 8], [3, 4]]) - >>> np.where([[0, 1], [1, 0]]) - (array([0, 1]), array([1, 0])) - - >>> x = np.arange(9.).reshape(3, 3) - >>> np.where( x > 5 ) - (array([2, 2, 2]), array([0, 1, 2])) - >>> x[np.where( x > 3.0 )] # Note: result is 1D. - array([ 4., 5., 6., 7., 8.]) - >>> np.where(x < 5, x, -1) # Note: broadcasting. - array([[ 0., 1., 2.], - [ 3., 4., -1.], - [-1., -1., -1.]]) - - Find the indices of elements of `x` that are in `goodvalues`. - - >>> goodvalues = [3, 4, 7] - >>> ix = np.isin(x, goodvalues) - >>> ix - array([[False, False, False], - [ True, True, False], - [False, True, False]]) - >>> np.where(ix) - (array([1, 1, 2]), array([0, 1, 1])) + The shapes of x, y, and the condition are broadcast together: + >>> x, y = np.ogrid[:3, :4] + >>> np.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 = np.array([[0, 1, 2], + ... [0, 2, 4], + ... [0, 3, 6]]) + >>> np.where(a < 4, a, -1) # -1 is broadcast + array([[ 0, 1, 2], + [ 0, 2, -1], + [ 0, 3, -1]]) """) diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py index 5b67a0dc5..373e0fde8 100644 --- a/numpy/core/fromnumeric.py +++ b/numpy/core/fromnumeric.py @@ -1615,16 +1615,16 @@ def nonzero(a): Examples -------- - >>> x = np.array([[1,0,0], [0,2,0], [1,1,0]]) + >>> x = np.array([[3, 0, 0], [0, 4, 0], [5, 6, 0]]) >>> x - array([[1, 0, 0], - [0, 2, 0], - [1, 1, 0]]) + array([[3, 0, 0], + [0, 4, 0], + [5, 6, 0]]) >>> np.nonzero(x) (array([0, 1, 2, 2]), array([0, 1, 0, 1])) >>> x[np.nonzero(x)] - array([1, 2, 1, 1]) + array([3, 4, 5, 6]) >>> np.transpose(np.nonzero(x)) array([[0, 0], [1, 1], @@ -1636,7 +1636,7 @@ def nonzero(a): boolean array and since False is interpreted as 0, np.nonzero(a > 3) yields the indices of the `a` where the condition is true. - >>> a = np.array([[1,2,3],[4,5,6],[7,8,9]]) + >>> a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) >>> a > 3 array([[False, False, False], [ True, True, True], @@ -1644,7 +1644,14 @@ def nonzero(a): >>> np.nonzero(a > 3) (array([1, 1, 1, 2, 2, 2]), array([0, 1, 2, 0, 1, 2])) - The ``nonzero`` method of the boolean array can also be called. + Using this result to index `a` is equivalent to using the mask directly: + + >>> a[np.nonzero(a > 3)] + array([4, 5, 6, 7, 8, 9]) + >>> a[a > 3] # prefer this spelling + array([4, 5, 6, 7, 8, 9]) + + ``nonzero`` can also be called as a method of the array. >>> (a > 3).nonzero() (array([1, 1, 1, 2, 2, 2]), array([0, 1, 2, 0, 1, 2])) diff --git a/numpy/lib/arraysetops.py b/numpy/lib/arraysetops.py index 4d3f35183..5880ea154 100644 --- a/numpy/lib/arraysetops.py +++ b/numpy/lib/arraysetops.py @@ -607,6 +607,14 @@ def isin(element, test_elements, assume_unique=False, invert=False): [ True, False]]) >>> element[mask] array([2, 4]) + + The indices of the matched values can be obtained with `nonzero`: + + >>> np.nonzero(mask) + (array([0, 1]), array([1, 0])) + + The test can also be inverted: + >>> mask = np.isin(element, test_elements, invert=True) >>> mask array([[ True, False], diff --git a/numpy/ma/core.py b/numpy/ma/core.py index 091ab4e20..5bfa51b12 100644 --- a/numpy/ma/core.py +++ b/numpy/ma/core.py @@ -7115,32 +7115,32 @@ size.__doc__ = np.size.__doc__ def where(condition, x=_NoValue, y=_NoValue): """ - Return a masked array with elements from x or y, depending on condition. + Return a masked array with elements from `x` or `y`, depending on condition. - Returns a masked array, shaped like condition, where the elements - are from `x` when `condition` is True, and from `y` otherwise. - If neither `x` nor `y` are given, the function returns a tuple of - indices where `condition` is True (the result of - ``condition.nonzero()``). + .. note:: + When only `condition` is provided, this function is identical to + `nonzero`. The rest of this documentation covers only the case where + all three arguments are provided. Parameters ---------- condition : array_like, bool - The condition to meet. For each True element, yield the corresponding - element from `x`, otherwise from `y`. + Where True, yield `x`, otherwise yield `y`. x, y : array_like, optional Values from which to choose. `x`, `y` and `condition` need to be broadcastable to some shape. Returns ------- - out : MaskedArray or tuple of ndarrays - The resulting masked array if `x` and `y` were given, otherwise - the result of ``condition.nonzero()``. + out : MaskedArray + An masked array with `masked` elements where the condition is masked, + elements from `x` where `condition` is True, and elements from `y` + elsewhere. See Also -------- numpy.where : Equivalent function in the top-level NumPy module. + nonzero : The function that is called when x and y are omitted Examples -------- @@ -7151,9 +7151,6 @@ def where(condition, x=_NoValue, y=_NoValue): [[0.0 -- 2.0] [-- 4.0 --] [6.0 -- 8.0]] - >>> np.ma.where(x > 5) # return the indices where x > 5 - (array([2, 2]), array([0, 2])) - >>> print(np.ma.where(x > 5, x, -3.1416)) [[-3.1416 -- -3.1416] [-- -3.1416 --] |