我沿某个轴获得指数。例如,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
?
答案 0 :(得分:0)
现在我自己找到了一个简单的解决方案:
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])
编辑然而,这种方法被限制为轴方向最多31个元素(32位?) - 因为np.choose
的(未记录的)限制。在许多情况下,这是可以的。
然而这是一个
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)