Adding take_along_axis and put_along_axis functions

classic Classic list List threaded Threaded
2 messages Options
Reply | Threaded
Open this post in threaded view
|

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, but as a new addition this probably should have gone by the mailing list first.

There was a lack of consensus in gh-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


_______________________________________________
NumPy-Discussion mailing list
[hidden email]
https://mail.python.org/mailman/listinfo/numpy-discussion
Reply | Threaded
Open this post in threaded view
|

Re: Adding take_along_axis and put_along_axis functions

Stephan Hoyer-2
As I'm sure I stated in the GItHub discussion, I strongly support adding these functions to NumPy. This logic is non-trivial to get right and is quite broadly useful.

These names also seem natural to me.

On Mon, May 28, 2018 at 8:07 PM Eric Wieser <[hidden email]> wrote:

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, but as a new addition this probably should have gone by the mailing list first.

There was a lack of consensus in gh-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

_______________________________________________
NumPy-Discussion mailing list
[hidden email]
https://mail.python.org/mailman/listinfo/numpy-discussion

_______________________________________________
NumPy-Discussion mailing list
[hidden email]
https://mail.python.org/mailman/listinfo/numpy-discussion