take_along_axis#

named_arrays.take_along_axis(a, indices, axis)#

Take values from the input array by matching indices along the given axis.

This is the named-array analogue of numpy.take_along_axis(). Unlike the numpy version, indices is broadcast against a by matching axis names, so indices only needs to define the axes it varies along.

Parameters:

  • a (ArrayT) – The source array to take values from.

  • indices (AbstractScalarArray | ArrayT) – The integer indices to take along axis. For scalars and uncertain scalars this is an instance of named_arrays.AbstractScalarArray. For vectors this may either be a scalar (the same indices are used for every component) or a vector of the same type as a (a separate set of indices for each component).

  • axis (str) – The axis of a along which the values are taken. The result replaces this axis with the axes of indices.

Return type:

ArrayT

See also

numpy.take_along_axis()

The equivalent numpy function.

numpy.argsort()

Produces indices suitable for this function.

named_arrays.AbstractArray.take_along_axis()

A method version of this function.

Examples

Sort an array by taking values along an axis using indices from numpy.argsort().

import numpy as np
import named_arrays as na

a = na.ScalarArray(np.array([3, 1, 2, 0]), axes="x")

indices = na.ScalarArray(np.argsort(a.ndarray), axes="x")

na.take_along_axis(a, indices, axis="x")

ScalarArray(
    ndarray=[0, 1, 2, 3],
    axes=('x',),
)