[Numpy-discussion] Adding take_along_axis and put_along_axis functions

[Eric Wieser]

These functions provide a vectorized way of using one array to look up
items in another. In particular, they extend the 1d:

a = np.array([4, 5, 6, 1, 2, 3])
b = np.array(["four", "five", "six", "one", "two", "three"])
i = a.argsort()
b_sorted = b[i]

To work for higher-dimensions:

a = np.array([[4, 1], [5, 2], [6, 3]])
b = np.array([["four", "one"],  ["five", "two"], ["six", "three"]])
i = a.argsort(axis=1)
b_sorted = np.take_along_axis(b, i, axis=1)

put_along_axis is the obvious but less useful dual to this operation,
inserting elements rather than extracting them. (Unlike put and take which
are not obvious duals).

These have been merged in gh-11105
<https://github.com/numpy/numpy/pull/11105>, but as a new addition this
probably should have gone by the mailing list first.

There was a lack of consensus in gh-8714
<https://github.com/numpy/numpy/pull/8714> about how best to generalize to
differing dimensions, so only the non-controversial case where the indices
and array have the same dimensions was implemented.

These names were chosen to mirror apply_along_axis, which behaves
similarly. Do they seem reasonable?

Eric
​
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.python.org/pipermail/numpy-discussion/attachments/20180528/9b5aa5b9/attachment.html>