任意の数の次元のためのNumpy Subarray

Question

Numpyでは、ND-Array Aがあるとしましょう。 A[...,0] または、最初の次元をスライスすることができます A[0]. 。この操作をすべての次元に一般化したい（最初でも最後でも）任意のnd-arrayにこれを望んでいるので、A3が3Dアレイで、A4が4Dアレイである場合、 func(A3, dim = 1, slice = 0) ください A3[ : , 0 , :] と func(A4, dim = 1, slice = 0) ください A4[ : , 0 , : , : ].

私はしばらくこれを探していますが、最終的に恐ろしいハックをせずにそれを行う方法を見つけました（関心のあるものが最後の場所にあるまで寸法を交換するなど）。だから私がここに投稿しているコードは私のニーズのために機能しますが

1）私はいつも自分自身をより良くするためのアドバイスを探しています

2）私が言ったように、私はしばらくこれを探していて、何も見つけられなかったので、他の人にとっては役に立つかもしれません。

def fancy_subarray(farray, fdim, fslice):
    # Return the farray slice at position fslice in dimension fdim

    ndim = farray.ndim

    # Handle negative dimension and slice indexing
    if fdim < 0:
        fdim += ndim
    if fslice < 0:
        fslice += v.shape[fdim]

    # Initilize slicing tuple
    obj = ()

    for i in range(ndim):
        if i == fdim:
            # Only element fslice in that dimension
            obj += (slice(fslice, fslice+1, 1),)
        else:
            # "Full" dimension
            obj += (slice(None,None,1),)

    return farray[obj].copy()

したがって、この小さな機能は、連結することによってスライスされたタプルを構築するだけです slice(None,None,1) 寸法の位置では、スライスしたくないと slice(fslice, fslice+1, 1) 関心のある次元で。サブアレイを返すよりも。負のインデックスを処理します。

これにはわずかな違いがあり、インデックスをまっすぐにします：A3が3x4x5の場合、 A3[:,0,:] 3x5になります fancy_subarray(A3, fdim = 1, fslice = 0) 3x1x5になります。また、関数は、バウンドディメンションと「自然」のインデックスからハンドルします。もしも fdim >= farray.ndim forループ内の条件が真実ではないので、関数は完全な配列を返すだけです fslice >= farray.shape[fdim] 返されたサブアレイのサイズは、ディメンションfdimで0です。

これは、もちろん、1つの次元で1つの要素を選択するだけでなく、簡単に拡張できます。

ありがとう！

Answer 1

特に処理するときは、あなたは物事を過剰に縮小していると思います fslice. 。あなたが単にした場合：

def fancy_subarray(farray, fdim, fslice):
    fdim += farray.ndim if fdim < 0 else 0
    index = ((slice(None),) * fdim + (fslice,) +
             (slice(None),) * (farray.ndim - fdim - 1))
    return farray[index]

次に、コードをよりコンパクトにするだけでなく、同じ関数が単一のインデックス、スライス、またはインデックスのリストを取得できます。

>>> fancy_subarray(a, 1, 2)
array([[10, 11, 12, 13, 14],
       [30, 31, 32, 33, 34],
       [50, 51, 52, 53, 54]])
>>> fancy_subarray(a, 1, slice(2,3))
array([[[10, 11, 12, 13, 14]],

       [[30, 31, 32, 33, 34]],

       [[50, 51, 52, 53, 54]]])
>>> fancy_subarray(a, 1, [2, 3])
array([[[10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19]],

       [[30, 31, 32, 33, 34],
        [35, 36, 37, 38, 39]],

       [[50, 51, 52, 53, 54],
        [55, 56, 57, 58, 59]]])