On 01-11-2019 09:51, Allan Haldane wrote:
my thought was to try `take` or `take_along_axis`:
ind = np.argmin(a, axis=1) np.take_along_axis(a, ind[:,None], axis=1)
But those functions tend to simply fall back to fancy indexing, and are pretty slow. On my system plain fancy indexing is fastest:
%timeit a[np.arange(N),ind] 1.58 µs ± 18.1 ns per loop %timeit np.take_along_axis(a, ind[:,None], axis=1) 6.49 µs ± 57.3 ns per loop %timeit np.min(a, axis=1) 9.51 µs ± 64.1 ns per loop
Probably `take_along_axis` was designed with uses like yours in mind, but it is not very optimized.
Hi Allan, after scanning the documentation once more I found `take_along_axis` and was hoping that it implements some smart trick that does not involve generating and indexing array, but apparently that is what it does. Given the current numpy primitives, I don't see a way to optimize it further and keep it generic. I think the direct fancy indexing is faster in your case because of overhead in handling the generic case and not because of algorithmic inefficiency (from the run times you report it seems that your test array was fairly small). Thank you. Cheers, Dan