DOC: Add take_along_axis to the see also section in argmin, argmax etc. #14799

mproszewska · 2019-10-29T17:46:17Z

This closes #6078.

…t and argpartition.

eric-wieser · 2019-10-29T18:37:40Z

numpy/core/fromnumeric.py

@@ -797,6 +797,7 @@ def argpartition(a, kth, axis=-1, kind='introselect', order=None):
    partition : Describes partition algorithms used.
    ndarray.partition : Inplace partition.
    argsort : Full indirect sort
+    take_along_axis : Take values from the input array by matching 1d index and data slices.


Let's see how another maintainer feels about this first, but perhaps:

Suggested change

take_along_axis : Take values from the input array by matching 1d index and data slices.

take_along_axis : Apply ``index_array`` from argpartition to an array as if by calling partition

And same for the below,

take_along_axis : Apply ``index_array`` from argsort to an array as if by calling sort take_along_axis : Apply ``np.expand_dims(index_array, axis)`` from argmin to an array as if by calling min take_along_axis : Apply ``np.expand_dims(index_array, axis)`` from argmax to an array as if by calling max

Probably those use cases are worth mentioning, but I wasn't sure if the "See also" section is a right place for that. e.g. argsort and argpartition refer to take_along_axis in "Returns". Perhaps it's better to just add an example.

An example for each of these showing how to use take_along_axis would be great, if you're willing to spend the time. If not, we can merge this as is, and leave the issue open for someone else to add the examples.

"See Also" seems appropriate here.

numpy/core/fromnumeric.py

eric-wieser · 2019-11-01T00:06:23Z

numpy/core/fromnumeric.py

@@ -1223,6 +1239,13 @@ def argmin(a, axis=None, out=None):
    >>> np.argmin(b)  # Only the first occurrence is returned.
    0

+    >>> x = np.array([[4,2,3],[1,0,3]])
+    >>> np.take_along_axis(x, np.expand_dims(np.argmin(x, axis=-1), axis=-1), axis=-1)


Possibly worth splitting these examples in two, either as:

Suggested change

>>> np.take_along_axis(x, np.expand_dims(np.argmin(x, axis=-1), axis=-1), axis=-1)

>>> index_array = np.argmin(x, axis=-1)

>>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1)

or

Suggested change

>>> np.take_along_axis(x, np.expand_dims(np.argmin(x, axis=-1), axis=-1), axis=-1)

>>> index_array = np.expand_dims(np.argmin(x, axis=-1), axis=-1) # emulate keepdims=True

>>> np.take_along_axis(x, index_array, axis=-1)

Co-Authored-By: Eric Wieser <wieser.eric@gmail.com>

numpy/core/fromnumeric.py

Co-Authored-By: Eric Wieser <wieser.eric@gmail.com>

eric-wieser

Looks good, will let someone else merge since I suggested a little too much of the wording to be impartial.

eric-wieser · 2019-11-03T21:15:50Z

numpy/core/fromnumeric.py

+    >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1) # same as np.min(x, axis=-1, keepdims=True)
+    array([[2],
+           [0]])
+    >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1).squeeze(axis=-1) # same as np.max(x, axis=-1)


Nit: pep8 wants two spaces before #, same on other lines.

Thanks, I was looking for that bug

mattip · 2019-11-03T21:31:02Z

numpy/core/fromnumeric.py

+           [3]])
+    >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1).squeeze(axis=-1) # same as np.max(x, axis=-1)
+    array([4, 3])
+


These lines are quite long. Maybe move the comment to stand alone in the line above the code

mattip · 2019-11-03T21:32:09Z

A small nit about line length, but otherwise this looks good

mattip · 2019-11-04T22:19:11Z

Thanks @mproszewska

Add take_along_axis to the see also section in argmin, argmax, argsor…

fbf8ae4

…t and argpartition.

eric-wieser added the 04 - Documentation label Oct 29, 2019

eric-wieser reviewed Oct 29, 2019

View reviewed changes

mproszewska added 2 commits November 1, 2019 00:40

Change description of method and add examples.

020d3f2

Better argpartition example.

508d384