ENH: Add a function that performs the indexing needed to map argsort to sort #8708

eric-wieser · 2017-02-27T23:15:50Z

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...]

The text was updated successfully, but these errors were encountered:

mhvk · 2017-02-28T01:19:56Z

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 · 2017-02-28T01:27:26Z

I think that changing argsort is the wrong thing to do - this is a compact and obvious representation of the data, and based on the implementation of np.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 · 2017-02-28T01:45:04Z

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 · 2017-02-28T08:27:31Z

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.take style) easily there. But I have serious doubts you can get great performance while making use of cython convenience.

eric-wieser · 2017-02-28T11:21:24Z

OTOH you can probably get half way (a bit like np.take style) easily there.

Can you elaborate on what you mean by that?

seberg · 2017-02-28T12:35:57Z

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 · 2017-02-28T12:40:01Z

Should I just submit a pure-python PR for now then?

seberg · 2017-02-28T12:50:26Z

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 · 2017-02-28T14:11:53Z

@eric-wieser -- you're code is missing one aspect, that to be able to mimic min and max it should be possible to keep or remove the axis. See my implementation in astropy's Time class: 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 the arg* methods: the call then becomes more simlar to max and min.

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]index methods) to accept just the output as it currently is, and let it accept axis and keepdims.

eric-wieser · 2017-02-28T14:15:11Z

your code is missing one aspect, that to be able to mimic min and max it should be possible to keep or remove the axis

This is the choice of argmin and argmax, not take_along_axis. These functions need to learn keepdim, not the indexer. See #8710.

Either way, this can be emulated with:

>>> ai = np.expand_dims(a.argmin(axis=axis), axis=axis)
>>> take_along_axis(a, ai, axis=axis)

This is a more powerful function than just taking keepdim, as it can also insert extra dims

eric-wieser · 2017-02-28T14:21:52Z

My overall preference, though, would be to use take to accept just the output as it currently is, and let it accept axis

take already does accept this output and an axis kwarg, but interprets them completely differently, and produces a different result. We can't change that now.

eric-wieser · 2017-02-28T14:25:26Z

@mhvk: here's an example of why I think this is a better api than just adding arguments to the arg methods

mhvk · 2017-02-28T14:33:07Z

Nothing beats actual code! I'm quite happy with take_along_axis too -- 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 to astropy.utils.compat.numpy rather than wait until our minimum version of numpy is high enough...).

eric-wieser · 2017-02-28T14:40:59Z

@mhvk: In case you missed it, the PR is at #8714

jason-s · 2017-10-31T21:54:07Z

Hmm. I just tried this take_along_axis:

A = np.array([[3,2,1],[4,0,6]])
take_along_axis(A,A.
8000
argsort(axis=-1),axis=-1)

and got this error:

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-127-e9c0d18e1e4c> in <module>()
      1 A = np.array([[3,2,1],[4,0,6]])
----> 2 take_along_axis(A,A.argsort(axis=-1),axis=-1)

<ipython-input-124-81abe9b2cb05> in take_along_axis(arr, ind, axis)
     27             inds.append(np.arange(n).reshape(ind_shape_dim))
     28 
---> 29     return arr[tuple(inds)]

IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (2,3) (3,1)

Oh, duh, never mind, the sample take_along_axis doesn't implement negative dimensions as relative to the end.

jason-s · 2017-10-31T21:59:02Z

this should fix:

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...]
    """
    if axis < 0:
       if axis >= -arr.ndim:
           axis += arr.ndim
       else:
           raise IndexError('axis out of range')
    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)]

eric-wieser · 2017-10-31T23:17:04Z

Using np.core.multiarray.normalize_axis_index(axis, arr.ndim) is a good way to spell that as of 1.13, but good catch.

The PR for this suggestion already handles negative indices correctly

crusaderky · 2018-05-13T21:26:23Z

Hello, any update on this?

mhvk · 2018-05-13T21:50:51Z

@crusaderky - there is a PR, see #8714. But your comment is a reminder to look at that and see what remains to be done -- it would still be good to have this functionality!

This is the reduced version that does not allow any insertion of extra dimensions

mhvk · 2018-05-29T00:44:29Z

fixed by #11105

oulenz · 2019-07-10T11:31:27Z

your code is missing one aspect, that to be able to mimic min and max it should be possible to keep or remove the axis

This is the choice of argmin and argmax, not take_along_axis. These functions need to learn keepdim, not the indexer. See #8710.

Either way, this can be emulated with:
>>> ai = np.expand_dims(a.argmin(axis=axis), axis=axis)
>>> take_along_axis(a, ai, axis=axis)
This is a more powerful function than just taking keepdim, as it can also insert extra dims

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 of keepdim = False? Having to expand and reduce the dimensions just to recover min and max from argmin and argmax is still quite cumbersome.

It seems to me it should be possible to modify take_along_axis to do this automatically if ai has one dimension fewer than a.

owlas · 2019-07-17T18:28:03Z

Awesome!

eric-wieser added the 01 - Enhancement label Feb 27, 2017

This was referenced Feb 27, 2017

BUG: fix ma.median for empty ndarrays #8705

Merged

WIP: Implement oindex #6075

Closed

eric-wieser mentioned this issue Feb 28, 2017

ENH: Implement take_along_axis as described in #8708 #8714

Closed

eric-wieser mentioned this issue Jul 10, 2017

Feature request: Incremental bincount and axis argument #9397

Open

eric-wieser mentioned this issue Jan 3, 2018

argsort does not work for multidimensional arrays #4724

Closed

eric-wieser added a commit to eric-wieser/numpy that referenced this issue May 16, 2018

ENH: Add (put|take)_along_axis as described in numpy#8708

7a87604

This is the reduced version that does not allow any insertion of extra dimensions

eric-wieser added a commit to eric-wieser/numpy that referenced this issue May 26, 2018

ENH: Add (put|take)_along_axis as described in numpy#8708

905e906

This is the reduced version that does not allow any insertion of extra dimensions

mhvk closed this as completed May 29, 2018

kgryte mentioned this issue Jun 27, 2024

RFC: add take_along_axis to take values along a specified dimension data-apis/array-api#808

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

ENH: Add a function that performs the indexing needed to map argsort to sort #8708

ENH: Add a function that performs the indexing needed to map argsort to sort #8708

ENH: Add a function that performs the indexing needed to map argsort to sort #8708

ENH: Add a function that performs the indexing needed to map argsort to sort #8708

Comments