임의의 차원 수에 대한 Numpy Subarray

Question

Numpy에서는 nd-array a가 있다고 가정 해 봅시다. A[...,0] 또는 일을하여 첫 번째 차원을 슬라이스 할 수 있습니다 A[0]. 나는 모든 차원 (첫 번째 또는 마지막이 아니라)에 대해이 작업을 일반화하고 싶고 A3이 3D 배열이고 A4가 4D 배열 인 경우 임의 N의 ND 배열에 대해 이것을 원합니다. 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 루프 내부의 IF 조건이 결코 사실이 아니고 IF이기 때문에 함수는 전체 배열을 반환합니다. fslice >= farray.shape[fdim] 반환 된 서브 어레이의 크기는 크기가 0입니다.

이것은 물론 한 차원에서 하나의 요소를 선택하는 것 이상으로 쉽게 확장 될 수 있습니다.

감사!

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