Repository navigation
ENH: Add a function that performs the indexing needed to map argsort to sort #8708
Description
There've been two prs recently ( #8703, #8678 ) that have needed this function, and in general it's needed for any kind of vectorized function returning indices along an axis (argsort, argpartition, argmin, argmax, count, ...)
A python implementation looks something like this:
def take_along_axis(arr, ind, axis):
"""
... here means a "pack" of dimensions, possibly empty
arr: array_like of shape (A..., M, B...)
source array
ind: array_like of shape (A..., K..., B...)
indices to take along each 1d slice of `arr`
axis: int
index of the axis with dimension M
out: array_like of shape (A..., K..., B...)
out[a..., k..., b...] = arr[a..., inds[a..., k..., b...], b...]
"""
ind_shape = (1,) * ind.ndim
ins_ndim = ind.ndim - (arr.ndim - 1) #inserted dimensions
dest_dims = list(range(axis)) + [None] + list(range(axis+ins_ndim, ind.ndim))
# could also call np.ix_ here with some dummy arguments, then throw those results away
inds = []
for dim, n in zip(dest_dims, arr.shape):
if dim is None:
inds.append(ind)
else:
ind_shape_dim = ind_shape[:dim] + (-1,) + ind_shape[dim+1:]
inds.append(np.arange(n).reshape(ind_shape_dim))
return arr[tuple(inds)]Which works as intended:
np.take_along_axis(a, a.argsort(axis=axis), axis=axis) == a.sort(axis=axis)
np.take_along_axis(a, a.argmin(axis=axis), axis=axis) == a.min(axis=axis)
np.take_along_axis(a, a.argmax(axis=axis), axis=axis) == a.max(axis=axis)As far as I can tell, np.take isn't suitable for this right now, as instead it computes:
out[a..., k..., b...] = arr[a..., inds[k...], b...]instead of
out[a..., k..., b...] = arr[a..., inds[a..., k..., b...], b...]Perhaps np.take(..., broadcast=True) should b a thing that means the above?
If not, there's perhaps enough similarity with apply_along_axis to justify the similar name:
apply_along_axis:out[a..., k..., b...] = f(arr[a..., :, b...])take_along_axis:out[a..., k..., b...] = arr[a..., inds[a..., k..., b...], b...]
Activity
mhvk commentedon Feb 28, 2017
See #6078 for what I think is the same issue; that has also some code snippets from myself and @shoyer as well as a comment by @jaimefrio that maybe the best approach would be for the
arg*methods to optionally return such a complete index.eric-wieser commentedon Feb 28, 2017
I think that changing
argsortis the wrong thing to do - this is a compact and obvious representation of the data, and based on the implementation ofnp.median, one that can appear from other sources.So I'd be strongly in favor of either making a method for this, or re-purposing an existing one. I also suspect that going through an index tuple leads to slower indexing than writing a C function to directly handle the loop described above, since looking up an index in
np.arange()is needlessly slow compared to directly using that index.eric-wieser commentedon Feb 28, 2017
Is this a good candidate for a Cython function, to implement this as just a
out[a..., k..., b...] = arr[a..., inds[a..., k..., b...], b...]loop over all indices?seberg commentedon Feb 28, 2017
Personally doubt it, since its probably a plan to use the numpy iteration API, and I am not sure there is any convenience in doing that in cython. OTOH you can probably get half way (a bit like
np.takestyle) easily there. But I have serious doubts you can get great performance while making use of cython convenience.eric-wieser commentedon Feb 28, 2017
Can you elaborate on what you mean by that?
seberg commentedon Feb 28, 2017
I mean performance on larger arrays is probably not very good, just like np.take comparing to advanced indexing. Except in some cases (here probably small dimensions or so), which don't optimize well then doing the fancier npyiter.
eric-wieser commentedon Feb 28, 2017
Should I just submit a pure-python PR for now then?
seberg commentedon Feb 28, 2017
If you like, didn't make up my mind what I think the best approach.
If you got time to waste, you could also check out npyiter stuff, add the dimension (inside npyiter) to the indexed array, iterate them together (+the result array) and then write the small inner loop necessary for copying and index checking.
mhvk commentedon Feb 28, 2017
@eric-wieser -- you're code is missing one aspect, that to be able to mimic
minandmaxit should be possible to keep or remove the axis. See my implementation in astropy'sTimeclass: https://github.com/astropy/astropy/blob/master/astropy/time/core.py#L915. In a sense, this is also an argument in favour of making this an option to thearg*methods: the call then becomes more simlar tomaxandmin.Note that yet another place where one could add this to is
unravel_index-- although that would not be the most obvious place a novice would look. That again is an argument for having it on the methods themselves.My overall preference, though, would be to use
take(or perhaps extend one of the new[ova]indexmethods) to accept just the output as it currently is, and let it acceptaxisandkeepdims.eric-wieser commentedon Feb 28, 2017
This is the choice of
argminandargmax, nottake_along_axis. These functions need to learnkeepdim, not the indexer. See #8710.Either way, this can be emulated with:
This is a more powerful function than just taking
keepdim, as it can also insert extra dimseric-wieser commentedon Feb 28, 2017
takealready does accept this output and anaxiskwarg, but interprets them completely differently, and produces a different result. We can't change that now.eric-wieser commentedon Feb 28, 2017
@mhvk: here's an example of why I think this is a better api than just adding arguments to the
argmethodsmhvk commentedon Feb 28, 2017
Nothing beats actual code! I'm quite happy with
take_along_axistoo -- really nice just to have the ability!! (also, from a practical point of view, easier to start using it, since we can just backport it toastropy.utils.compat.numpyrather than wait until our minimum version of numpy is high enough...).16 remaining items
mhvk commentedon May 29, 2018
fixed by #11105
oulenz commentedon Jul 10, 2019
Correct me if I'm missing something, but while what you say is true in case you want the behaviour of
keepdim = True, what if you want the behaviour ofkeepdim = False? Having to expand and reduce the dimensions just to recoverminandmaxfromargminandargmaxis still quite cumbersome.It seems to me it should be possible to modify
take_along_axisto do this automatically ifaihas one dimension fewer thana.owlas commentedon Jul 17, 2019
Awesome!