numpy ND配列：ある軸に沿ってindex'ed要素を取る（ND選択）

Question

特定の軸に沿ってインデックスを取得します。たとえば、次のような2Dおよび軸= -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次元だけ減らすにはどうすればよいですか？

>>> a[arange(len(axisinds)), axisinds] 
array([ 0.49970414, 0.4263927 , -0.4647834 , -0.29381079, 0.15729681]) 

しかし、これは非常に特別なに限定される：2Dおよび軸の

= -1は（各行の絶対max'esを例アレイを得るために）このように行うことができます1または0の結果ディメンション。 どのようにすればndimとaxis？

Answer 1

は今、私はシンプルなソリューションを自分自身を見つけた：

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]) 

編集：しかし、このアプローチは、（32bit版？）軸方向に最大31の要素に限定されることが判明 - 理由（文書化されていない）の制限は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)