8000 DOC: describe the expansion of take and apply_along_axis in detail by eric-wieser · Pull Request #9946 · numpy/numpy · GitHub
[go: up one dir, main page]
Skip to content

DOC: describe the expansion of take and apply_along_axis in detail #9946

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Nov 28, 2017
Merged

DOC: describe the expansion of take and apply_along_axis in detail #9946

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
DOC: describe the expansion of take and apply_along_axis in detail 
Extracted from gh-8714

[ci-skip]
  • Loading branch information
8000
@eric-wieser
eric-wieser committed Nov 22, 2017
commit 21ef1383cb4f6e27af188a6da5cdca93cff1bd07
44 changes: 37 additions & 7 deletions numpy/core/fromnumeric.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,15 +66,28 @@ def take(a, indices, axis=None, out=None, mode='raise'):
"""
Take elements from an array along an axis.

This function does the same thing as "fancy" indexing (indexing arrays
using arrays); however, it can be easier to use if you need elements
along a given axis.
When axis is not None, this function does the same thing as "fancy"
indexing (indexing arrays using arrays); however, it can be easier to use
if you need elements along a given axis. A call such as
``np.take(arr, indices, axis=3)`` is equivalent to
``arr[:,:,:,indices,...]``.

Explained without fancy indexing, this is equivalent to the following use
of `ndindex`, which sets each of ``ii``, ``jj``, and ``kk`` to a tuple of
indices::

Ni, Nk = a.shape[:axis], a.shape[axis+1:]
Nj = indices.shape
for ii in ndindex(Ni):
for jj in ndindex(Nj):
for kk in ndindex(Nk):
out[ii + jj + kk] = a[ii + (indices[jj],) + kk]
Copy link
Member Author
@eric-wieser eric-wieser Nov 21, 2017 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Alternatively:

out[(*ii, *jj, *kk)] = a[(*ii, indices[jj], *kk)]

It's a shame the unpacking generalizations don't apply to tuples within indexing brackets

Copy link
Member
@ahaldane ahaldane Nov 21, 2017 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, while we are using s_, another way to express take is

arr[(s_[:],)*axis + (indices, Ellipsis)]

which is essentially what my comment in the intro of the docstring was saying about a[:,:,:,indices,...].

What about adding something like this to the docstring? :

     This function does the same thing as "fancy" indexing (indexing arrays
     using arrays); however, it can be easier to use if you need elements
     along a given axis and supports `out` and `mode` arguments.

     That is, a call such as ``np.take(arr, indices, axis=3)`` is equivalent to
     ``arr[:,:,:,indices,...]``. More explicitly, `take` is equivalent to the 
     following use of `ndindex` which sets each of ``ii``, ``jj``, and ``kk``
     to a tuple of indices::

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added, with some rephrasing


Parameters
----------
a : array_like
a : array_like (Ni..., M, Nk...)
The source array.
indices : array_like
indices : array_like (Nj...)
The indices of the values to extract.

.. versionadded:: 1.8.0
Expand All @@ -83,7 +96,7 @@ def take(a, indices, axis=None, out=None, mode='raise'):
axis : int, optional
The axis over which to select values. By default, the flattened
input array is used.
out : ndarray, optional
out : ndarray, optional (Ni..., Nj..., Nk...)
If provided, the result will be placed in this array. It should
be of the appropriate shape and dtype.
mode : {'raise', 'wrap', 'clip'}, optional
Expand All @@ -99,14 +112,31 @@ def take(a, indices, axis=None, out=None, mode='raise'):

Returns
-------
subarray : ndarray
out : ndarray (Ni..., Nj..., Nk...)
The returned array has the same type as `a`.

See Also
--------
compress : Take elements using a boolean mask
ndarray.take : equivalent method

Notes
-----

By eliminating the inner loop in the description above, and using `s_` to
build simple slice objects, `take` can be expressed in terms of applying
fancy indexing to each 1-d slice::

Ni, Nk = a.shape[:axis], a.shape[axis+1:]
for ii in ndindex(Ni):
for kk in ndindex(Nj):
out[ii + s_[...,] + kk] = a[ii + s_[:,] + kk][indices]
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As in the above comment, could be:

out[(*ii, ..., *kk)] = a[(*ii, s_[:], *kk)]

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think there is a little less mental overhead for tuple-addition than for tuple-unpacking, especially for people new to python.


For this reason, it is equivalent to (but faster than) the following use
of `apply_along_axis`::

out = np.apply_along_axis(lambda a_1d: a_1d[indices], axis, a)

Examples
--------
>>> a = [4, 3, 5, 7, 6, 8]
Expand Down
28 changes: 23 additions & 5 deletions numpy/lib/shape_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,14 +27,32 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):
Execute `func1d(a, *args)` where `func1d` operates on 1-D arrays and `a`
is a 1-D slice of `arr` along `axis`.

This is equivalent to (but faster than) the following use of `ndindex` and
`s_`, which sets each of ``ii``, ``jj``, and ``kk`` to a tuple of indices::

Ni, Nk = a.shape[:axis], a.shape[axis+1:]
for ii in ndindex(Ni):
for kk in ndindex(Nk):
f = func1d(arr[ii + s_[:,] + kk])
Nj = f.shape
for jj in ndindex(Nj):
out[ii + jj + kk] = f[jj]

Equivalently, eliminating the inner loop, this can be expressed as::

Ni, Nk = a.shape[:axis], a.shape[axis+1:]
for ii in ndindex(Ni):
for kk in ndindex(Nk):
out[ii + s_[...,] + kk] = func1d(arr[ii + s_[:,] + kk])
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am tempted to skip the first longer version, and only leave this one, which I think is clearer.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The longer version is more explicit, and is easier to compare to take IMO. I'd be happy to switch the order of these though.


Parameters
----------
func1d : function
func1d : function (M,) -> (Nj...)
This function should accept 1-D arrays. It is applied to 1-D
slices of `arr` along the specified axis.
axis : integer
Axis along which `arr` is sliced.
arr : ndarray
arr : ndarray (Ni..., M, Nk...)
Input array.
args : any
Additional arguments to `func1d`.
Expand All @@ -46,11 +64,11 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):

Returns
-------
apply_along_axis : ndarray
The output array. The shape of `outarr` is identical to the shape of
out : ndarray (Ni..., Nj..., Nk...)
The output array. The shape of `out` is identical to the shape of
`arr`, except along the `axis` dimension. This axis is removed, and
replaced with new dimensions equal to the shape of the return value
of `func1d`. So if `func1d` returns a scalar `outarr` will have one
of `func1d`. So if `func1d` returns a scalar `out` will have one
fewer dimensions than `arr`.

See Also
Expand Down
0