summaryrefslogtreecommitdiff
path: root/numpy
diff options
context:
space:
mode:
Diffstat (limited to 'numpy')
-rw-r--r--numpy/add_newdocs.py77
-rw-r--r--numpy/core/fromnumeric.py21
-rw-r--r--numpy/core/src/multiarray/einsum.c.src28
-rw-r--r--numpy/core/tests/test_einsum.py5
-rw-r--r--numpy/lib/arraysetops.py8
-rw-r--r--numpy/ma/core.py25
6 files changed, 87 insertions, 77 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/core/src/multiarray/einsum.c.src b/numpy/core/src/multiarray/einsum.c.src
index 33184d99a..1765982a0 100644
--- a/numpy/core/src/multiarray/einsum.c.src
+++ b/numpy/core/src/multiarray/einsum.c.src
@@ -2767,11 +2767,11 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
goto fail;
}
- /* Initialize the output to all zeros and reset the iterator */
+ /* Initialize the output to all zeros */
ret = NpyIter_GetOperandArray(iter)[nop];
- Py_INCREF(ret);
- PyArray_AssignZero(ret, NULL);
-
+ if (PyArray_AssignZero(ret, NULL) < 0) {
+ goto fail;
+ }
/***************************/
/*
@@ -2785,16 +2785,12 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
case 1:
if (ndim == 2) {
if (unbuffered_loop_nop1_ndim2(iter) < 0) {
- Py_DECREF(ret);
- ret = NULL;
goto fail;
}
goto finish;
}
else if (ndim == 3) {
if (unbuffered_loop_nop1_ndim3(iter) < 0) {
- Py_DECREF(ret);
- ret = NULL;
goto fail;
}
goto finish;
@@ -2803,16 +2799,12 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
case 2:
if (ndim == 2) {
if (unbuffered_loop_nop2_ndim2(iter) < 0) {
- Py_DECREF(ret);
- ret = NULL;
goto fail;
}
goto finish;
}
else if (ndim == 3) {
if (unbuffered_loop_nop2_ndim3(iter) < 0) {
- Py_DECREF(ret);
- ret = NULL;
goto fail;
}
goto finish;
@@ -2823,7 +2815,6 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
/***************************/
if (NpyIter_Reset(iter, NULL) != NPY_SUCCEED) {
- Py_DECREF(ret);
goto fail;
}
@@ -2845,8 +2836,6 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
if (sop == NULL) {
PyErr_SetString(PyExc_TypeError,
"invalid data type for einsum");
- Py_DECREF(ret);
- ret = NULL;
}
else if (NpyIter_GetIterSize(iter) != 0) {
NpyIter_IterNextFunc *iternext;
@@ -2858,7 +2847,6 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
iternext = NpyIter_GetIterNext(iter, NULL);
if (iternext == NULL) {
NpyIter_Deallocate(iter);
- Py_DECREF(ret);
goto fail;
}
dataptr = NpyIter_GetDataPtrArray(iter);
@@ -2874,12 +2862,16 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
/* If the API was needed, it may have thrown an error */
if (NpyIter_IterationNeedsAPI(iter) && PyErr_Occurred()) {
- Py_DECREF(ret);
- ret = NULL;
+ goto fail;
}
}
finish:
+ if (out != NULL) {
+ ret = out;
+ }
+ Py_INCREF(ret);
+
NpyIter_Deallocate(iter);
for (iop = 0; iop < nop; ++iop) {
Py_DECREF(op[iop]);
diff --git a/numpy/core/tests/test_einsum.py b/numpy/core/tests/test_einsum.py
index 647738831..a72079218 100644
--- a/numpy/core/tests/test_einsum.py
+++ b/numpy/core/tests/test_einsum.py
@@ -730,6 +730,11 @@ class TestEinSum(object):
res = np.einsum('...ij,...jk->...ik', a, a, out=out)
assert_equal(res, tgt)
+ def test_out_is_res(self):
+ a = np.arange(9).reshape(3, 3)
+ res = np.einsum('...ij,...jk->...ik', a, a, out=a)
+ assert res is a
+
def optimize_compare(self, string):
# Tests all paths of the optimization function against
# conventional einsum
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 --]
View Lorry Controller Status