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
_______________________________________________
NumPy-Discussion mailing list
NumPy-Discussion@python.org
https://mail.python.org/mailman/listinfo/numpy-discussion

Reply via email to