[英]numpy ND array: take index'ed elements along a certain axis (ND choose)
我沿某個軸獲取索引。 例如2D和axis = -1這樣的例子:
>>> axis = -1
>>> a = rand(5, 3) - 0.5; a
array([[ 0.49970414, -0.14251437, 0.2881351 ],
[ 0.3280437 , 0.33766112, 0.4263927 ],
[ 0.37377502, 0.05392274, -0.4647834 ],
[-0.09461463, -0.25347861, -0.29381079],
[-0.09642799, 0.15729681, 0.06048399]])
>>> axisinds = a.__abs__().argmax(axis); axisinds
array([0, 2, 2, 2, 1])
現在如何通過沿該軸獲取索引元素來將數組縮小1維?
對於2D和axis = -1,可以這樣進行(為了獲得示例數組中每一行的絕對最大值):
>>> a[arange(len(axisinds)), axisinds]
array([ 0.49970414, 0.4263927 , -0.4647834 , -0.29381079, 0.15729681])
但這是非常特殊的,並且限於1或0個結果維度。 對於任何ndim
和axis
如何?
現在我自己找到了一個簡單的解決方案:
def choose_axis(inds, a, axis=-1):
return np.choose(inds, np.rollaxis(a, axis))
>>> choose_axis(axisinds, a, -1)
array([ 0.49970414, 0.4263927 , -0.4647834 , -0.29381079, 0.15729681])
編輯:但是,由於np.choose
的(未記錄)限制,這種方法最終僅限於在軸方向(32位?)上最多31個元素。 在很多情況下都可以。
但這是一個
def choose_axis(inds, a, axis=-1):
# handles any number & size of dimensions, and any axis
if (axis + 1) % a.ndim: # move axis to last dim
a = np.moveaxis(a, axis, -1) # = np.rollaxis(a, axis, a.ndim)
shape = a.shape
a = a.reshape(-1, shape[-1]) # 2D
a = a[np.arange(inds.size), inds.ravel()] # effective reduction
return a.reshape(shape[:-1])
因此,ND絕對最小值示例可以像這樣完成:
def absminND(a, axis=-1):
inds = a.__abs__().argmin(axis)
if axis is None:
return a.ravel()[inds]
return choose_axis(inds, a)
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.