python numpy中axis的含义_python numpy take_along_axis用法及代码示例

通过匹配一维索引和数据切片从输入数组中获取值。

这将迭代沿索引和数据数组中沿指定轴定向的匹配1d切片，并使用前者在后者中查找值。这些切片的长度可以不同。

沿轴返回索引的函数，例如argsort和argpartition，为此功能生成合适的索引。

1.15.0版中的新功能。

参数：arr:ndarray (Ni…, M, Nk…)：

源数组

indices:ndarray (Ni…, J, Nk…)：

沿arr的每个1d切片取的指标。这必须匹配arr的尺寸，但是尺寸Ni和Nj仅需要针对arr进行广播。

axis:int：

沿其取一维切片的轴。如果axis为None，则将输入数组视为首先被展平为1d，以确保与sort和argsort。

返回值：out:ndarray (Ni…, J, Nk…)：

索引结果。

注意：

这等效于(但比以下方法的使用更快)ndindex和s_，它设置了每个ii和kk到索引元组：

Ni, M, Nk = a.shape[:axis], a.shape[axis], a.shape[axis+1:]

J = indices.shape[axis] # Need not equal M

out = np.empty(Ni + (J,) + Nk)

for ii in ndindex(Ni):

for kk in ndindex(Nk):

a_1d = a [ii + s_[:,] + kk]

indices_1d = indices[ii + s_[:,] + kk]

out_1d = out [ii + s_[:,] + kk]

for j in range(J):

out_1d[j] = a_1d[indices_1d[j]]

同样，消除内部循环，最后两行将是：

out_1d[:] = a_1d[indices_1d]

例子：

对于此样本数组

>>> a = np.array([[10, 30, 20], [60, 40, 50]])

我们可以直接使用sort或argsort进行排序，并且可以使用此函数

>>> np.sort(a, axis=1)

array([[10, 20, 30],

[40, 50, 60]])

>>> ai = np.argsort(a, axis=1); ai

array([[0, 2, 1],

[1, 2, 0]])

>>> np.take_along_axis(a, ai, axis=1)

array([[10, 20, 30],

[40, 50, 60]])

如果扩展尺寸，则最大值和最小值也一样：

>>> np.expand_dims(np.max(a, axis=1), axis=1)

array([[30],

[60]])

>>> ai = np.expand_dims(np.argmax(a, axis=1), axis=1)

>>> ai

array([[1],

[0]])

>>> np.take_along_axis(a, ai, axis=1)

array([[30],

[60]])

如果要同时获取最大值和最小值，则可以先堆叠索引

>>> ai_min = np.expand_dims(np.argmin(a, axis=1), axis=1)

>>> ai_max = np.expand_dims(np.argmax(a, axis=1), axis=1)

>>> ai = np.concatenate([ai_min, ai_max], axis=1)

>>> ai

array([[0, 1],

[1, 0]])

>>> np.take_along_axis(a, ai, axis=1)

array([[10, 30],

[40, 60]])